/**
 * pi-router: DeepSeek V4 Flash 专用任务感知思维模式路由扩展。
 *
 * 移植自 dsh-routing-suite 的 mode-boost v0.1.0(实测驱动,单一事实源见 ./core.ts):
 * - **Flash-only 门控**:非 Flash 模型零干预——不注入 persona、不收窄工具、
 *   不注入引导,且 pi_router_* 工具从工具面隐藏(模型看不到)。
 * - **首轮锚定**:Flash 会话首轮只暴露分类匹配的核心工具,首次工具调用后
 *   恢复用户原配置(实测:可调用的 schema 面决定行为,不是工具名文本)。
 * - **近距离引导**:weak(内路由)带每条真实用户消息后注入一条固定引导
 *   (缓存中性;第 3 轮起 boost 重分类;简单任务快速收敛 / 复杂任务有向深度)。
 * - **寒暄让位**:问候/无任务消息完全不干预(实测:deep persona 套在聊天上
 *   → 338 块思维链)。
 * - **自优化工具**:pi_router_status / pi_router_mode / pi_router_subagent。
 *
 * 缓存原则(继承实测):persona 同一会话内静态不变(除非 override),引导走
 * 消息层注入;system prompt 前缀只变化一次。
 *
 * 零运行时依赖:`import type` 全部擦除,唯一运行时依赖 typebox(pi 扩展
 * 加载器内建解析)。
 */

