import { execFileSync } from "node:child_process";
import { copyFileSync, existsSync, readdirSync, readFileSync, realpathSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { CONFIG_DIR, THEME_NAME } from "./config.ts";

const REL = "dist/modes/interactive/components/assistant-message.js";
const REL_USER = "dist/modes/interactive/components/user-message.js";

const FLAG_ANCHOR = "        const hasVisibleContent = message.content.some(";
const MD_ANCHOR = `                this.contentContainer.addChild(new Markdown(content.text.trim(), this.outputPad, 0, this.markdownTheme, undefined, {
                    transform: createMarkdownTransform("assistant", this.isStreaming, this.markdownTransformers),
                }));`;

/**
 * Only a *finalized* message with no tool call is a final answer.
 *
 * `stopReason` is unreliable as a "is it done" signal: providers can populate it
 * on a message that is still streaming, which made intermediate text briefly
 * flash with the final-answer background. Upstream's own `isStreaming` flag is
 * the honest signal - it is true between message_start and message_end and
 * false for history/replay renders.
 */
const FLAG_INJECT =
  '        const isIntermediate = this.isStreaming || message.stopReason === "aborted" || message.stopReason === "error" || message.content.some((c) => c.type === "toolCall");\n';

// Runtime helpers injected into the component module. The patched code runs
// inside pi's own bundle and cannot import from this extension, so the config
// is read from disk (cached for 1s to avoid per-render file reads).
const HELPER_ANCHOR = "export class AssistantMessageComponent extends Container {";
const HELPERS = `import { createRequire as __mbsCreateRequire } from "node:module";
const __mbsRequire = __mbsCreateRequire(import.meta.url);
const __mbsConfigPath = __mbsRequire("node:path").join(__mbsRequire("node:os").homedir(), ${JSON.stringify(CONFIG_DIR)}, "material-black-simple.json");
let __mbsCache = { at: 0, value: { overlay: false } };
function __mbsConfig() {
    const now = Date.now();
    if (now - __mbsCache.at < 1000) return __mbsCache.value;
    let value = { overlay: false };
    try {
        value = { overlay: false, ...JSON.parse(__mbsRequire("node:fs").readFileSync(__mbsConfigPath, "utf8")) };
    }
    catch { }
    __mbsCache = { at: now, value };
    return value;
}
/** Styling applies to material_black_simple, or to any theme when overlay is on. */
function __mbsActive() {
    const cfg = __mbsConfig();
    return cfg.overlay || theme.name === "${THEME_NAME}";
}
function __mbsHex(hex) {
    const n = parseInt(hex.slice(1), 16);
    return [(n >> 16) & 255, (n >> 8) & 255, n & 255];
}
function __mbsFg(override, fallbackKey) {
    if (override)
        return (text) => \`\\x1b[38;2;\${__mbsHex(override).join(";")}m\${text}\\x1b[39m\`;
    return (text) => theme.fg(fallbackKey, text);
}
function __mbsBg(override, fallbackKey) {
    if (override)
        return (text) => \`\\x1b[48;2;\${__mbsHex(override).join(";")}m\${text}\\x1b[49m\`;
    return fallbackKey ? (text) => theme.bg(fallbackKey, text) : undefined;
}
`;

const HELPER_INJECT = `${HELPERS}${HELPER_ANCHOR}`;

// user-message.js: same helpers, plus overrides for the user bubble.
const USER_ANCHOR = "export class UserMessageComponent extends Container {";
const USER_HELPER_INJECT = `${HELPERS}${USER_ANCHOR}`;
const USER_BOX_ANCHOR = `        const contentBox = new Box(this.outputPad, 1, (content) => theme.bg("userMessageBg", content));
        contentBox.addChild(new Markdown(this.text, 0, 0, this.markdownTheme, {
            color: (content) => theme.fg("userMessageText", content),
        }, {`;
const USER_BOX_INJECT = `        const __mbsCfg = __mbsConfig();
        const __mbsOn = __mbsActive();
        const contentBox = new Box(this.outputPad, 1, __mbsOn && __mbsCfg.inputBg
            ? __mbsBg(__mbsCfg.inputBg, "userMessageBg")
            : (content) => theme.bg("userMessageBg", content));
        contentBox.addChild(new Markdown(this.text, 0, 0, this.markdownTheme, {
            color: __mbsOn && __mbsCfg.inputFontColor
                ? __mbsFg(__mbsCfg.inputFontColor, "userMessageText")
                : (content) => theme.fg("userMessageText", content),
        }, {`;

const MD_REPLACEMENT = `                const __mbsCfg = __mbsConfig();
                this.contentContainer.addChild(new Markdown(content.text.trim(), this.outputPad, 0, this.markdownTheme, !__mbsActive()
                    ? undefined
                    : isIntermediate
                        ? {
                            color: __mbsFg(__mbsCfg.intermediateFontColor, "muted"),
                            bgColor: __mbsBg(__mbsCfg.intermediateBg, undefined),
                        }
                        : {
                            color: __mbsFg(__mbsCfg.finalFontColor, "text"),
                            bgColor: __mbsBg(__mbsCfg.finalOutputBg, "selectedBg"),
                        }, {
                    transform: createMarkdownTransform("assistant", this.isStreaming, this.markdownTransformers),
                }));`;

export type PatchOutcome =
  | { status: "patched"; file: string }
  | { status: "already"; file: string }
  | { status: "failed"; file: string; reason: string };

// Locate the running CLI, PATH install, and local copies.
export function findTargets(): string[] {
  const roots = new Set<string>();

  const entries = process.argv[1] ? [process.argv[1]] : [];
  try {
    entries.push(execFileSync("which", ["pi"], { encoding: "utf8" }).trim());
  } catch {
    // The running CLI can still be found when pi is absent from PATH.
  }

  for (const entry of entries) {
    let bin: string;
    try {
      bin = realpathSync(entry);
    } catch {
      continue;
    }
    for (let dir = dirname(bin); ; dir = dirname(dir)) {
      if (existsSync(join(dir, REL))) {
        roots.add(dir);
        break;
      }
      if (dirname(dir) === dir) break;
    }
    // Homebrew can use a wrapper outside the package directory.
    const pkg = join(dirname(dirname(bin)), "libexec/lib/node_modules/@earendil-works/pi-coding-agent");
    if (existsSync(join(pkg, REL))) roots.add(pkg);
  }

  for (const dir of ["agent", "cc"]) {
    const pkg = join(homedir(), CONFIG_DIR, dir, "node_modules/@earendil-works/pi-coding-agent");
    if (existsSync(join(pkg, REL))) roots.add(pkg);
  }

  return [...roots].flatMap((root) => {
    const chunks = join(root, "dist/bundle/chunks");
    const bundled = existsSync(chunks)
      ? readdirSync(chunks).filter((name) => name.endsWith(".js"))
        .map((name) => join(chunks, name))
        .filter((file) => readFileSync(file, "utf8").includes("AssistantMessageComponent=class extends Container{"))
      : [];
    return [join(root, REL), ...bundled];
  });
}

const BUNDLE_MD_ANCHOR = 'new Markdown(content.text.trim(),this.outputPad,0,this.markdownTheme,void 0,{transform:createMarkdownTransform("assistant",this.isStreaming,this.markdownTransformers)})';
const BUNDLE_USER_BG_ANCHOR = 'new Box(this.outputPad,1,content=>theme.bg("userMessageBg",content))';
const BUNDLE_USER_FG_ANCHOR = 'new Markdown(this.text,0,0,this.markdownTheme,{color:content=>theme.fg("userMessageText",content)}';
const BUNDLE_HELPERS = `${HELPERS}
function __mbsAssistantStyle(message, isStreaming) {
    if (!__mbsActive()) return undefined;
    const cfg = __mbsConfig();
    const isIntermediate = isStreaming || message.stopReason === "aborted" || message.stopReason === "error" || message.content.some((c) => c.type === "toolCall");
    return isIntermediate
        ? { color: __mbsFg(cfg.intermediateFontColor, "muted"), bgColor: __mbsBg(cfg.intermediateBg, undefined) }
        : { color: __mbsFg(cfg.finalFontColor, "text"), bgColor: __mbsBg(cfg.finalOutputBg, "selectedBg") };
}
`;

export function patchBundle(file: string): PatchOutcome {
  try {
    const src = readFileSync(file, "utf8");
    if (src.includes("function __mbsAssistantStyle(")) return { status: "already", file };
    for (const anchor of [BUNDLE_MD_ANCHOR, BUNDLE_USER_BG_ANCHOR, BUNDLE_USER_FG_ANCHOR]) {
      if (src.split(anchor).length !== 2) {
        return { status: "failed", file, reason: "expected one renderer anchor in runtime bundle (upstream changed)" };
      }
    }
    const next = BUNDLE_HELPERS + src
      .replace(BUNDLE_MD_ANCHOR, BUNDLE_MD_ANCHOR.replace("this.markdownTheme,void 0,", "this.markdownTheme,__mbsAssistantStyle(message,this.isStreaming),"))
      .replace(BUNDLE_USER_BG_ANCHOR, 'new Box(this.outputPad,1,__mbsActive()?__mbsBg(__mbsConfig().inputBg,"userMessageBg"):content=>theme.bg("userMessageBg",content))')
      .replace(BUNDLE_USER_FG_ANCHOR, 'new Markdown(this.text,0,0,this.markdownTheme,{color:__mbsActive()?__mbsFg(__mbsConfig().inputFontColor,"userMessageText"):content=>theme.fg("userMessageText",content)}');
    if (!existsSync(`${file}.orig`)) copyFileSync(file, `${file}.orig`);
    writeFileSync(file, next);
    return { status: "patched", file };
  } catch (error) {
    return { status: "failed", file, reason: error instanceof Error ? error.message : String(error) };
  }
}

export function patchFile(file: string): PatchOutcome {
  let src: string;
  try {
    src = readFileSync(file, "utf8");
  } catch (e) {
    return { status: "failed", file, reason: `unreadable: ${(e as Error).message}` };
  }

  if (src.includes("isIntermediate")) return { status: "already", file };

  for (const [name, anchor] of [
    ["class header", HELPER_ANCHOR],
    ["hasVisibleContent", FLAG_ANCHOR],
    ["Markdown call", MD_ANCHOR],
  ] as const) {
    const n = src.split(anchor).length - 1;
    if (n !== 1) {
      return { status: "failed", file, reason: `expected 1 "${name}" anchor, found ${n} (upstream changed)` };
    }
  }

  const next = src
    .replace(HELPER_ANCHOR, HELPER_INJECT)
    .replace(FLAG_ANCHOR, FLAG_INJECT + FLAG_ANCHOR)
    .replace(MD_ANCHOR, MD_REPLACEMENT);

  try {
    if (!existsSync(`${file}.orig`)) copyFileSync(file, `${file}.orig`);
    writeFileSync(file, next);
  } catch (e) {
    return { status: "failed", file, reason: `not writable: ${(e as Error).message}` };
  }

  return { status: "patched", file };
}

/** Patch user-message.js so input bg/font overrides apply. */
function patchUser(file: string): PatchOutcome {
  let src: string;
  try {
    src = readFileSync(file, "utf8");
  } catch (e) {
    return { status: "failed", file, reason: `unreadable: ${(e as Error).message}` };
  }

  if (src.includes("__mbsConfig")) return { status: "already", file };

  for (const [name, anchor] of [
    ["class header", USER_ANCHOR],
    ["content box", USER_BOX_ANCHOR],
  ] as const) {
    const n = src.split(anchor).length - 1;
    if (n !== 1) {
      return { status: "failed", file, reason: `expected 1 "${name}" anchor, found ${n} (upstream changed)` };
    }
  }

  const next = src.replace(USER_ANCHOR, USER_HELPER_INJECT).replace(USER_BOX_ANCHOR, USER_BOX_INJECT);

  try {
    if (!existsSync(`${file}.orig`)) copyFileSync(file, `${file}.orig`);
    writeFileSync(file, next);
  } catch (e) {
    return { status: "failed", file, reason: `not writable: ${(e as Error).message}` };
  }

  return { status: "patched", file };
}

export function applyPatch(): PatchOutcome[] {
  return findTargets().flatMap((file) => file.endsWith(REL)
    ? [patchFile(file), patchUser(file.replace(REL, REL_USER))]
    : [patchBundle(file)]);
}

export function revertPatch(): PatchOutcome[] {
  return findTargets()
    .flatMap((f) => f.endsWith(REL) ? [f, f.replace(REL, REL_USER)] : [f])
    .map((file) => {
    const orig = `${file}.orig`;
    if (!existsSync(orig)) return { status: "failed", file, reason: "no .orig backup found" };
    try {
      copyFileSync(orig, file);
      return { status: "patched", file };
    } catch (e) {
      return { status: "failed", file, reason: (e as Error).message };
    }
  });
}