import { defineCommand } from "citty";
import {
  apiStreamNdjson,
  type ApiRequestInit,
  type ApiRequestOptions,
} from "../lib/api";
import { ctx } from "../lib/context";
import { createLogger } from "../lib/logger";

export type LogApiEvent =
  | {
      type: "log";
      targetName: string;
      pod: string;
      container: string;
      message: string;
    }
  | { type: "heartbeat"; timestamp: string }
  | {
      type: "error";
      message: string;
      pod?: string;
      container?: string;
      retryable: boolean;
    };

export type LogsApiStream = <T>(
  path: string,
  init?: ApiRequestInit,
  options?: ApiRequestOptions,
) => AsyncIterable<T>;

export type LogEventWriter = (
  prefix: string,
  message: string,
  error: boolean,
) => void;

function defaultWriter(): LogEventWriter {
  const loggers = new Map<string, ReturnType<typeof createLogger>>();
  return (prefix, message, error) => {
    let logger = loggers.get(prefix);
    if (!logger) {
      logger = createLogger(prefix);
      loggers.set(prefix, logger);
    }
    if (error) logger.warn`${message}`;
    else logger`${message}`;
  };
}

export function parseLogTailLines(value: string): number {
  const tailLines = Number(value);
  if (!/^\d+$/.test(value) || !Number.isSafeInteger(tailLines))
    throw new Error("-n must be a non-negative integer");
  return tailLines;
}

export async function runLogs(
  project: string,
  service: string | undefined,
  follow: boolean,
  signal?: AbortSignal,
  stream: LogsApiStream = apiStreamNdjson,
  write: LogEventWriter = defaultWriter(),
  tailLines?: number,
): Promise<void> {
  if (
    tailLines !== undefined &&
    (!Number.isSafeInteger(tailLines) || tailLines < 0)
  )
    throw new Error("-n must be a non-negative integer");
  const query = new URLSearchParams();
  if (service) query.set("service", service);
  if (follow) query.set("follow", "true");
  if (tailLines !== undefined) query.set("tailLines", String(tailLines));
  const suffix = query.size ? `?${query}` : "";

  for await (const event of stream<LogApiEvent>(
    `/workspaces/${encodeURIComponent(project)}/logs${suffix}`,
    { signal },
    { timeoutMs: follow ? 0 : undefined },
  )) {
    if (event.type === "heartbeat") continue;
    if (event.type === "error") {
      write(service ?? event.pod ?? project, event.message, true);
      continue;
    }
    write(event.pod || event.targetName || service || project, event.message, false);
  }
}

export const logs = defineCommand({
  meta: {
    name: "logs",
    description: "Show deployment logs",
  },
  args: {
    follow: {
      type: "boolean",
      alias: "f",
      description: "Follow log output",
    },
    tail: {
      type: "string",
      alias: "n",
      description: "Show the last N lines per container",
    },
  },
  async run({ args }) {
    const service = args._[0];
    const controller = new AbortController();
    const abort = () => controller.abort();
    process.on("SIGINT", abort);
    process.on("SIGTERM", abort);
    try {
      await runLogs(
        ctx().project,
        service,
        Boolean(args.follow),
        controller.signal,
        apiStreamNdjson,
        defaultWriter(),
        args.tail === undefined ? undefined : parseLogTailLines(args.tail),
      );
    } catch (error) {
      if (!controller.signal.aborted) throw error;
    } finally {
      process.off("SIGINT", abort);
      process.off("SIGTERM", abort);
    }
  },
});