import { Buffer } from "node:buffer";
import { unlinkSync } from "node:fs";
import {
  CustomEditor,
  type ExtensionAPI,
  type KeybindingsManager,
} from "@earendil-works/pi-coding-agent";
import {
  AttachmentStore,
  findAtomicDeleteRange,
  segmentWithAtomicImages,
  splitImagePlaceholders,
} from "./src/attachments.ts";
import { loadPiClipboardImage } from "./src/clipboard-file.ts";

interface EditorSegmentationAccess {
  segment: (text: string, mode?: "word" | "grapheme") => Intl.SegmentData[];
}

interface EditorStateAccess {
  state: { lines: string[]; cursorLine: number; cursorCol: number };
  pushUndoSnapshot?: () => void;
  setCursorCol?: (column: number) => void;
  lastAction?: unknown;
  historyIndex?: number;
}

class AtomicImageEditor extends CustomEditor {
  private readonly editorKeybindings: KeybindingsManager;
  private readonly attachments: AttachmentStore;

  constructor(
    tui: ConstructorParameters<typeof CustomEditor>[0],
    theme: ConstructorParameters<typeof CustomEditor>[1],
    editorKeybindings: KeybindingsManager,
    attachments: AttachmentStore,
  ) {
    super(tui, theme, editorKeybindings);
    this.editorKeybindings = editorKeybindings;
    this.attachments = attachments;

    const editor = this as unknown as EditorSegmentationAccess;
    editor.segment = (text, mode = "grapheme") =>
      segmentWithAtomicImages(text, mode, (placeholder) => attachments.has(placeholder));
  }

  override insertTextAtCursor(text: string): void {
    const image = loadPiClipboardImage(text);
    if (!image) {
      super.insertTextAtCursor(text);
      return;
    }

    const attachment = this.attachments.add(image.bytes, image.mimeType);
    super.insertTextAtCursor(attachment.placeholder);

    try {
      unlinkSync(image.path);
    } catch {
      // The attachment is already in memory; stale temp-file cleanup is best effort.
    }
  }

  override handleInput(data: string): void {
    if (this.handleAtomicDelete(data)) return;
    super.handleInput(data);
  }

  private handleAtomicDelete(data: string): boolean {
    const backward = this.editorKeybindings.matches(data, "tui.editor.deleteCharBackward");
    const forward = this.editorKeybindings.matches(data, "tui.editor.deleteCharForward");
    if (!backward && !forward) return false;

    const cursor = this.getCursor();
    const line = this.getLines()[cursor.line] ?? "";
    const range = findAtomicDeleteRange(
      line,
      cursor.col,
      backward ? "backward" : "forward",
      (placeholder) => this.attachments.has(placeholder),
    );
    if (!range) return false;

    const editor = this as unknown as EditorStateAccess;
    editor.pushUndoSnapshot?.();
    editor.state.lines[cursor.line] = line.slice(0, range.start) + line.slice(range.end);
    editor.state.cursorLine = cursor.line;
    if (editor.setCursorCol) editor.setCursorCol(range.start);
    else editor.state.cursorCol = range.start;
    editor.lastAction = null;
    editor.historyIndex = -1;
    this.onChange?.(this.getText());
    this.tui.requestRender();
    return true;
  }
}

export default function atomicImages(pi: ExtensionAPI): void {
  const attachments = new AttachmentStore();

  pi.on("session_start", (_event, ctx) => {
    attachments.clear();
    if (!ctx.hasUI) return;

    ctx.ui.setEditorComponent(
      (tui, theme, keybindings) =>
        new AtomicImageEditor(tui, theme, keybindings, attachments),
    );
  });

  pi.on("input", (event, ctx) => {
    if (event.source !== "interactive") return { action: "continue" as const };

    const selected = attachments.matching(event.text);
    if (selected.length === 0) return { action: "continue" as const };

    if (!ctx.model?.input.includes("image")) {
      ctx.ui.setEditorText(event.text);
      const modelName = ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : "当前模型";
      ctx.ui.notify(
        `pi-atomic-images: ${modelName} 未声明图片输入支持，请切换到 images=yes 的模型`,
        "warning",
      );
      return { action: "handled" as const };
    }

    return {
      action: "transform" as const,
      // Keep the placeholder in the stored message so the submitted user turn
      // still renders like the editor. The context hook below hides it only
      // from the model-facing copy.
      text: event.text,
      images: [
        ...(event.images ?? []),
        ...selected.map((attachment) => ({
          type: "image" as const,
          mimeType: attachment.mimeType,
          data: Buffer.from(attachment.bytes).toString("base64"),
        })),
      ],
    };
  });

  pi.on("context", (event) => ({
    messages: event.messages.map((message) => {
      if (message.role !== "user" || !Array.isArray(message.content)) return message;
      const imageParts = message.content.filter((part) => part.type === "image");
      if (imageParts.length === 0) return message;

      const placeholderCount = message.content.reduce((count, part) => {
        if (part.type !== "text") return count;
        return (
          count +
          splitImagePlaceholders(part.text).filter((segment) => segment.type === "image").length
        );
      }, 0);
      if (placeholderCount === 0 || placeholderCount > imageParts.length) return message;

      // Pi stores input-event images after the message text. Images created by
      // this extension therefore occupy the final N image slots, where N is
      // the number of live placeholders. Any pre-existing image blocks retain
      // their original positions.
      const attachedImages = imageParts.slice(imageParts.length - placeholderCount);
      const attachedImageSet = new Set(attachedImages);
      const projected: typeof message.content = [];
      let attachmentIndex = 0;
      let textTemplate: (typeof message.content)[number] | undefined;

      for (const part of message.content) {
        if (part.type === "image") {
          if (!attachedImageSet.has(part)) projected.push(part);
          continue;
        }
        if (part.type !== "text") {
          projected.push(part);
          continue;
        }

        textTemplate ??= part;
        for (const segment of splitImagePlaceholders(part.text)) {
          if (segment.type === "image") {
            const image = attachedImages[attachmentIndex++];
            if (image) projected.push(image);
          } else if (segment.text.length > 0) {
            projected.push({ ...part, text: segment.text });
          }
        }
      }

      const hasPromptText = projected.some(
        (part) => part.type === "text" && part.text.trim().length > 0,
      );
      if (!hasPromptText && textTemplate?.type === "text") {
        projected.push({ ...textTemplate, text: "请直接查看随消息附带的图片。" });
      }

      return {
        ...message,
        content: projected,
      };
    }),
  }));
}