import { Buffer } from "node:buffer";
import { extname } from "node:path";
import type { Context } from "@earendil-works/chord";
import { BACKGROUND_CONTEXT, withAbortSignal } from "@earendil-works/chord/context";
import { Type } from "@earendil-works/pi-ai";
import { type FileError, getOrThrow } from "@earendil-works/pi-durable/env";
import ignore, { type Ignore } from "ignore";
import { defineExtension, type Environment, type Tool } from "@berry/api";
import {
  decodeText,
  fitting,
  MAX_BYTES,
  MAX_LINES,
  MAX_TEXT_BYTES,
  planEdit,
  renderRead,
  type TextFile,
} from "./hashline.ts";
import { hasWorkspace, workspace } from "./workspace.ts";
import { hasCommand, powershellCommand } from "./commands.ts";

type EditMode = "naive" | "hashline" | "both";
type State = { editMode: EditMode };

/** Only old history has these: the harness keeps the working directory now. */
type Cd = { type: "cd"; cwd: string };
export type Event = Cd | { type: "editMode"; editMode: EditMode };

const IMAGE_TYPES: Record<string, string> = {
  ".png": "image/png",
  ".jpg": "image/jpeg",
  ".jpeg": "image/jpeg",
  ".gif": "image/gif",
  ".webp": "image/webp",
  ".bmp": "image/bmp",
};
const MAX_IMAGE_BYTES = 5 * 1024 * 1024;

class TooLarge extends Error {}

/** The whole of the regular file at `path`, which must be no larger than `limit`. */
async function readBytes(
  env: Environment,
  path: string,
  limit: number,
  context: Context,
): Promise<Uint8Array> {
  const reader = getOrThrow(await env.openBinaryReader(path, undefined, context));
  try {
    const { size } = getOrThrow(await reader.info(context));
    if (size > limit) throw new TooLarge(`${path} is larger than ${limit} bytes`);
    return getOrThrow(await reader.read(0, size, context));
  } finally {
    await reader.close(context);
  }
}

async function readText(env: Environment, path: string, context: Context): Promise<TextFile> {
  try {
    return decodeText(await readBytes(env, path, MAX_TEXT_BYTES, context));
  } catch (error) {
    if (error instanceof TooLarge) throw new Error("File exceeds the 50 MiB edit/read limit.");
    throw error;
  }
}

const base64 = (bytes: Uint8Array) =>
  Buffer.from(bytes.buffer, bytes.byteOffset, bytes.byteLength).toString("base64");

const mutations = new Map<string, Promise<unknown>>();

/** Serialize changes to one file, whatever path (or symlink) names it on whichever machine. */
async function mutating<T>(env: Environment, path: string, change: () => Promise<T>): Promise<T> {
  const canonical = await env.canonicalPath(path, BACKGROUND_CONTEXT);
  const key = `${env.id}\0${canonical.ok ? canonical.value : path}`;
  const result = (mutations.get(key) ?? Promise.resolve()).then(change);
  const settled = result.catch(() => {});
  mutations.set(key, settled);
  void settled.then(() => mutations.get(key) === settled && mutations.delete(key));
  return await result;
}

/** The last `MAX_LINES` lines and `MAX_BYTES` of output, keeping at least the final line. */
function tail(output: string): string {
  const ending = output.endsWith("\n") ? "\n" : "";
  const lines = output.slice(0, output.length - ending.length).split("\n").slice(-MAX_LINES);
  const kept = Math.max(1, fitting(lines.toReversed()));
  return lines.slice(-kept).join("\n") + ending;
}

const LIST_LIMIT = 1000;
const OPAQUE_DIRS = new Set([".git", ".jj", "node_modules"]);

function unavailable(error: FileError): boolean {
  return ["not_found", "not_directory", "permission_denied", "invalid"].includes(error.code);
}

