import { spawn } from "node:child_process";
import { mkdtemp, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";
import { Type } from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { filePath, hash, readText } from "./hashline.ts";

const binary = fileURLToPath(new URL("../../libexec/visualize-code", import.meta.url));
const sequence = Type.Array(Type.String({ pattern: "^[A-Za-z]{4}$" }), { minItems: 1 });

function matches(anchors: string[], sequence: string[]): number[] {
  return anchors.flatMap((_, i) => (sequence.every((anchor, j) => anchor === anchors[i + j]) ? [i] : []));
}

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "visualize_code",
    label: "visualize_code",
    description:
      "Render a range of a UTF-8 text file as a syntax-highlighted image, with each line's hash " +
      "in the gutter (the same hashes read shows and edit accepts) and ↪ marking soft-wrapped " +
      "continuations. start is a consecutive hash sequence that must match exactly once; the range " +
      "begins at its first line. end is a consecutive hash sequence; the range ends at the last line " +
      "of its first match at or after start. Without start, the range begins at the top of the file; " +
      "without end, it continues to the end of the file. Ranges that don't fit in one image are cut " +
      "off at a line boundary, with a start to continue from.",
    promptSnippet: "Render file ranges, selected by content hashes, as syntax-highlighted images",
    parameters: Type.Object({
      path: Type.String(),
      start: Type.Optional(sequence),
      end: Type.Optional(sequence),
    }, { additionalProperties: false }),
    async execute(_id, args, signal, _onUpdate, ctx) {
      const path = filePath(args.path, ctx.cwd);
      const { text, lines } = await readText(path, signal);
      if (!lines.length) throw new Error("File is empty.");
      const anchors = lines.map(hash);

      const starts = args.start ? matches(anchors, args.start) : [0];
      if (starts.length !== 1) {
        throw new Error(starts.length
          ? "start is ambiguous; add adjacent hashes."
          : "start hashes not found (stale reference); read the file again.");
      }
      const start = starts[0];
      let end = lines.length;
      if (args.end) {
        const match = matches(anchors, args.end).find((i) => i >= start);
        if (match === undefined) {
          throw new Error("end hashes not found at or after start (stale reference?); read the file again.");
        }
        end = match + args.end.length;
      }
      // Bounds argv; still far more lines than fit in one image.
      const labels = anchors.slice(start, Math.min(end, start + 2000));

      const dir = await mkdtemp(join(tmpdir(), "visualize-code-"));
      try {
        const image = join(dir, "code.png");
        const shown = await new Promise<number>((resolve, reject) => {
          const child = spawn(binary, [image, path, String(start + 1), ...labels], { signal });
          let stdout = "";
          let stderr = "";
          child.stdout.on("data", (chunk) => (stdout += chunk));
          child.stderr.on("data", (chunk) => (stderr += chunk));
          child.on("error", (error: NodeJS.ErrnoException) => reject(error.code === "ENOENT"
            ? new Error(`${binary} is missing; build src/visualize-code and copy the binary there.`)
            : error));
          child.on("close", (code) => code === 0
            ? resolve(Number(stdout))
            : reject(new Error(stderr.trim() || `visualize-code exited with code ${code}`)));
          // Failures surface through the exit status, so an early exit shouldn't crash us with EPIPE.
          child.stdin.on("error", () => {});
          child.stdin.end(text);
        });

        const next = start + shown;
        let summary = `Lines ${start + 1}-${next} of ${lines.length}.`;
        if (next < end) {
          const continuation = [1, 2, 3, 4, 5, 6, 7, 8]
            .map((length) => anchors.slice(next, next + length))
            .find((candidate) => matches(anchors, candidate).length === 1);
          summary += continuation
            ? ` Continue with start: ${JSON.stringify(continuation)}${args.end ? " and the same end" : ""}.`
            : ` No short hash sequence uniquely identifies line ${next + 1}; use read with offset=${next + 1} to continue.`;
        }
        return {
          content: [
            { type: "text", text: summary },
            { type: "image", data: (await readFile(image)).toString("base64"), mimeType: "image/png" },
          ],
          details: {},
        };
      } finally {
        await rm(dir, { recursive: true, force: true });
      }
    },
  });
}