import {
  AssistantMessageComponent,
  type ExtensionAPI,
  type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { detectMathBlocks, latexToTerminalText } from "./latex.ts";
import { layoutMathTreeBox, wrapTextCells, type CellBox } from "./text-layout.ts";

export default async function terminalMathExtension(pi: ExtensionAPI) {
  const patchOwner = Symbol("terminal-math-extension-instance");
  let enabled = false;

  installAssistantTranscriptRenderer(patchOwner);

  const setEnabled = async (next: boolean, ctx: ExtensionContext) => {
    if (next) {
      await getMathJax();
      enableAssistantTranscriptRenderer(patchOwner);
    } else {
      disableAssistantTranscriptRenderer(patchOwner);
    }
    enabled = next;
    ctx.ui.setStatus("terminal-math", enabled ? ctx.ui.theme.fg("accent", "∑ math") : undefined);
    ctx.ui.notify(`Terminal math ${enabled ? "enabled" : "disabled"}.`, "info");
  };

  pi.registerCommand("math", {
    description: "Toggle terminal math rendering (on/off)",
    handler: async (args, ctx) => {
      const action = args.trim().toLowerCase();
      if (action && !["on", "off", "toggle"].includes(action)) {
        ctx.ui.notify("Usage: /math [on|off|toggle]", "warning");
        return;
      }
      await setEnabled(action === "on" || (action !== "off" && !enabled), ctx);
    },
  });

  pi.on("session_start", (_event, ctx) => {
    enabled = false;
    disableAssistantTranscriptRenderer(patchOwner);
    if (ctx.mode === "tui") {
      setAssistantTranscriptTheme(ctx.ui.theme);
      ctx.ui.setStatus("terminal-math", undefined);
    }
  });

  pi.on("session_shutdown", () => {
    enabled = false;
    uninstallAssistantTranscriptRenderer(patchOwner);
  });

  pi.on("before_agent_start", async (event, ctx) => {
    if (ctx.mode !== "tui" || !enabled) return undefined;
    return {
      systemPrompt:
        event.systemPrompt +
        `\n\nTerminal math rendering guidance:\n- Use LaTeX display math for substantial equations.\n- Prefer simple equation blocks over dense inline math.\n- Avoid unnecessary LaTeX when plain text is clearer.`,
    };
  });
}

interface MathJaxParser {
  texToMml: (tex: string, options: { display: boolean }) => unknown;
}

let initializedMathJax: MathJaxParser | undefined;
let mathjax: Promise<MathJaxParser> | undefined;

async function getMathJax() {
  mathjax ??= (async () => {
    const [{ mathjax }, { TeX }, { liteAdaptor }, { RegisterHTMLHandler }, _allPackages, { STATE }] = await Promise.all([
      import("mathjax-full/js/mathjax.js"),
      import("mathjax-full/js/input/tex.js"),
      import("mathjax-full/js/adaptors/liteAdaptor.js"),
      import("mathjax-full/js/handlers/html.js"),
      import("mathjax-full/js/input/tex/AllPackages.js"),
      import("mathjax-full/js/core/MathItem.js"),
    ]);

    RegisterHTMLHandler(liteAdaptor());
    const tex = new TeX({ packages: ["base", "ams", "newcommand"] });
    const document = mathjax.document("", { InputJax: tex });
    initializedMathJax = {
      texToMml: (source: string, options: { display: boolean }) =>
        document.convert(source, { ...options, end: STATE.COMPILED }),
    };
    return initializedMathJax;
  })();
  return mathjax;
}

export async function renderTexToTextCells(
  tex: string,
  display: boolean,
  maxWidthCells = display ? 60 : 40,
  maxHeightCells = display ? 8 : 5,
): Promise<string[]> {
  await getMathJax();
  return renderTexToTextCellsSync(tex, display, maxWidthCells, maxHeightCells);
}

function renderTexToTextCellsSync(
  tex: string,
  display: boolean,
  maxWidthCells = display ? 60 : 40,
  maxHeightCells = display ? 8 : 5,
): string[] {
  const width = Math.max(1, Math.floor(maxWidthCells));
  const height = Math.max(1, Math.floor(maxHeightCells));
  if (!initializedMathJax) throw new Error("MathJax is not initialized");
  return renderTexToCellBoxSync(tex, display, width, height).lines;
}

function renderTexToCellBoxSync(
  tex: string,
  display: boolean,
  maxWidthCells: number,
  maxHeightCells: number,
): CellBox {
  const width = Math.max(1, Math.floor(maxWidthCells));
  const height = Math.max(1, Math.floor(maxHeightCells));
  if (!initializedMathJax) throw new Error("MathJax is not initialized");
  const box = layoutMathTreeBox(initializedMathJax.texToMml(tex, { display }));
  if (box.lines.length === 0) throw new Error("MathJax produced an empty layout");
  if (box.lines.length <= height && box.lines.every((line) => line.length <= width)) return box;

  const fallback = wrapTextCells(latexToTerminalText(tex), width);
  if (fallback.length > height) throw new Error("Math layout exceeds terminal preview bounds");
  return { lines: fallback, baseline: 0 };
}

const ASSISTANT_PATCH = Symbol.for("terminal-math:assistant-transcript-renderer");

interface AssistantRenderState {
  display: "source" | "math";
  message: any;
  processed: boolean;
  themeVersion?: number;
  width?: number;
}

interface AssistantPatch {
  originalRender: (this: AssistantMessageComponent, width: number) => string[];
  originalUpdateContent: (this: AssistantMessageComponent, message: any) => void;
  states: Map<AssistantMessageComponent, AssistantRenderState>;
  owners: Set<symbol>;
  enabledOwners: Set<symbol>;
  wrapperRender: (this: AssistantMessageComponent, width: number) => string[];
  wrapperUpdateContent: (this: AssistantMessageComponent, message: any) => void;
  theme?: { fg(color: "text", text: string): string };
  themeVersion: number;
}

function installAssistantTranscriptRenderer(owner: symbol): void {
  const globalState = globalThis as typeof globalThis & { [ASSISTANT_PATCH]?: AssistantPatch };
  const existing = globalState[ASSISTANT_PATCH];
  if (existing) {
    existing.owners.add(owner);
    return;
  }

  const prototype = AssistantMessageComponent.prototype;
  const patch = {
    originalRender: prototype.render,
    originalUpdateContent: prototype.updateContent,
    states: new Map<AssistantMessageComponent, AssistantRenderState>(),
    owners: new Set([owner]),
    enabledOwners: new Set<symbol>(),
    themeVersion: 0,
  } as AssistantPatch;

  patch.wrapperUpdateContent = function updateContent(message: any): void {
    patch.states.set(this, { display: "source", message, processed: false });
    patch.originalUpdateContent.call(this, message);
  };
  patch.wrapperRender = function render(width: number): string[] {
    const state = patch.states.get(this);
    if (!state) return patch.originalRender.call(this, width);

    if (patch.enabledOwners.size === 0) {
      if (state.display === "math") {
        patch.originalUpdateContent.call(this, state.message);
        (this as any).lastMessage = state.message;
        state.display = "source";
      }
      state.processed = false;
      return patch.originalRender.call(this, width);
    }

    if (state.processed && state.width === width && state.themeVersion === patch.themeVersion) {
      return patch.originalRender.call(this, width);
    }

    try {
      const displayMessage = cloneMessageForDisplay(state.message, width, patch.theme);
      if (displayMessage !== state.message) {
        patch.originalUpdateContent.call(this, displayMessage);
        (this as any).lastMessage = state.message;
        state.display = "math";
      } else if (state.display === "math") {
        patch.originalUpdateContent.call(this, state.message);
        (this as any).lastMessage = state.message;
        state.display = "source";
      }
      state.processed = true;
      state.width = width;
      state.themeVersion = patch.themeVersion;
      return patch.originalRender.call(this, width);
    } catch {
      patch.originalUpdateContent.call(this, state.message);
      (this as any).lastMessage = state.message;
      state.display = "source";
      state.processed = true;
      state.width = width;
      state.themeVersion = patch.themeVersion;
      return patch.originalRender.call(this, width);
    }
  };
  prototype.updateContent = patch.wrapperUpdateContent;
  prototype.render = patch.wrapperRender;
  globalState[ASSISTANT_PATCH] = patch;
}

function enableAssistantTranscriptRenderer(owner: symbol): void {
  const globalState = globalThis as typeof globalThis & { [ASSISTANT_PATCH]?: AssistantPatch };
  globalState[ASSISTANT_PATCH]?.enabledOwners.add(owner);
}

function disableAssistantTranscriptRenderer(owner: symbol): void {
  const globalState = globalThis as typeof globalThis & { [ASSISTANT_PATCH]?: AssistantPatch };
  globalState[ASSISTANT_PATCH]?.enabledOwners.delete(owner);
}

function uninstallAssistantTranscriptRenderer(owner: symbol): void {
  const globalState = globalThis as typeof globalThis & { [ASSISTANT_PATCH]?: AssistantPatch };
  const patch = globalState[ASSISTANT_PATCH];
  if (!patch) return;
  patch.owners.delete(owner);
  patch.enabledOwners.delete(owner);
  if (patch.owners.size > 0) return;

  for (const [component, state] of patch.states) {
    if (state.display !== "math") continue;
    patch.originalUpdateContent.call(component, state.message);
    (component as any).lastMessage = state.message;
  }

  const prototype = AssistantMessageComponent.prototype;
  const ownsRender = prototype.render === patch.wrapperRender;
  const ownsUpdate = prototype.updateContent === patch.wrapperUpdateContent;
  if (ownsRender) prototype.render = patch.originalRender;
  if (ownsUpdate) prototype.updateContent = patch.originalUpdateContent;
  if (ownsRender && ownsUpdate) delete globalState[ASSISTANT_PATCH];
}

function setAssistantTranscriptTheme(theme: { fg(color: "text", text: string): string }): void {
  const globalState = globalThis as typeof globalThis & { [ASSISTANT_PATCH]?: AssistantPatch };
  const patch = globalState[ASSISTANT_PATCH];
  if (!patch || patch.theme === theme) return;
  patch.theme = theme;
  patch.themeVersion += 1;
}

function cloneMessageForDisplay(
  message: any,
  width: number,
  theme?: { fg(color: "text", text: string): string },
): any {
  let changed = false;
  const content = message.content.map((part: any) => {
    if (part?.type !== "text") return part;
    const text = renderMathInText(part.text, Math.max(1, width - 2), theme);
    if (text === part.text) return part;
    changed = true;
    return { ...part, text };
  });
  return changed ? { ...message, content } : message;
}

function renderMathInText(
  text: string,
  width: number,
  theme?: { fg(color: "text", text: string): string },
): string {
  const blocks = detectMathBlocks(text);
  let rendered = text;
  for (let index = blocks.length - 1; index >= 0; index -= 1) {
    const block = blocks[index]!;
    const box = renderTexToCellBoxSync(
      block.tex,
      block.display,
      Math.min(width, block.display ? 60 : 40),
      block.display ? 8 : 5,
    );
    const cells = box.lines
      .map((line) => line.replaceAll(" ", "\u00a0"))
      .map((line) => theme?.fg("text", line) ?? line);

    if (!block.display && cells.length > 1) {
      const lineStart = rendered.lastIndexOf("\n", block.range[0] - 1) + 1;
      const newline = rendered.indexOf("\n", block.range[1]);
      const lineEnd = newline === -1 ? rendered.length : newline;
      const prefix = rendered.slice(lineStart, block.range[0]);
      const suffix = rendered.slice(block.range[1], lineEnd);
      const composed = cells
        .map((line, row) =>
          row === box.baseline
            ? `${prefix}${line}${suffix}`
            : `${"\u00a0".repeat(prefix.length)}${line}`,
        )
        .join("\n");
      rendered = rendered.slice(0, lineStart) + composed + rendered.slice(lineEnd);
      continue;
    }

    const layout = cells.join("\n");
    const replacement = block.display ? `\n\n${layout}\n\n` : layout;
    rendered = rendered.slice(0, block.range[0]) + replacement + rendered.slice(block.range[1]);
  }
  return rendered;
}