// parallel-batch.ts —— 并行批观察 · 正式版
// 零依赖，无需安装。运行:
//   pi -e /path/to/parallel-batch.ts
// 或放入 ~/.pi/agent/extensions/ 全局生效（任意目录开 pi 都有）
//
// 命令:
//   /pb-mode notify|widget   切换显示模式（选择持久化，重开 pi 保持）
//   /pb-demo                 模拟两批并行工具，用于自测（勿与真实执行同时跑）
//
// 模式:
//   widget: 输入框上方卡片，逐工具 ✓/⟳ 实时状态，批完成自动消失
//   notify: 消息流底部一行灰字，批开始/完成各一条（会残留在消息流中）
//
// 真实行为: 模型在一轮回复中并行调用多个工具时，按当前模式显示批次。
// 批的边界 = 一次 LLM 响应发出的一批工具调用。

import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { truncateToWidth } from "@earendil-works/pi-tui";

const CONF_PATH = join(homedir(), ".pi", "agent", "parallel-batch.json");
const WIDGET_ID = "pb";
const MIN_WIDGET_VISIBLE_MS = 1_000;

type Mode = "widget" | "notify";

type ToolItem = {
  id: string;
  name: string;
  args: string;
  status: "running" | "done";
  isError: boolean;
};
type Batch = {
  id: number;
  tools: ToolItem[];
  startTs: number;
};

// ── 状态 ────────────────────────────────────────────────────

let currentBatch: Batch | null = null;
let nextId = 1;
let tui: any = null;
let theme: any = null;
let lastCtx: any = null;
let mode: Mode;
let onBatchStart: ((b: Batch) => void) | null = null;
let onBatchDone: ((b: Batch) => void) | null = null;
let clearBatchTimer: ReturnType<typeof setTimeout> | null = null;

const sleep = (ms: number) => new Promise<void>((r) => setTimeout(r, ms));
const pad = (s: string, n: number) => (s + " ".repeat(n)).slice(0, n);
const trunc = (s: string, n: number) =>
  s.length <= n ? s : s.slice(0, n - 1) + "…";

// ── 配置持久化 ──────────────────────────────────────────────

function loadConfig(): Mode {
  try {
    const m = JSON.parse(readFileSync(CONF_PATH, "utf8")).mode;
    return m === "notify" ? "notify" : "widget";
  } catch {
    return "widget";
  }
}

function saveConfig(m: Mode) {
  try {
    mkdirSync(dirname(CONF_PATH), { recursive: true });
    writeFileSync(CONF_PATH, JSON.stringify({ mode: m }, null, 2));
  } catch {
    /* 配置写失败不影响运行 */
  }
}

// ── 批次状态机 ──────────────────────────────────────────────
// 规则: 只要当前批还有工具在跑，新工具就归入该批;
// 当前批全部结束后，下一个工具自动开启新批（批 = 一次 LLM 响应的一批调用）。

function hasLiveTool(): boolean {
  return !!currentBatch && currentBatch.tools.some((t) => t.status === "running");
}

function openBatch(): Batch {
  if (hasLiveTool()) return currentBatch!;
  const b: Batch = { id: nextId++, tools: [], startTs: Date.now() };
  currentBatch = b;
  onBatchStart?.(b);
  return b;
}

function startTool(id: string, name: string, args: string) {
  openBatch().tools.push({ id, name, args, status: "running", isError: false });
  publish();
}

function finishTool(id: string, isError = false) {
  const b = currentBatch;
  if (!b) return;
  const t = b.tools.find((x) => x.id === id);
  if (!t) return;
  t.status = "done";
  t.isError = isError;
  publish();
  if (!hasLiveTool()) {
    onBatchDone?.(b);
    scheduleBatchClear(b);
  }
}

// 快速工具也要留下足够时间，让 TUI 至少完成一次可见渲染。
function scheduleBatchClear(b: Batch) {
  if (clearBatchTimer) clearTimeout(clearBatchTimer);
  const remaining = Math.max(0, MIN_WIDGET_VISIBLE_MS - (Date.now() - b.startTs));
  clearBatchTimer = setTimeout(() => {
    if (currentBatch === b && !hasLiveTool()) {
      currentBatch = null;
      if (mode === "widget" && lastCtx) clearWidget(lastCtx);
      publish();
    }
    clearBatchTimer = null;
  }, remaining);
}

function shortArgs(a: any): string {
  if (!a || typeof a !== "object") return "";
  for (const k of ["command", "path", "pattern", "url", "text"]) {
    if (typeof a[k] === "string") return trunc(a[k], 36);
  }
  try {
    return trunc(JSON.stringify(a), 36);
  } catch {
    return "";
  }
}

function namesOf(b: Batch): string {
  return b.tools.map((t) => trunc(t.name, 12)).join(" · ");
}

// ── widget 渲染（输入框上方卡片）────────────────────────────