async function list(env: Environment, path: string, depth: number, context: Context) {
  const found: string[] = [];
  type Rule = { prefix: string; matcher: Ignore };
  async function visit(
    dir: string,
    prefix: string,
    level: number,
    inherited: Rule[],
  ): Promise<void> {
    const entries = getOrThrow(await env.listDir(dir, context)).sort((a, b) =>
      a.name.localeCompare(b.name)
    );
    let rules = inherited;
    const gitignore = entries.find((entry) => entry.name === ".gitignore");
    if (gitignore) {
      const source = await env.readTextFile(gitignore.path, context);
      if (source.ok) rules = [...rules, { prefix, matcher: ignore().add(source.value) }];
      else if (!unavailable(source.error)) throw source.error;
    }
    for (const entry of entries) {
      if (found.length >= LIST_LIMIT) return;
      const directory = entry.kind === "directory";
      const shown = prefix + entry.name;
      const candidate = directory ? `${shown}/` : shown;
      let ignored = false;
      for (const rule of rules) {
        const result = rule.matcher.test(candidate.slice(rule.prefix.length));
        if (result.ignored) ignored = true;
        else if (result.unignored) ignored = false;
      }
      if (ignored) continue;
      found.push(directory ? `${shown}/` : entry.kind === "symlink" ? `${shown}@` : shown);
      if (directory && level < depth && !OPAQUE_DIRS.has(entry.name)) {
        await visit(entry.path, `${shown}/`, level + 1, rules);
      }
    }
  }
  await visit(getOrThrow(await env.absolutePath(path, context)), "", 1, []);
  if (found.length >= LIST_LIMIT) {
    found.push(`[Stopped at ${LIST_LIMIT} entries; list a subdirectory or lower depth.]`);
  }
  return found.join("\n");
}

/** `AGENTS.md` of the working directory and its ancestors, outermost first. */
async function projectContext(env: Environment, context: Context) {
  const files: { path: string; content: string }[] = [];
  for (let dir = env.cwd;;) {
    for (const parts of [[dir, ".berry", "AGENTS.md"], [dir, "AGENTS.md"]]) {
      const path = getOrThrow(await env.joinPath(parts, context));
      const content = await env.readTextFile(path, context);
      if (content.ok) {
        files.unshift({ path, content: content.value.trim() });
        break;
      }
    }
    const parent = getOrThrow(await env.joinPath([dir, ".."], context));
    if (parent === dir) return files;
    dir = parent;
  }
}

const hashes = Type.Array(Type.String({ pattern: "^[A-Za-z]{4}$" }));

const modes = Type.Union([
  Type.Literal("naive"),
  Type.Literal("hashline"),
  Type.Literal("both"),
], {
  title: "Editing mode",
  description:
    "Naive uses exact text replacements; hashline uses line anchors; both offers either.",
  default: "both",
});
const hashEdit = Type.Object({
  before: Type.Optional(hashes),
  old: hashes,
  after: Type.Optional(hashes),
  new: Type.Array(Type.String({ pattern: "^[^\\r\\n\\u0000]*$" })),
}, { additionalProperties: false });
const naiveEdit = Type.Object({
  oldText: Type.String({ pattern: "^[^\\u0000]*$" }),
  newText: Type.String({ pattern: "^[^\\u0000]*$" }),
}, { additionalProperties: false });
const editForms = { naive: naiveEdit, hashline: hashEdit, both: Type.Union([hashEdit, naiveEdit]) };

