import { randomUUID } from "node:crypto";
import type { FakeOAuthControl } from "@umt/agent";
import { defineTask } from "@umt/tasks";
import { z } from "zod";
import { registerScheduledJob } from "../schedules/registry.ts";
import type { ServerContext, ServerModule } from "../types.ts";
import { type OAuthOptions, OAuthSessions } from "./oauth.ts";
import { refreshPiTokens } from "./refresh.ts";
import { registerPiRoutes } from "./routes.ts";

export { type RunUpstream, resolveRunUpstream } from "./service.ts";

/**
 * W7-A Pi 模型上游:pi_upstream CRUD、能力表、测试、模型表、服务端订阅登录(挂起式 interaction)、
 * 持久任务 refreshPiTokens(周期登记为定时任务 pi.token_refresh;OAuth 过期前刷新,行锁内写回)。runAgent 经 ./service.ts 的 resolveRunUpstream 取上游。
 * 用法见 docs/runbook/agent.md「上游与登录」。
 */
export interface PiModuleOptions {
  /** provider = fake 的模拟 OAuth(测试注入) */
  fakeOAuth?: FakeOAuthControl;
  oauth?: Omit<OAuthOptions, "client">;
  /** 定时任务 pi.token_refresh 的默认周期(默认 10 分钟;<=0 不登记) */
  refreshIntervalMs?: number;
  /** 剩余有效期少于它就刷新(默认 30 分钟) */
  refreshWindowMs?: number;
}

export function createPiModule(
  opts: PiModuleOptions = {},
): ServerModule & { sessions(ctx: ServerContext): OAuthSessions } {
  const client = opts.fakeOAuth ? { fakeOAuth: opts.fakeOAuth } : undefined;
  const sessions = new WeakMap<ServerContext, OAuthSessions>();
  const sessionsFor = (ctx: ServerContext) => {
    let s = sessions.get(ctx);
    if (!s) {
      s = new OAuthSessions(ctx, { ...opts.oauth, ...(client ? { client } : {}) });
      sessions.set(ctx, s);
    }
    return s;
  };
  const live: OAuthSessions[] = [];
  const interval = opts.refreshIntervalMs ?? 10 * 60_000;
  const windowMs = opts.refreshWindowMs ?? 30 * 60_000;

  return {
    name: "pi",
    sessions: sessionsFor,
    registerRoutes(app, ctx) {
      registerPiRoutes(app, ctx, {
        oauth: () => {
          const s = sessionsFor(ctx);
          if (!live.includes(s)) live.push(s);
          return s;
        },
        ...(client ? { client } : {}),
      });
    },
    registerTasks(registry, ctx) {
      registry.register(
        defineTask({
          name: "refreshPiTokens",
          version: 1,
          inputSchema: z.object({}).passthrough(),
          retryPolicy: { maximumAttempts: 2, initialInterval: "10s", backoffCoefficient: 2, maximumInterval: "60s" },
          handler: async (t) =>
            t.step.run("refresh", () => refreshPiTokens(ctx, { windowMs, ...(client ? { client } : {}) })),
        }),
      );
    },
    /** 刷新周期登记为定时任务 pi.token_refresh(维护组,由 schedules 调度器驱动),执行体只入队持久任务 */
    start(ctx) {
      if (interval <= 0) return;
      registerScheduledJob(ctx, {
        key: "pi.token_refresh",
        defaults: {
          name: "上游登录刷新",
          description: "订阅登录即将过期的 Agent 上游,提前刷新 token",
          group: "maintain",
          schedule: { type: "interval", everyMs: interval },
        },
        run: async (c) => {
          const r = await ctx.engine.enqueue(
            "refreshPiTokens",
            {},
            { idempotencyKey: `refreshPiTokens:${c.slot}`, traceId: randomUUID() },
          );
          return { summary: r.deduplicated ? "本周期已入队" : "已入队刷新" };
        },
      });
    },
    async stop() {
      await Promise.all(live.splice(0).map((s) => s.shutdown()));
    },
  };
}

export const piModule = createPiModule();