import { defineCommand } from "citty";
import { Listr } from "listr2";
import type { ComposeSpecification } from "../schema/docker.d";
import type { KuberResource } from "../types";
import { ctx } from "../lib/context";
import {
  assertManagedNamespace,
  applyResource,
  deleteResource,
  getStaleResources,
  sortResources,
} from "../lib/apply";
import { buildServices } from "../lib/build";
import { composeToKubernetes, type KubernetesResource } from "../lib/convert";
import {
  getComposePostgresClaims,
  reconcilePostgresClaims,
} from "../lib/database";
import { getComposeS3Claims, reconcileS3Claims } from "../lib/storage";
import { restartDeployment, waitForDeploymentRollout } from "../lib/shared";

type UpContext = {
  compose?: ComposeSpecification;
  changedDeployments?: string[];
  serviceEnv?: Record<string, Record<string, string>>;
  resources?: KubernetesResource[];
  staleResources?: KubernetesResource[];
};

export function mergeServiceEnv(
  current: Record<string, Record<string, string>> | undefined,
  next: Record<string, Record<string, string>>,
): Record<string, Record<string, string>> {
  const merged = { ...current };
  for (const [service, environment] of Object.entries(next)) {
    merged[service] = { ...merged[service], ...environment };
  }
  return merged;
}

export function getDeploymentNames(resources: KubernetesResource[]): string[] {
  return resources
    .filter((resource) => resource.kind === "Deployment")
    .map((resource) => resource.metadata?.name)
    .filter((name): name is string => Boolean(name));
}

export async function runUp(build: boolean) {
  const { project, compose, cwd, config, hookContext: getHookContext } = ctx();
  await assertManagedNamespace(project);

  const taskCtx = await new Listr<UpContext>(
    [
      {
        title: "Read compose",
        task: async (taskCtx, task) => {
          taskCtx.compose = await compose();
          if (build) {
            await config.preBuild?.(taskCtx.compose, await getHookContext());
          }
          task.output = `${Object.keys(taskCtx.compose.services ?? {}).length} services`;
        },
      },
      {
        title: "Build images",
        rendererOptions: {
          outputBar: 10,
          persistentOutput: true,
        },
        enabled: async () =>
          build &&
          Object.values((await compose()).services ?? {}).some(
            (service) => service.build,
          ),
        task: async (taskCtx, task) => {
          const result = await buildServices(
            project,
            taskCtx.compose!,
            cwd,
            {
              progress: (message) => {
                task.output = message;
              },
              stream: task.stdout(),
            },
            config,
          );
          taskCtx.changedDeployments = result.changed;
          await config.postBuild?.(result, await getHookContext());
          task.output = `Built ${result.built.length} image${result.built.length === 1 ? "" : "s"}, ${result.changed.length} changed`;
        },
      },
      {
        title: "Reconcile databases",
        enabled: async () =>
          getComposePostgresClaims(await compose()).length > 0,
        task: async (taskCtx, task) => {
          taskCtx.serviceEnv = mergeServiceEnv(
            taskCtx.serviceEnv,
            await reconcilePostgresClaims(project, taskCtx.compose!),
          );
          task.output = `${Object.keys(taskCtx.serviceEnv).length} service${Object.keys(taskCtx.serviceEnv).length === 1 ? "" : "s"}`;
        },
      },
      {
        title: "Reconcile S3 storage",
        enabled: async () => getComposeS3Claims(await compose()).length > 0,
        task: async (taskCtx, task) => {
          taskCtx.serviceEnv = mergeServiceEnv(
            taskCtx.serviceEnv,
            await reconcileS3Claims(project, taskCtx.compose!),
          );
          task.output = `${Object.keys(taskCtx.serviceEnv).length} service${Object.keys(taskCtx.serviceEnv).length === 1 ? "" : "s"}`;
        },
      },
      {
        title: "Render manifests",
        task: async (taskCtx, task) => {
          taskCtx.resources = await composeToKubernetes(
            project,
            taskCtx.compose!,
            cwd,
            taskCtx.serviceEnv,
            config.registry,
          );
          await config.postRender?.(
            taskCtx.resources as KuberResource[],
            await getHookContext(),
          );
          task.output = `${taskCtx.resources.length} resources`;
        },
      },
      {
        title: "Plan reconciliation",
        task: async (taskCtx, task) => {
          taskCtx.staleResources = (await getStaleResources(
            project,
            taskCtx.resources!,
          )) as KubernetesResource[];
          task.output = `${taskCtx.staleResources.length} stale resource${taskCtx.staleResources.length === 1 ? "" : "s"}`;
        },
      },
      {
        title: "Apply resources",
        task: (taskCtx, task) => {
          const resources = sortResources(taskCtx.resources!);

          task.output = `${resources.length} resources queued`;
          return task.newListr(
            resources.map((resource) => ({
              title: `${resource.kind} ${resource.metadata?.name}`,
              task: () => applyResource(resource),
            })),
            { concurrent: false, exitOnError: true },
          );
        },
      },
      {
        title: "Run post-apply hook",
        enabled: () => Boolean(config.postApply),
        task: async (taskCtx) => {
          await config.postApply?.(
            taskCtx.resources! as KuberResource[],
            await getHookContext(),
          );
        },
      },
    ],
    { rendererOptions: { collapseErrors: false } },
  ).run();

  const changedDeployments = taskCtx.changedDeployments ?? [];
  if (changedDeployments.length > 0) {
    await new Listr([
      {
        title: "Restart changed deployments",
        task: (_taskCtx, task) => {
          task.output = `${changedDeployments.length} deployments queued`;
          return task.newListr(
            changedDeployments.map((name) => ({
              title: `Deployment ${name}`,
              task: () => restartDeployment(name),
            })),
            { concurrent: false, exitOnError: true },
          );
        },
      },
    ]).run();
  }

  const deployments = getDeploymentNames(taskCtx.resources!);
  if (deployments.length > 0) {
    await new Listr([
      {
        title: "Wait for rollout",
        task: (_taskCtx, task) => {
          task.output = `${deployments.length} deployments queued`;
          return task.newListr(
            deployments.map((name) => ({
              title: `Deployment ${name}`,
              task: () =>
                waitForDeploymentRollout(name, config.rolloutTimeoutMs),
            })),
            { concurrent: false, exitOnError: true },
          );
        },
      },
    ]).run();
  }

  if (taskCtx.staleResources && taskCtx.staleResources.length > 0) {
    const resources = sortResources(taskCtx.staleResources).reverse();
    await new Listr([
      {
        title: "Delete stale resources",
        task: (_taskCtx, task) => {
          task.output = `${resources.length} resources queued`;
          return task.newListr(
            resources.map((resource) => ({
              title: `${resource.kind} ${resource.metadata?.name}`,
              task: () => deleteResource(resource),
            })),
            { concurrent: false, exitOnError: true },
          );
        },
      },
    ]).run();
  }
}

export const up = defineCommand({
  meta: {
    name: "up",
    description: "Create and start deployments",
  },
  args: {
    build: {
      type: "boolean",
      default: true,
      alias: "b",
      description: "Build images before starting deployments",
      negativeDescription: "Don't build an image, even if it's policy",
    },
  },
  async run({ args }) {
    await runUp(args.build);
  },
});