/**
 * @inobit/pi-undo — 撤销：把最近一次发送的输入撤回到输入框并从对话中移除
 * 单次/轮，原子 abort，快捷键 alt+u（委托 /undo 命令管道）
 *
 * 队列非空（执行中）：等价官方 dequeue（alt+up）——取回全部排队文本到编辑器、
 * 清空 steer/followUp 队列，不中断当前轮；实现见 capturedTui/findCustomEditor。
 */

import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { findLastUserEntry } from "./history.ts";
import { loadConfig, SHORTCUT as DEFAULT_SHORTCUT } from "./config.ts";

type SessionId = string;
const SHORTCUT = DEFAULT_SHORTCUT;

// 硬撤销成功后追加的哨兵条目类型：不进 LLM 上下文、TUI 不渲染，
// 仅用于把分支终点固定到磁盘（resume 按「文件最后一条 entry」重建 leaf）
const UNDO_PIN_ENTRY = "pi-undo-pin";

// abort 后等 idle 的总预算与轮询节拍（原 50*200ms=10s 魔法数收敛）
const WAIT_MS = 3000;
const POLL_MS = 50;
const sleep = (ms: number): Promise<void> => new Promise((r) => setTimeout(r, ms));

function sessionIdOf(ctx: ExtensionContext): SessionId {
  try {
    return ctx.sessionManager.getSessionId() ?? "__global__";
  } catch {
    return "__global__";
  }
}

/**
 * 等待 agent 回到 idle。
 * - 有 waitForIdle（/undo 命令上下文）：事件驱动 + WAIT_MS 超时兜底
 * - 无 waitForIdle（受限上下文）：deadline 轮询，每 POLL_MS 检一次
 */
async function waitUntilIdle(ctx: ExtensionContext): Promise<void> {
  const maybe = ctx as unknown as { waitForIdle?: () => Promise<void> };
  if (typeof maybe.waitForIdle === "function") {
    await Promise.race([maybe.waitForIdle(), sleep(WAIT_MS)]);
    return;
  }
  const deadline = Date.now() + WAIT_MS;
  while (Date.now() < deadline) {
    const isIdle = typeof (ctx as unknown as { isIdle?: () => boolean }).isIdle === "function"
      ? (ctx as unknown as { isIdle: () => boolean }).isIdle()
      : true;
    if (isIdle) return;
    await sleep(POLL_MS);
  }
}

/**
 * 宿主 TUI 引用与 CustomEditor 定位（队列取回专用）。
 *
 * 背景：官方 dequeue（alt+up）绑定在 TUI 内部的 CustomEditor 上
 * （editor.onAction("app.message.dequeue", ...)），扩展 API 没有编程式入口。
 * 这里通过 setWidget 的组件工厂拿到 TUI 引用，撤销时在组件树中找到
 * CustomEditor 并直接调用其 actionHandlers 里的 dequeue handler——
 * 与按下 alt+up 执行的是同一个函数对象，语义逐字节一致。
 */
const capturedTui: { current: unknown } = { current: undefined };

interface EditorLike {
  actionHandlers?: Map<string, () => void>;
}

/** 从 TUI 组件树递归查找带 actionHandlers 的 CustomEditor */
function findCustomEditor(node: unknown, depth = 0): EditorLike | undefined {
  if (!node || typeof node !== "object" || depth > 6) return undefined;
  const n = node as Record<string, unknown>;
  // 精确匹配：只命中注册了官方 dequeue 的编辑器，避免诱饵 Map 误命中
  if (n.actionHandlers instanceof Map && (n.actionHandlers as Map<string, unknown>).has("app.message.dequeue")) return n as EditorLike;
  const children = n.children;
  if (Array.isArray(children)) {
    for (const child of children) {
      const found = findCustomEditor(child, depth + 1);
      if (found) return found;
    }
  }
  return undefined;
}