function buildLines(width: number): string[] {
  if (!currentBatch) return []; // minimum display time elapsed → hide
  const top = currentBatch;
  const title = ` Batch #${top.id} ×${top.tools.length} `;
  const iconOf = (t: ToolItem) =>
    t.status === "done"
      ? t.isError
        ? theme?.fg("error", "✗")
        : theme?.fg("success", "✓")
      : theme?.fg("accent", "⟳");
  let maxArgs = 0;
  for (const t of top.tools) {
    maxArgs = Math.max(maxArgs, trunc(t.args, 36).length);
  }
  const innerW = 1 + 1 + 10 + 1 + maxArgs; // icon + sep + name + sep + args
  const rows = top.tools.map(
    (t) =>
      `│ ${iconOf(t)} ${pad(trunc(t.name, 10), 10)} ${theme?.fg(
        "dim",
        pad(trunc(t.args, 36), maxArgs)
      )} │`
  );
  const fill = Math.max(0, innerW + 1 - title.length);
  return [
    `${theme?.fg("muted", "╭─")}${theme?.fg("accent", title)}${theme?.fg(
      "muted",
      "─".repeat(fill) + "╮"
    )}`,
    ...rows,
    theme?.fg("muted", `╰${"─".repeat(innerW + 2)}╯`),
  ].map((line) => truncateToWidth(line, width, ""));
}

function publish() {
  if (tui) tui.requestRender();
}

function ensureWidget(ctx: any) {
  ctx.ui.setWidget(WIDGET_ID, (t: any, th: any) => {
    tui = t;
    theme = th;
    return { render: (width: number) => buildLines(width), invalidate: () => {} };
  });
}

function clearWidget(ctx: any) {
  ctx.ui.setWidget(WIDGET_ID, undefined);
}

// ── notify 渲染（消息流底部一行灰字）────────────────────────

function notifyStart(b: Batch) {
  lastCtx?.ui.notify(`⚡ Batch #${b.id} ×${b.tools.length}  ${namesOf(b)}`, "info");
}

function notifyDone(b: Batch) {
  const ok = !b.tools.some((t) => t.isError);
  lastCtx?.ui.notify(
    `${ok ? "✓" : "✗"} Batch #${b.id} ×${b.tools.length} ${ok ? "done" : "failed"} · ${(
      (Date.now() - b.startTs) /
      1000
    ).toFixed(1)}s`,
    "info"
  );
}

// ── 模式绑定 ────────────────────────────────────────────────

function bindRenderers(ctx: any) {
  if (mode === "widget") {
    ensureWidget(ctx);
    onBatchStart = null;
    onBatchDone = null;
  } else {
    clearWidget(ctx);
    lastCtx = ctx;
    onBatchStart = notifyStart;
    onBatchDone = notifyDone;
  }
}

function setMode(m: Mode, ctx: any) {
  mode = m;
  bindRenderers(ctx);
  saveConfig(m);
}

// ── 模拟波（/pb-demo 自测用）──────────────────────────────────

async function playWave() {
  startTool("a1", "ffgrep", '"parallel" src/');
  startTool("a2", "read", "src/foo.ts");
  startTool("a3", "fffind", "profile");
  await sleep(600);
  finishTool("a1");
  await sleep(900);
  finishTool("a2");
  await sleep(600);
  finishTool("a3");
  await sleep(1500);

  startTool("b1", "bash", "ls -la");
  startTool("b2", "edit", "README.md");
  await sleep(800);
  finishTool("b1");
  await sleep(800);
  finishTool("b2", true);
}

// ── 入口 ────────────────────────────────────────────────────

export default function (pi: any) {
  mode = loadConfig();

  pi.on("session_start", async (_e: any, ctx: any) => {
    lastCtx = ctx;
    bindRenderers(ctx); // 重开 pi 按持久化的模式恢复
  });
  pi.on("session_shutdown", async () => {
    if (clearBatchTimer) clearTimeout(clearBatchTimer);
    clearBatchTimer = null;
  });

  // 真实事件: 模型并行调用工具时实时显示
  pi.on("tool_execution_start", async (e: any, ctx: any) => {
    lastCtx = ctx;
    bindRenderers(ctx);
    startTool(e.toolCallId, e.toolName, shortArgs(e.args));
  });
  pi.on("tool_execution_end", async (e: any) => {
    finishTool(e.toolCallId, !!e.isError);
  });

  // 切换模式
  pi.registerCommand("pb-mode", {
    description: "Switch parallel-batch display mode: widget | notify",
    handler: async (args: string | undefined, ctx: any) => {
      const m = (args ?? "").trim() as Mode;
      if (m !== "widget" && m !== "notify") {
        ctx.ui.notify(
          `Current mode: ${mode} (usage: /pb-mode widget|notify)`,
          "info"
        );
        return;
      }
      setMode(m, ctx);
      ctx.ui.notify(`Parallel-batch mode → ${m}`, "info");
    },
  });

  // self-test: simulate two parallel tool batches
  pi.registerCommand("pb-demo", {
    description: "Simulate two parallel tool batches (self-test only)",
    handler: async (_args: string | undefined, ctx: any) => {
      lastCtx = ctx;
      bindRenderers(ctx);
      await playWave();
    },
  });
}