import type {
  AgentMessage,
  ExtensionAPI,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import {
  bandFor,
  classifyTask,
  clamp01,
  coreFor,
  guideFor,
  isChatTask,
  isFlashModel,
  messageText,
  parseMode,
  personaFor,
  testinessFor,
  MODE_WEAK,
} from "./core.ts";

/** Marker used to idempotently locate our persona inside the system prompt. */
const PERSONA_MARKER = "# Reasoning-mode persona (pi-router)";

/** Custom entry type used to persist per-session router state. */
const STATE_ENTRY = "pi-router";

/** Tools owned by this extension (hidden from non-Flash sessions). */
const ROUTER_TOOLS = [
  "pi_router_status",
  "pi_router_mode",
  "pi_router_subagent",
];

/** Persisted override. null is persisted explicitly so `auto` survives restarts. */
interface PersistedState {
  override: number | "weak" | null;
}

interface SessionState {
  /** First REAL user message text of the session (mode anchor, resume-safe). */
  firstUserText: string | null;
  /** Explicit override from pi_router_mode (null = task classification). */
  override: number | "weak" | null;
  /** Whether the first-turn tool narrowing already happened (session-scoped, once). */
  firstTurnNarrowed: boolean;
  /** Active narrowing lease; ends at first tool call / agent settle / shutdown. */
  lease: { core: string[]; full: string[] } | null;
  /** Timestamps of user messages we already appended guidance for (dedupe). */
  guidedTimestamps: Set<number>;
}

function newState(): SessionState {
  return {
    firstUserText: null,
    override: null,
    firstTurnNarrowed: false,
    lease: null,
    guidedTimestamps: new Set(),
  };
}

/** Intersect core candidates with the user's active tool set, then add shell. */
function resolveCore(full: string[], mode: number | "weak"): string[] {
  const available = new Set(full);
  const core = new Set<string>();
  for (const name of coreFor(mode)) {
    if (available.has(name)) core.add(name);
  }
  if (available.has("bash")) core.add("bash");
  else if (available.has("pwsh")) core.add("pwsh");
  // 兜底:core 为空(用户禁用了全部核心工具)时至少保留读写面。
  if (core.size === 0) {
    for (const name of ["read", "write", "edit"]) {
      if (available.has(name)) core.add(name);
    }
  }
  return [...core];
}

/** Compare-and-set restore: only write back `full` if the surface is still
 *  exactly our `core` (another extension may have changed it meanwhile). */
function restoreLease(pi: ExtensionAPI, state: SessionState): void {
  if (state.lease === null) return;
  const current = pi.getActiveTools();
  const { core, full } = state.lease;
  const unchanged =
    current.length === core.length &&
    core.every((name) => current.includes(name));
  if (unchanged) {
    pi.setActiveTools(full);
  }
  state.lease = null;
}

export default function (pi: ExtensionAPI) {
  const states = new Map<string, SessionState>();
  let routerToolsHidden = false;

  function stateOf(sessionId: string | undefined): SessionState | undefined {
    return sessionId === undefined
      ? undefined
      : (states.get(sessionId) ?? undefined);
  }

  /** Hide pi_router_* from the tool surface (non-Flash zero-intervention). */
  function hideRouterTools(): void {
    if (routerToolsHidden) return;
    const active = pi.getActiveTools();
    const next = active.filter((name) => !ROUTER_TOOLS.includes(name));
    if (next.length !== active.length) {
      pi.setActiveTools(next);
      routerToolsHidden = true;
    }
  }

  /** Restore pi_router_* to the tool surface (Flash sessions). */
  function showRouterTools(): void {
    if (!routerToolsHidden) return;
    const active = pi.getActiveTools();
    pi.setActiveTools([
      ...active,
      ...ROUTER_TOOLS.filter((name) => !active.includes(name)),
    ]);
    routerToolsHidden = false;
  }

  /** Restore every outstanding lease (global setActiveTools semantics). */
  function restoreAllLeases(): void {
    for (const [, state] of states) {
      restoreLease(pi, state);
    }
  }

  /** Rebuild per-session state from the CURRENT BRANCH (resume/tree-safe). */
  function rebuildState(
    sessionId: string,
    ctx: {
      sessionManager: {
        getBranch(): {
          type: string;
          customType?: string;
          data?: unknown;
          message?: { role: string; content: unknown };
        }[];
      };
    },
  ): void {
    const state = newState();
    // 恢复 override(取当前分支最后一条 pi-router 条目;null 显式清除)
    for (const entry of ctx.sessionManager.getBranch()) {
      if (entry.type === "custom" && entry.customType === STATE_ENTRY) {
        const data = entry.data as PersistedState | undefined;
        if (data !== undefined && "override" in data)
          state.override = data.override;
      }
      // 找第一条真实 user 消息作为模式锚
      if (state.firstUserText === null && entry.type === "message") {
        const msg = entry.message;
        if (msg !== undefined && msg.role === "user") {
          const text = messageText(msg.content).trim();
          if (text) state.firstUserText = text;
        }
      }
    }
    states.set(sessionId, state);
  }

  // ── session_start / session_tree: 从当前分支重建状态 ──────────────────
  pi.on("session_start", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    if (sessionId === undefined) return;
    rebuildState(sessionId, ctx);
    // 启动即检查模型:非 Flash → 隐藏 router 工具
    if (ctx.model && !isFlashModel(ctx.model.id)) hideRouterTools();
  });

  pi.on("session_tree", (_event, ctx) => {
    // /tree 导航不触发 session_start,但分支切换必须重建状态
    const sessionId = ctx.sessionManager.getSessionId();
    if (sessionId === undefined) return;
    rebuildState(sessionId, ctx);
  });

  // ── session_shutdown: 恢复工具面 + 清理状态 ──────────────────────────
  pi.on("session_shutdown", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    if (sessionId === undefined) return;
    const state = states.get(sessionId);
    if (state !== undefined) restoreLease(pi, state);
    states.delete(sessionId);
  });

  // ── model_select: 门控随模型切换 ──────────────────────────────────────
  pi.on("model_select", (event) => {
    if (event.model && isFlashModel(event.model.id)) {
      showRouterTools();
    } else {
      restoreAllLeases();
      hideRouterTools();
    }
  });

  // ── agent_settled: 结束未恢复的首轮租约(首轮无工具调用也恢复) ────────
  pi.on("agent_settled", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    const state = stateOf(sessionId);
    if (state !== undefined) restoreLease(pi, state);
  });

  // ── before_agent_start: 核心路由点 ─────────────────────────────────────
  // 门控:非 Flash / 寒暄 → 零干预。Flash + 任务 → 注入静态 persona,
  // 首轮收窄工具面(mode 由 override ?? 会话首条消息分类决定)。
  pi.on("before_agent_start", (event, ctx) => {
    const model = ctx.model;
    if (!model || !isFlashModel(model.id)) return;
    showRouterTools(); // 防御:Flash 会话工具必须可见
    if (isChatTask(event.prompt)) return; // 寒暄让位,不干预

    const sessionId = ctx.sessionManager.getSessionId();
    const state = stateOf(sessionId);
    if (state === undefined) return;
    const firstText = state.firstUserText ?? event.prompt.trim(); // 会话首条真实消息
    state.firstUserText = firstText;
    const mode = state.override ?? classifyTask(firstText);
    const persona = personaFor(mode);

    // 首轮工具面收窄(仅会话一次;租约在首个工具调用/agent_settled 结束)
    if (!state.firstTurnNarrowed) {
      state.firstTurnNarrowed = true;
      const full = pi.getActiveTools();
      const core = resolveCore(full, mode);
      state.lease = { core, full };
      pi.setActiveTools(core);
    }

    // persona 幂等注入。基底优先取 setActiveTools 后的最新链式 prompt
    // (工具列表与收窄后的 schema 一致),取不到再退回事件快照。
    //
    // 真实语义(pi 源码核对,2026-08-17):before_agent_start 传入的
    // systemPrompt 是 _baseSystemPrompt——每次 agent turn 由
    // _rebuildSystemPrompt 重新组装,不含扩展注入内容。因此每个 turn
    // 都会注入一次 persona(这是必要行为,否则第二轮 persona 消失);
    // 不会叠加,因为 base 每次重建。includes 检查是防御性幂等:若未来
    // pi 语义变化导致 base 含 marker,此处跳过避免重复追加。
    const base = ctx.getSystemPrompt() ?? event.systemPrompt;
    if (!base.includes(PERSONA_MARKER)) {
      return { systemPrompt: `${base}\n\n${PERSONA_MARKER}\n${persona}` };
    }
  });

  // ── tool_call: 首次工具调用 → 结束收窄租约(首轮锚定放行) ─────────────
  pi.on("tool_call", (_event, ctx) => {
    const sessionId = ctx.sessionManager.getSessionId();
    const state = stateOf(sessionId);
    if (state !== undefined) restoreLease(pi, state);
  });

  // ── context: 近距离引导(仅 weak 带;每条真实 user 消息后一条) ─────────
  // 在每次 LLM 调用前检查:末条消息是未被引导过的 user 消息 → 追加引导。
  // 消息层注入,不落盘、不污染 session 文件,严格对应 DSH 的 inbox.append。
  // 去重键 = 消息 timestamp(唯一);round 计数排除已引导消息(即引导消息
  // 若进入历史也不会被计入)。
  pi.on("context", (event, ctx) => {
    const model = ctx.model;
    if (!model || !isFlashModel(model.id)) return;
    const sessionId = ctx.sessionManager.getSessionId();
    const state = stateOf(sessionId);
    if (state === undefined) return;
    const mode =
      state.override ??
      (state.firstUserText === null ? null : classifyTask(state.firstUserText));
    if (mode === null || bandFor(mode) !== "weak") return; // 强带无需引导

    const messages = event.messages;
    const last = messages[messages.length - 1];
    if (last === undefined || last.role !== "user") return;
    const text = messageText(last.content).trim();
    if (!text) return;
    if (state.guidedTimestamps.has(last.timestamp)) return; // 这条已引导过

    // round = 上下文中的真实用户消息序号(context messages 是完整上下文,
    // 注入的 guide 不落盘不会混入;防御性排除 'Router:' 开头的消息)
    const round = messages.filter(
      (m) =>
        m.role === "user" &&
        !messageText(m.content).trim().startsWith("Router:"),
    ).length;
    const guide = guideFor(round, text);
    state.guidedTimestamps.add(last.timestamp);
    return {
      messages: [
        ...messages,
        {
          role: "user",
          content: [{ type: "text", text: guide }],
          timestamp: Date.now(),
        } satisfies AgentMessage,
      ],
    };
  });

  // ── 自优化工具 ─────────────────────────────────────────────────────────

  pi.registerTool({
    name: "pi_router_status",
    label: "Pi Router Status",
    description:
      "Show this session's pi-router state: active model, mode, band, persona, first-turn core tools, and whether an override is active. Only meaningful for Flash-family models.",
    promptSnippet:
      "Show the current reasoning-mode routing state of this session",
    parameters: Type.Object({}),
    async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
      const model = ctx.model;
      const sessionId = ctx.sessionManager.getSessionId();
      const state = stateOf(sessionId);
      if (state === undefined) {
        return {
          content: [{ type: "text", text: "no session state" }],
          details: {},
        };
      }
      const mode =
        state.override ??
        (state.firstUserText === null
          ? null
          : classifyTask(state.firstUserText));
      const lines = [
        `model=${model ? `${model.provider}/${model.id}` : "unknown"}`,
        `flash-only-gate=${model ? (isFlashModel(model.id) ? "active" : "bypassed (non-Flash: zero intervention)") : "unknown"}`,
        `mode=${mode === null ? "unset" : typeof mode === "string" ? mode : mode.toFixed(2)} (band=${mode === null ? "-" : bandFor(mode)})`,
        `persona=${mode === null ? "-" : personaFor(mode).replace(/\n/g, " / ")}`,
        `core=[${state.lease ? state.lease.core.join(", ") : "(full surface restored)"}]`,
        `testiness=${mode === null ? "-" : testinessFor(mode)}`,
        `override=${state.override === null ? "no" : typeof state.override === "string" ? state.override : state.override.toFixed(2)}`,
      ];
      return {
        content: [{ type: "text", text: lines.join("\n") }],
        details: {},
      };
    },
  });

  pi.registerTool({
    name: "pi_router_mode",
    label: "Pi Router Mode",
    description:
      "Set this session's reasoning mode (Flash sessions only): spec (plan-first) / weak (internal routing, model decides per task) / mixed (transition, trap) / react (doer). Accepts band names, 0-100, or 0.0-1.0; use auto to return to task classification. The next request applies it. Non-Flash sessions are unaffected.",
    promptSnippet:
      "Override the reasoning mode of this session (spec/weak/mixed/react/auto)",
    parameters: Type.Object({
      mode: Type.String({
        description:
          "band name (spec / weak / mixed / react), a 0-100 number, a 0.0-1.0 number, or auto to clear the override",
      }),
    }),
    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      if (ctx.model && !isFlashModel(ctx.model.id)) {
        return {
          content: [
            {
              type: "text",
              text: "pi-router is Flash-only: non-Flash sessions are not routed",
            },
          ],
          details: {},
        };
      }
      const parsed = parseMode(params.mode as string);
      if (parsed === null) {
        return {
          content: [
            {
              type: "text",
              text: `invalid mode "${params.mode}": use spec/weak/mixed/react, 0-100, 0.0-1.0, or auto`,
            },
          ],
          details: {},
        };
      }
      const sessionId = ctx.sessionManager.getSessionId();
      const state = stateOf(sessionId);
      if (state === undefined) {
        return {
          content: [{ type: "text", text: "no session state" }],
          details: {},
        };
      }
      // auto → null;null 显式持久化,resume 后正确清除旧 override
      state.override =
        parsed === "auto"
          ? null
          : parsed === "weak"
            ? MODE_WEAK
            : clamp01(parsed);
      pi.appendEntry(STATE_ENTRY, {
        override: state.override,
      } satisfies PersistedState);
      const current =
        state.override ??
        (state.firstUserText === null
          ? null
          : classifyTask(state.firstUserText));
      const text = `mode=${current === null ? "unset" : typeof current === "string" ? current : current.toFixed(2)} (band=${current === null ? "-" : bandFor(current)}) — next request applies`;
      return { content: [{ type: "text", text }], details: {} };
    },
  });

  pi.registerTool({
    name: "pi_router_subagent",
    label: "Pi Router Subagent",
    description:
      "Run one task in a DIFFERENT reasoning mode than this session, in a fresh isolated context (own system prompt). The current session trajectory is untouched. Mode: spec (plan-first) / weak (internal routing) / react (doer) / balanced. Returns the answer text.",
    promptSnippet:
      "Run a task in a different reasoning mode with an isolated system prompt",
    parameters: Type.Object({
      mode: Type.String({
        description: "spec / weak / react / balanced (or 0-100)",
      }),
      task: Type.String({
        description: "the task to hand to the mode-isolated subagent",
      }),
      maxTokens: Type.Optional(
        Type.Number({ description: "output cap (default 1024)" }),
      ),
    }),
    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      if (ctx.model && !isFlashModel(ctx.model.id)) {
        return {
          content: [
            {
              type: "text",
              text: "pi-router is Flash-only: non-Flash sessions are not routed",
            },
          ],
          details: {},
        };
      }
      const parsed = parseMode(params.mode as string);
      if (parsed === null || parsed === "auto") {
        return {
          content: [{ type: "text", text: `invalid mode "${params.mode}"` }],
          details: {},
        };
      }
      const model = ctx.model;
      if (!model) {
        return {
          content: [{ type: "text", text: "no active model" }],
          details: {},
        };
      }
      let maxTokens = Number(params.maxTokens || 1024);
      if (!Number.isFinite(maxTokens) || maxTokens <= 0) maxTokens = 1024;
      maxTokens = Math.min(maxTokens, model.maxTokens ?? maxTokens);
      const persona = personaFor(parsed);
      const result = await ctx.modelRegistry.complete(
        model,
        {
          systemPrompt: persona,
          messages: [
            {
              role: "user",
              content: [{ type: "text", text: String(params.task) }],
              timestamp: Date.now(),
            },
          ],
        },
        { maxTokens, signal },
      );
      if (result.stopReason === "error" || result.stopReason === "aborted") {
        throw new Error(
          `subagent ${result.stopReason}: ${messageText(result.content).trim() || "no output"}`,
        );
      }
      const text = messageText(result.content).trim();
      const head = text.slice(0, 3000);
      return {
        content: [
          {
            type: "text",
            text: `[mode-subagent ${bandFor(parsed)}]\n${head}${text.length > 3000 ? "\n…(truncated)" : ""}`,
          },
        ],
        details: { mode: bandFor(parsed), length: text.length },
        usage: result.usage,
      };
    },
  });
}