import type { ExtensionAPI, ExtensionUIContext } from "@oh-my-pi/pi-coding-agent";
import { mkdir, writeFile } from "node:fs/promises";
import { spawn } from "node:child_process";
import path from "node:path";

import { devflowHome } from "./paths.ts";
import { aggregateUsageReport, renderUsageHtml, renderUsageReport, type UsageRange } from "./usage-stats.ts";

export interface PlatformExtensionOptions {
  sessionRoot: string;
  ompSessionRoot: string;
  projectId: string;
  /** Read the persisted role assignment so optional roles can fail closed. */
  getConfiguredRole?: (role: string) => string | undefined;
}

type UsageFormat = "table" | "json" | "html";
type UsageInput = UsageRange & { source?: "project" | "omp"; format?: UsageFormat };

function commandUsageInput(args: string): UsageInput {
  const tokens = args.trim().split(/\s+/).filter(Boolean);
  const read = (flag: string): string | undefined => {
    const index = tokens.indexOf(flag);
    return index >= 0 ? tokens[index + 1] : undefined;
  };
  const source = read("--source") ?? (tokens[0] === "omp" || tokens[0] === "project" ? tokens[0] : undefined);
  return {
    source: source === "omp" ? "omp" : "project",
    session: read("--session"),
    from: read("--from"),
    to: read("--to"),
    models: read("--model")?.split(",").map(value => value.trim()).filter(Boolean),
    format: (read("--format") as UsageFormat | undefined) ?? "table",
  };
}

function localDate(offset = "+08:00"): string {
  const formatter = new Intl.DateTimeFormat("en-CA", { timeZone: offset === "+08:00" ? "Asia/Shanghai" : "UTC", year: "numeric", month: "2-digit", day: "2-digit" });
  const parts = Object.fromEntries(formatter.formatToParts(new Date()).filter(part => part.type !== "literal").map(part => [part.type, part.value]));
  return `${parts.year}-${parts.month}-${parts.day}`;
}

function daysAgo(days: number): string {
  const date = new Date();
  date.setDate(date.getDate() - days);
  return date.toISOString().slice(0, 10);
}

function openHtmlReport(file: string): void {
  const command = process.platform === "darwin" ? "open" : process.platform === "win32" ? "cmd" : "xdg-open";
  const args = process.platform === "win32" ? ["/c", "start", "", file] : [file];
  const child = spawn(command, args, { detached: true, stdio: "ignore" });
  child.unref();
}