/** 注册零高 widget 以捕获 TUI 引用（render 返回空数组，不影响布局） */
function registerTuiCapture(ctx: ExtensionContext): void {
  try {
    ctx.ui.setWidget("pi-undo-tui-capture", (tui: unknown) => {
      capturedTui.current = tui;
      return {
        render: (_width: number): string[] => [],
        invalidate: (): void => {},
        dispose: (): void => {},
      };
    });
  } catch {}
}

export default function (pi: ExtensionAPI) {
  const canUndoBySid = new Map<SessionId, boolean>();
  const pendingBySid = new Set<SessionId>();

  const getCanUndo = (sid: SessionId): boolean => (canUndoBySid.has(sid) ? canUndoBySid.get(sid)! : true);
  const setCanUndo = (sid: SessionId, v: boolean): void => { canUndoBySid.set(sid, v); };

  const doUndo = async (ctx: ExtensionContext): Promise<void> => {
    if (!ctx.hasUI) return;
    const sid = sessionIdOf(ctx);

    const safeNotify = (msg: string, level: "info" | "warning" | "error" = "warning"): void => {
      try { ctx.ui.notify(msg, level); } catch {}
    };

    if (pendingBySid.has(sid)) {
      safeNotify("Undo already in progress, try again", "warning");
      return;
    }
    if (!getCanUndo(sid)) {
      safeNotify("Already undone for this turn. Send a new message to undo again.", "warning");
      return;
    }
    let draft = "";
    try { draft = ctx.ui.getEditorText() ?? ""; } catch {}
    if (draft.trim() !== "") {
      safeNotify("Editor has draft, clear it first", "warning");
      return;
    }

    // 撤销路径：加互斥，防止并发绕过单次/轮。
    // 注意：不做镜像软撤——扩展拿不到宿主队列内容，pop 副本无法真正移除排队的
    // steer/followUp 消息，只会造成「假撤回」。队列非空时走官方 dequeue 直调
    // （见下方 hasPending 分支）；队列空时执行中 abort 后硬撤刚发的 user 消息。
    pendingBySid.add(sid);
    try {
      const isIdleFn = typeof (ctx as unknown as { isIdle?: () => boolean }).isIdle === "function"
        ? () => (ctx as unknown as { isIdle: () => boolean }).isIdle()
        : () => true;
      const abortFn = typeof (ctx as unknown as { abort?: () => void }).abort === "function"
        ? () => (ctx as unknown as { abort: () => void }).abort()
        : () => {};

      // 队列非空时等价于官方 dequeue（alt+up）：取回全部排队文本到编辑器并清空
      // steer/followUp 队列，不中断当前轮、不动会话树。通过 TUI 引用直达
      // CustomEditor.actionHandlers 中注册的官方 handler，避免模拟按键序列。
      try {
        const hasPendingFn = (ctx as unknown as { hasPendingMessages?: () => boolean }).hasPendingMessages;
        if (typeof hasPendingFn === "function") {
          let hasPending: boolean;
          try {
            hasPending = hasPendingFn.call(ctx);
          } catch (e) {
            // 无法判定队列状态时保守返回，不坠入 abort 分支
            safeNotify(`Cannot determine queue state: ${e instanceof Error ? e.message : String(e)} \u2014 try again or use the dequeue shortcut (alt+up; alt+q on Windows)`, "warning");
            return;
          }
          if (hasPending) {
            const tui = capturedTui.current;
            // TUI 引用缺失：单独提示
            if (tui === undefined) {
              safeNotify("Host TUI reference unavailable \u2014 press the dequeue shortcut (alt+up; alt+q on Windows) to recall queued messages", "warning");
              return;
            }
            const editor = findCustomEditor(tui);
            const dequeueHandler =
              editor?.actionHandlers instanceof Map
                ? editor.actionHandlers.get("app.message.dequeue")
                : undefined;
            if (typeof dequeueHandler === "function") {
              try { dequeueHandler(); } catch (e) {
                safeNotify(`Failed to recall queued messages: ${e instanceof Error ? e.message : String(e)}`, "warning");
              }
            } else {
              // 宿主结构漂移：未找到编辑器或 dequeue action 已改名
              safeNotify("Host editor structure changed \u2014 press the dequeue shortcut (alt+up; alt+q on Windows) to recall queued messages", "warning");
            }
            return;
          }
        }
      } catch {}

      if (!isIdleFn()) {
        try { abortFn(); } catch {}
        await waitUntilIdle(ctx);
        if (!isIdleFn()) {
          safeNotify("Abort did not settle, try again", "warning");
          return;
        }
        // 草稿守卫优先级最高、适用于所有分支：abort 后编辑器非空（含异常宿主把
        // 排队消息回灌进编辑器的情况）一律视为草稿——提示后直接返回，不硬撤、
        // 不合并回填；未消耗单次/轮，用户清空草稿后可重试。
        let d2 = "";
        try { d2 = ctx.ui.getEditorText() ?? ""; } catch {}
        if (d2.trim() !== "") {
          safeNotify("Editor has draft, clear it first", "warning");
          return;
        }
      }

      let branch: readonly unknown[] = [];
      try { branch = ctx.sessionManager.getBranch() as readonly unknown[]; } catch { branch = []; }
      const found = findLastUserEntry(branch as never);
      if (!found) {
        safeNotify("No message to undo", "warning");
        return;
      }

      // 回填编辑器：仅写回被撤消息文本（草稿守卫已在上方拦截一切非空场景）
      const fillEditorWithUndoText = (): void => {
        try { ctx.ui.setEditorText(found.text); } catch {}
      };

      const anyCtx = ctx as unknown as Record<string, unknown>;
      const sm = ctx.sessionManager as unknown as Record<string, unknown>;

      // 读取当前 leaf 指针；能力缺失或异常时返回 undefined（无法观测移动，退化为仅信任结果标志）
      const getLeafIdSafe = (): string | null | undefined => {
        try {
          const fn = (ctx.sessionManager as unknown as { getLeafId?: () => string | null }).getLeafId;
          return typeof fn === "function" ? fn.call(ctx.sessionManager) : undefined;
        } catch {
          return undefined;
        }
      };

      // 哨兵落盘：把生效分支的终点写进文件末尾，否则 resume 会按旧末条重建 leaf 导致复活
      const appendPinSafely = (): void => {
        try {
          pi.appendEntry(UNDO_PIN_ENTRY, { v: 1, at: found.parentId ?? null });
        } catch {}
      };

      try {
        if (typeof anyCtx.navigateTree !== "function") {
          // 受限上下文（如部分宿主版本的快捷键环境）没有导航能力：
          // 显式失败优于静默降级——sm.branch() 只动内存指针，UI/context/磁盘三者都不会更新
          safeNotify("Hard undo unavailable here — use /undo instead", "warning");
          return;
        }

        const leafBefore = getLeafIdSafe();
        let target = found.entryId;
        if (leafBefore !== undefined && leafBefore === found.entryId) {
          // no-op 陷阱：目标恰为当前 leaf 时 navigateTree 会静默早退（返回 cancelled:false 但什么都不做）。
          // 典型场景：崩溃 resume 后文件末条恰为该 user 消息。改以 parent 为目标触发真实导航。
          if (found.parentId != null) {
            target = found.parentId;
          } else {
            // 特例：首条消息本身是 leaf，没有任何 entry 可作导航目标。
            // 语义：resetLeaf + 根级哨兵保证磁盘正确；进程内视图无法在此重建，显式告知用户。
            if (typeof sm.resetLeaf !== "function") throw new Error("No hard revert capability");
            (sm.resetLeaf as () => void)();
            appendPinSafely();
            fillEditorWithUndoText();
            try { ctx.ui.setStatus("pi-undo", " "); ctx.ui.setStatus("pi-undo", undefined); } catch {}
            setCanUndo(sid, false);
            safeNotify("Undone on disk. In-memory view not refreshed — restart or /new to apply.", "warning");
            return;
          }
        }

        const result = await (anyCtx.navigateTree as (id: string, opts: unknown) => Promise<unknown>)(
          target,
          { summarize: false },
        );
        // 结果字段不可信（no-op 早退同样返回 cancelled:false），以 leaf 是否实际移动为准；
        // getLeafId 不可用时退化为仅信任 cancelled 标志
        const cancelled = (result as { cancelled?: boolean } | undefined)?.cancelled === true;
        const moved = leafBefore === undefined
          ? !cancelled
          : getLeafIdSafe() !== leafBefore;
        if (cancelled || !moved) {
          // 导航未生效：会话保持原状。不回填编辑器、不消耗单次/轮、绝不落哨兵
          // （此时落哨兵会把待撤销消息钉回生效路径，主动制造 resume 复活）
          safeNotify("Undo did not take effect, try again", "warning");
          return;
        }

        // navigateTree 已在内部 setEditorText（当编辑器空时）；此处显式回填，
        // 并保留宿主 abort 时回灌的排队文本
        fillEditorWithUndoText();
        try { ctx.ui.setStatus("pi-undo", " "); ctx.ui.setStatus("pi-undo", undefined); } catch {}
        appendPinSafely();
        setCanUndo(sid, false);
        return;
      } catch (e) {
        fillEditorWithUndoText();
        setCanUndo(sid, false);
        safeNotify(`Hard undo failed: ${e instanceof Error ? e.message : String(e)}`, "warning");
        return;
      }
    } finally {
      pendingBySid.delete(sid);
    }
  };

  pi.on("before_agent_start", async (_e, ctx) => {
    setCanUndo(sessionIdOf(ctx), true);
  });

  pi.on("session_start", async (_e, ctx) => {
    const sid = sessionIdOf(ctx);
    setCanUndo(sid, true);
    pendingBySid.delete(sid);
    if (ctx.hasUI && ctx.mode === "tui") registerTuiCapture(ctx);
  });
  pi.on("session_shutdown", async (_e, ctx) => {
    const sid = sessionIdOf(ctx);
    canUndoBySid.delete(sid);
    pendingBySid.delete(sid);
    // 生命周期清理：释放捕获的 TUI 引用，避免跨会话残留
    capturedTui.current = undefined;
  });

  // 快捷键可配：~/.pi/agent/extensions/pi-undo/config.json {"shortcut":"alt+u"}，需 /reload
  let shortcut: string = SHORTCUT;
  try {
    const cfg = loadConfig(process.cwd());
    if (cfg.shortcut) shortcut = cfg.shortcut;
  } catch {}

  pi.registerCommand("undo", {
    description: `Undo last prompt to editor (hard revert, single per turn; recalls queued messages first). Shortcut: ${shortcut}`,
    handler: async (_a, ctx) => { await doUndo(ctx); },
  });

  try {
    pi.registerShortcut(shortcut as unknown as import("@earendil-works/pi-tui").KeyId, {
      description: "Undo last prompt to editor (delegates to /undo)",
      handler: async () => {
        // 宿主给快捷键注入的上下文缺少 navigateTree 等 session-control 能力（类型层却同为
        // ExtensionContext），直接跑 doUndo 只能拿到精简 ctx。委托 /undo 走命令派发管道后
        // 与手敲 /undo 字面上同一条执行路径，行为一致性由构造保证。
        // expandPromptTemplates 必须显式 true（宿主默认 false）：命中扩展命令后立即返回，不触发 LLM turn。
        pi.sendUserMessage("/undo", { expandPromptTemplates: true });
      },
    });
  } catch (e) {
    try { console.warn(`[pi-undo] shortcut ${shortcut} failed: ${String(e)}`); } catch {}
  }
}