export default defineExtension<State, Event>({
  name: "code",
  sortOrder: 10,
  defaultEnabled: true,
  formerly: ["coding"],
  description: "edit files and run shell commands",
  initial: () => ({ editMode: "both" }),
  apply: (state, event) => event.type === "editMode" ? { editMode: event.editMode } : state,

  setup(ext) {
    ext.promptProse({
      default: "guidelines:\n" + [
        "when changing multiple separate locations in one file, use one edit call with multiple entries in edits[] instead of multiple edit calls",
        "if you need to start a background process, use tmux",
        "when a command is missing, you can get it via `nix shell` or `nix run`",
        "when producing source code, avoid comments that explain the plain mechanics of the code, " +
        "but preserve terse comments that explain subtle motivation. " +
        "(e.g. '// work around a layout bug in firefox' and not '// set the style X to Y')",
        "we will prioritise elegance, simplicity, and long-term maintainability of code. " +
        "we have license to _work slower_ in order to do it right the first time",
        "additionally wrt simplicity: don't outline simple functions if they're only going to be used in one place",
        "avoid tautological tests: tests should not need to change when details of their implementation changes",
        "$USER is likely using jj instead of git. pass --git flags to jj diff to work around the lack of terminal color",
      ].map((g) => `- ${g}`).join("\n"),
    });
    ext.configuration({
      parameters: Type.Object({ editMode: modes }),
      values: (state) => state,
      events: ({ editMode }, state) =>
        editMode === state.editMode ? [] : [{ type: "editMode", editMode }],
    });
    ext.prompt("project", async ({ env }) => {
      if (env === undefined) return undefined;
      const context = await projectContext(env, BACKGROUND_CONTEXT);
      if (context.length === 0) return undefined;
      return "# project context\n\nproject-specific instructions and guidelines:" +
        context.map(({ path, content }) => `\n\n## ${path}\n\n${content}`).join("");
    });

    ext.tool({
      name: "write",
      expose: "both",
      description: "Create or overwrite a file, creating parent directories.",
      parameters: Type.Object({ path: Type.String(), content: Type.String() }),
      returns: Type.String(),
      async execute({ path, content }, ctx) {
        const { env } = ctx;
        const context = withAbortSignal(ctx.signal, BACKGROUND_CONTEXT);
        const absolute = getOrThrow(await env.absolutePath(path, context));
        return await mutating(env, absolute, async () => {
          getOrThrow(await env.writeFile(absolute, content, context));
          return `Wrote ${content.split("\n").length} lines to ${absolute}.`;
        });
      },
    });

    ext.tool({
      name: "list",
      expose: "both",
      description:
        "List a directory (default: the working directory) as relative paths, one per line, with / after " +
        "directories and @ after symlinks. Skips paths ignored by .gitignore. depth (default 1) descends " +
        "except into .git, .jj, and " +
        `node_modules. Stops at ${LIST_LIMIT} entries.`,
      parameters: Type.Object({
        path: Type.Optional(Type.String()),
        depth: Type.Optional(Type.Integer({ minimum: 1, maximum: 10 })),
      }),
      returns: Type.String(),
      execute: ({ path, depth }, ctx) =>
        list(ctx.env, path ?? ".", depth ?? 1, withAbortSignal(ctx.signal, BACKGROUND_CONTEXT)),
    });

    const shellParameters = Type.Object({
      command: Type.String(),
      timeout: Type.Number({ exclusiveMinimum: 0, description: "Seconds." }),
    });
    const shell = (name: "bash" | "powershell"): Tool<State, Event, typeof shellParameters, {
      output: string;
      exitCode: number;
      fullOutputPath?: string;
    }> => ({
      name,
      expose: "both",
      description:
        `Run a ${name === "bash" ? "bash" : "PowerShell"} command in the working directory. ` +
        (name === "powershell"
          ? "Uses PS7 when available, otherwise PS5; no profile or interactive input. "
          : "") +
        `Returns the exit code and the last ${MAX_LINES} lines / ${
          MAX_BYTES / 1024
        } KiB of combined ` +
        "stdout and stderr, with fullOutputPath holding all of a longer output. A nonzero exit code " +
        "does not reject; a timeout does.",
      parameters: shellParameters,
      returns: Type.Object({
        output: Type.String(),
        exitCode: Type.Integer(),
        fullOutputPath: Type.Optional(Type.String()),
      }),
      async execute({ command, timeout }, ctx) {
        let output = "";
        const invocation = name === "powershell"
          ? await powershellCommand(ctx.env, command)
          : command;
        const result = await ctx.env.exec(invocation, {
          timeout,
          onOutput: (text) => {
            ctx.output(text);
            output += text;
            if (output.length > 4 * MAX_BYTES) output = output.slice(-2 * MAX_BYTES);
            // direct-call output loses escape characters in the harness's display sanitization
            ctx.progress({ rawOutput: tail(output) });
          },
          spill: { afterBytes: MAX_BYTES, afterLines: MAX_LINES },
        }, withAbortSignal(ctx.signal, BACKGROUND_CONTEXT));
        if (!result.ok) {
          if (result.error.code === "timeout") {
            throw new Error(`Command timed out after ${timeout} seconds.`);
          }
          if (result.error.code === "aborted") throw new Error("Command aborted.");
          throw result.error;
        }
        const { exitCode, spillPath } = result.value;
        return {
          output: tail(output),
          exitCode,
          ...(spillPath === undefined ? {} : { fullOutputPath: spillPath }),
        };
      },
      text: ({ output, exitCode, fullOutputPath }) =>
        [
          output.trimEnd() || "(no output)",
          `[exit code ${exitCode}${fullOutputPath ? `; full output: ${fullOutputPath}` : ""}]`,
        ].join("\n\n"),
    });
    ext.tool(shell("bash"));

    return (conv) => {
      const { editMode } = conv.state;
      const anchored = editMode !== "naive";

      conv.tool({
        name: "read",
        expose: "both",
        description: (anchored
          ? "Read UTF-8 text as hash│content, each hash four letters of the exact line content, " +
            "independent of position; equal lines share a hash. "
          : "Read UTF-8 text. ") +
          `At most ${MAX_LINES} lines / ${
            MAX_BYTES / 1024
          } KiB; page with offset (1-indexed line) and ` +
          "limit (line count). Images (png, jpeg, gif, webp, bmp) resolve to an image for image().",
        parameters: Type.Object({
          path: Type.String(),
          offset: Type.Optional(Type.Integer({ minimum: 1 })),
          limit: Type.Optional(Type.Integer({ minimum: 1 })),
        }),
        returns: Type.Union([
          Type.String(),
          Type.Object({
            type: Type.Literal("image"),
            data: Type.String(),
            mimeType: Type.String(),
          }),
        ]),
        async execute({ path, offset, limit }, ctx) {
          ctx.progress({ editMode });
          const context = withAbortSignal(ctx.signal, BACKGROUND_CONTEXT);
          const absolute = getOrThrow(await ctx.env.absolutePath(path, context));
          const mimeType = IMAGE_TYPES[extname(absolute).toLowerCase()];
          if (mimeType !== undefined) {
            const bytes = await readBytes(ctx.env, absolute, MAX_IMAGE_BYTES, context).catch(
              (error) => {
                if (error instanceof TooLarge) {
                  throw new Error("Image exceeds 5 MiB; scale it down with bash first.");
                }
                throw error;
              },
            );
            return { type: "image", data: base64(bytes), mimeType };
          }
          return renderRead(
            await readText(ctx.env, absolute, context),
            offset,
            limit,
            anchored,
          );
        },
      });

      conv.tool({
        name: "edit",
        expose: "both",
        description: "Edit a file. Each of edits[] is " +
          (editMode === "both" ? "either " : "") +
          (editMode !== "hashline"
            ? "{ oldText, newText }, replacing literal text that must occur exactly once"
            : "") +
          (editMode === "both" ? ", or " : "") +
          (anchored
            ? "{ before?, old, after?, new } with line hashes from read, where before + old + after " +
              "must match exactly once and only old is replaced (old: [] inserts, new: [] deletes; " +
              "list every removed line's hash; new lines can't contain CR/LF/NUL)"
            : "") +
          ". Add adjacent context to disambiguate; an empty match is allowed only for an empty file. " +
          "All edits match the original file and are validated before writing; ambiguous, stale, or " +
          "overlapping ones fail. Batch a file's edits in one call. Returns a diff" +
          (anchored ? " with hashes on changed rows." : "."),
        parameters: Type.Object({
          path: Type.String(),
          edits: Type.Array(editForms[editMode], { minItems: 1 }),
        }, { additionalProperties: false }),
        returns: Type.String(),
        async execute({ path, edits }, ctx) {
          ctx.progress({ editMode });
          const { env } = ctx;
          const context = withAbortSignal(ctx.signal, BACKGROUND_CONTEXT);
          const absolute = getOrThrow(await env.absolutePath(path, context));
          return await mutating(env, absolute, async () => {
            const file = await readText(env, absolute, context);
            const { content, diff } = planEdit(file, edits, anchored);
            if (content === file.raw) return "No changes made.";
            ctx.signal.throwIfAborted();
            const current = await readBytes(env, absolute, MAX_TEXT_BYTES, context);
            if (!Buffer.from(current).equals(Buffer.from(file.raw))) {
              throw new Error("File changed while preparing the edit; read it again.");
            }
            ctx.signal.throwIfAborted();
            getOrThrow(await env.writeFile(absolute, content, context));
            const shown = diff.slice(0, Math.min(MAX_LINES, fitting(diff)));
            const truncated = shown.length < diff.length
              ? "\n[Diff truncated; read for more context.]"
              : "";
            return `Applied ${edits.length} edit(s).\n\n${shown.join("\n")}${truncated}`;
          });
        },
      });

      if (conv.env?.platform === "win32") {
        conv.tool(shell("powershell"));
      }

      conv.command({
        name: "workspace",
        enabled: ({ env }) => env ? hasWorkspace(env) : false,
        description: "create and switch to a jj workspace (under $BERRY_WORKSPACES_DIR)",
        async execute(args, ctx) {
          if (args.trim()) throw new Error("Usage: /workspace");
          return await ctx.cd(await workspace(ctx.env, ctx.tid, BACKGROUND_CONTEXT), {
            inheritProject: true,
          });
        },
      });

      conv.prompt(async ({ env }) => {
        const commands = env
          ? (await Promise.all(
            ["tmux", "nix", "jj"].map(async (name) =>
              await hasCommand(env, name) ? name : undefined
            ),
          )).filter((name) => name !== undefined)
          : [];
        return [
          ...(env ? [`current working directory: ${env.cwd}`] : []),
          ...(env?.platform === "win32"
            ? ["use powershell for commands in this Windows environment"]
            : []),
          "@-references to files (e.g. @path/to/file.ts) are relative to the working directory unless " +
          "absolute: read them directly, without the @, rather than searching for them.",
          ...(env ? [`available workflow commands: ${commands.join(", ") || "none"}`] : []),
        ].join("\n\n");
      });
    };
  },
});