import { z } from "zod";
import { renderMarkdown } from "../markdown";
import type { MarkdownContext } from "../markdown";
import type { Notebook, NotebookCell, NotebookOutput } from "~/lib/code-intel";

const MAX_CELLS = 500;
const MAX_TEXT_CHARS = 20_000;
const MAX_IMAGE_BYTES = 4 * 1024 * 1024;
const IMAGE_MIMES = ["image/png", "image/jpeg", "image/gif"];
const ANSI = /\u001b\[[0-9;]*[A-Za-z]/g;

const multiline = z
  .union([z.string(), z.array(z.string())])
  .transform((value) => (Array.isArray(value) ? value.join("") : value));

const outputSchema = z
  .object({
    output_type: z.string(),
    name: z.string().optional(),
    text: multiline.optional(),
    data: z.record(z.string(), z.unknown()).optional(),
    traceback: z.array(z.string()).optional(),
    ename: z.string().optional(),
    evalue: z.string().optional(),
  })
  .loose();

const cellSchema = z
  .object({
    cell_type: z.string(),
    source: multiline.default(""),
    execution_count: z.number().nullable().optional(),
    outputs: z.array(outputSchema).optional(),
  })
  .loose();

const notebookSchema = z
  .object({
    cells: z.array(cellSchema),
    metadata: z
      .object({
        language_info: z.object({ name: z.string().optional() }).loose().optional(),
        kernelspec: z.object({ language: z.string().optional() }).loose().optional(),
      })
      .loose()
      .optional(),
  })
  .loose();

function clip(text: string): string {
  return text.length > MAX_TEXT_CHARS ? `${text.slice(0, MAX_TEXT_CHARS)}\n…` : text;
}

function textOf(value: unknown): string | null {
  const parsed = multiline.safeParse(value);
  return parsed.success ? parsed.data : null;
}

type WithoutId<T> = T extends unknown ? Omit<T, "id"> : never;

function convertOutput(
  output: z.infer<typeof outputSchema>,
  budget: { image: number },
): WithoutId<NotebookOutput> | null {
  if (output.output_type === "stream") {
    return { kind: "text", text: clip(output.text ?? ""), stderr: output.name === "stderr" };
  }
  if (output.output_type === "error") {
    const trace = output.traceback?.join("\n") ?? `${output.ename ?? "Error"}: ${output.evalue ?? ""}`;
    return { kind: "error", text: clip(trace.replace(ANSI, "")) };
  }
  const data = output.data ?? {};
  for (const mime of IMAGE_MIMES) {
    const image = textOf(data[mime]);
    if (image && image.length + budget.image <= MAX_IMAGE_BYTES) {
      budget.image += image.length;
      return { kind: "image", mime, data: image.replace(/\s/g, "") };
    }
  }
  const plain = textOf(data["text/plain"]);
  return plain === null ? null : { kind: "text", text: clip(plain), stderr: false };
}

/** Converts nbformat 4 JSON into display cells; markdown goes through the sanitizing renderer and HTML outputs are dropped. */
export function renderNotebook(source: string, context: MarkdownContext): Notebook | null {
  let json: unknown;
  try {
    json = JSON.parse(source);
  } catch {
    return null;
  }
  const parsed = notebookSchema.safeParse(json);
  if (!parsed.success) return null;
  const budget = { image: 0 };
  const cells = parsed.data.cells.slice(0, MAX_CELLS).map((cell, id): NotebookCell => {
    if (cell.cell_type === "markdown") return { id, kind: "markdown", html: renderMarkdown(cell.source, context) };
    if (cell.cell_type === "code") {
      const outputs = (cell.outputs ?? []).flatMap((output) => convertOutput(output, budget) ?? []);
      return {
        id,
        kind: "code",
        source: clip(cell.source),
        executionCount: cell.execution_count ?? null,
        outputs: outputs.map((output, outputId) => ({ ...output, id: outputId })),
      };
    }
    return { id, kind: "raw", text: clip(cell.source) };
  });
  const metadata = parsed.data.metadata;
  return {
    language: metadata?.language_info?.name ?? metadata?.kernelspec?.language ?? "python",
    cells,
    truncated: parsed.data.cells.length > MAX_CELLS,
  };
}