export function createPlatformExtension(options: PlatformExtensionOptions) {
  return function devflowPlatform(pi: ExtensionAPI) {
    const { z } = pi.zod;
    pi.setLabel("Devflow Platform Runtime");
    let modelBeforeVisionTurn: Parameters<ExtensionAPI["setModel"]>[0] | undefined;
    let visionTurnActive = false;

    pi.on("input", async (event, ctx) => {
      if (event.source !== "interactive" || !event.images?.length) return;

      const configuredVision = options.getConfiguredRole?.("vision")?.trim();
      if (!configuredVision) {
        ctx.ui.notify("检测到图片，但尚未配置 VISION 模型。请先在 /model 中配置 VISION，再重新提交。", "warning");
        if (ctx.hasUI) await ctx.ui.input("VISION 未配置", "请先在 /model 中配置 VISION，完成后重新提交本次图片");
        return { handled: true };
      }

      const visionModel = ctx.models.resolve("@vision");
      if (!visionModel || !visionModel.input.includes("image")) {
        ctx.ui.notify(`VISION 配置无法解析为支持图片输入的模型：${configuredVision}`, "error");
        if (ctx.hasUI) await ctx.ui.input("VISION 配置无效", "请选择支持图片输入的模型后重新提交");
        return { handled: true };
      }

      modelBeforeVisionTurn = ctx.models.current();
      if (!(await pi.setModel(visionModel))) {
        modelBeforeVisionTurn = undefined;
        ctx.ui.notify(`无法切换到 VISION 模型：${visionModel.provider}/${visionModel.id}`, "error");
        return { handled: true };
      }
      visionTurnActive = true;
      ctx.ui.notify(`本回合已切换到 VISION：${visionModel.provider}/${visionModel.id}`, "info");
      return;
    });

    pi.on("turn_end", async (_event, ctx) => {
      if (!visionTurnActive || !modelBeforeVisionTurn) return;
      const previousModel = modelBeforeVisionTurn;
      modelBeforeVisionTurn = undefined;
      visionTurnActive = false;
      await pi.setModel(previousModel);
      ctx.ui.notify(`图片回合结束，已恢复模型：${previousModel.provider}/${previousModel.id}`, "info");
    });
    const usageParameters = z.object({
      source: z.enum(["project", "omp"]).default("project"),
      session: z.string().optional(),
      from: z.string().optional(),
      to: z.string().optional(),
      models: z.array(z.string()).optional(),
      format: z.enum(["table", "json", "html"]).default("table"),
    });
    const usageTool = {
      label: "Devflow Usage Stats",
      description:
        "Read token and cost totals. source=project (default) reads only the current Devflow project; source=omp explicitly reads ~/.omp history. Session buckets preserve their real relative paths and do not merge unrelated sessions into main.",
      parameters: usageParameters,
      approval: "read" as const,
      async execute(_id: string, params: unknown) {
        const input = params as UsageInput;
        const source = input.source ?? "project";
        const report = await aggregateUsageReport(
          source === "omp" ? options.ompSessionRoot : options.sessionRoot,
          source,
          { session: input.session, from: input.from, to: input.to, models: input.models },
        );
        return {
          content: [
            {
              type: "text" as const,
              text: input.format === "json" ? JSON.stringify({ projectId: options.projectId, ...report }, null, 2) : input.format === "html" ? renderUsageHtml(report) : renderUsageReport(report),
            },
          ],
          details: { projectId: options.projectId, ...report },
        };
      },
    };
    for (const name of ["devflow_usage_stats", "usage_stats"]) {
      pi.registerTool({
        name,
        ...usageTool,
      });
    }
    async function pickUsageInput(ctx: { ui: ExtensionUIContext; modelRegistry: { getAvailable(): Array<{ provider: string; id: string }> } }): Promise<UsageInput | undefined> {
      const sourceChoice = await ctx.ui.select("使用统计：数据范围", [
        { label: "当前 Devflow 项目", description: "只读取当前项目会话（推荐）" },
        { label: "OMP 历史", description: "读取 ~/.omp/agent/sessions" },
      ]);
      if (!sourceChoice) return undefined;
      const source = sourceChoice === "OMP 历史" ? "omp" : "project" as const;
      const root = source === "omp" ? options.ompSessionRoot : options.sessionRoot;
      const overview = await aggregateUsageReport(root, source);
      const sessionChoice = await ctx.ui.select("使用统计：会话", ["全部会话", ...Object.keys(overview.perSession)]);
      if (!sessionChoice) return undefined;
      // The selector is for filtering a report, so it must list every model
      // configured in OMP — not merely the model names found in old sessions.
      const registeredModels = [...new Set(ctx.modelRegistry.getAvailable().map(model => `${model.provider}/${model.id}`))].sort();
      const modelChoice = await ctx.ui.select("使用统计：模型", ["全部模型", ...registeredModels]);
      if (!modelChoice) return undefined;
      const period = await ctx.ui.select("使用统计：时间范围", ["全部时间", "今天", "最近 7 天", "最近 30 天", "自定义日期"]);
      if (!period) return undefined;
      let from: string | undefined;
      let to: string | undefined;
      if (period === "今天") from = to = localDate();
      if (period === "最近 7 天") { from = daysAgo(6); to = localDate(); }
      if (period === "最近 30 天") { from = daysAgo(29); to = localDate(); }
      if (period === "自定义日期") {
        from = await ctx.ui.input("使用统计：开始日期", "YYYY-MM-DD（留空表示不限）");
        if (from === undefined) return undefined;
        to = await ctx.ui.input("使用统计：结束日期", "YYYY-MM-DD（留空表示不限）");
        if (to === undefined) return undefined;
      }
      const output = await ctx.ui.select("使用统计：输出方式", [
        { label: "终端表格", description: "在当前终端显示" },
        { label: "HTML 报告", description: "生成可在浏览器打开的单文件报告" },
        { label: "JSON", description: "用于自动化处理" },
      ]);
      if (!output) return undefined;
      return { source, session: sessionChoice === "全部会话" ? "all" : sessionChoice, models: modelChoice === "全部模型" ? undefined : [modelChoice], from: from || undefined, to: to || undefined, format: output === "HTML 报告" ? "html" : output === "JSON" ? "json" : "table" };
    }

    async function renderCommand(input: UsageInput, ctx: { ui: { notify(message: string, type?: "info" | "warning" | "error"): void } }) {
      const source = input.source ?? "project";
      const report = await aggregateUsageReport(source === "omp" ? options.ompSessionRoot : options.sessionRoot, source, input);
      if (input.format === "html") {
        const reportDir = path.join(devflowHome(), "reports", options.projectId);
        const destination = path.join(reportDir, `usage-${new Date().toISOString().replaceAll(":", "-")}.html`);
        await mkdir(reportDir, { recursive: true });
        await writeFile(destination, renderUsageHtml(report), "utf8");
        openHtmlReport(destination);
        ctx.ui.notify(`已生成并在默认浏览器打开 HTML 使用报告：${destination}`, "info");
        return;
      }
      ctx.ui.notify(input.format === "json" ? JSON.stringify({ projectId: options.projectId, ...report }, null, 2).slice(0, 12000) : renderUsageReport(report).slice(0, 12000), "info");
    }

    pi.registerCommand("usage-stats", {
      description: "交互式查看 token/成本：范围、会话、模型、时间与 HTML 报告；也支持 --source/--session/--model/--format 参数。",
      handler: async (args, ctx) => {
        const input = args.trim() ? commandUsageInput(args) : await pickUsageInput(ctx);
        if (!input) return;
        await renderCommand(input, ctx);
      },
    });
    pi.registerCommand("devflow-usage", {
      description: "Alias for /usage-stats (project-scoped)",
      handler: async (_args, ctx) => {
        const report = await aggregateUsageReport(options.sessionRoot, "project");
        ctx.ui.notify(renderUsageReport(report), "info");
      },
    });
  };
}