import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { z } from "zod";
import { defineTask, type OpenWorkflowEngine, TaskRegistry } from "../src/index.ts";
import {
  connectDb,
  dropSchemas,
  migrateBizSchema,
  pgAvailable,
  resetSchemas,
  setupEngine,
  type TestDb,
  testSchemas,
  waitFor,
} from "./helpers.ts";

// worker 分组(io / agent):带 versions 的 worker 只认领自己的任务,不会误认领后以"未注册"失败再退避。

const schemas = testSchemas("groups");
const hasPg = await pgAvailable();

/** 两个"进程"各一套注册表(任务集相同,handler 记下自己是哪一组) */
function registryFor(group: string, log: string[]) {
  const mk = (name: string) =>
    defineTask({
      name,
      version: 1,
      inputSchema: z.object({ i: z.number() }),
      handler: async (ctx) => {
        log.push(`${group}:${name}`);
        return { group, i: ctx.input.i };
      },
    });
  return new TaskRegistry().register(mk("ioTask"), mk("runAgent"), mk("resolveDecision"));
}

describe.skipIf(!hasPg)("worker 分组认领", () => {
  let db: TestDb;
  let io: OpenWorkflowEngine;
  let agent: OpenWorkflowEngine;
  const log: string[] = [];

  beforeAll(async () => {
    await resetSchemas(schemas);
    db = connectDb(schemas.biz);
    await migrateBizSchema(db, schemas.biz);
    io = await setupEngine(db, schemas, registryFor("io", log));
    agent = await setupEngine(db, schemas, registryFor("agent", log));
    await io.startWorker({ versions: ["ioTask"], concurrency: 4 });
    await agent.startWorker({ versions: ["runAgent", "resolveDecision"], concurrency: 1 });
  }, 30_000);

  afterAll(async () => {
    await io?.stop();
    await agent?.stop();
    await db?.destroy();
    await dropSchemas(schemas);
  }, 60_000);

  it("io 组只跑 ioTask,agent 组只跑 runAgent/resolveDecision;每个 run 一次认领即完成", async () => {
    const runs: { runId: string; task: string }[] = [];
    for (let i = 0; i < 6; i++) {
      for (const task of ["ioTask", "runAgent", "resolveDecision"]) {
        const r = await io.enqueue(task, { i }, { idempotencyKey: `${task}-${i}`, traceId: `t-${i}` });
        runs.push({ runId: r.runId, task });
      }
    }
    for (const r of runs) {
      const v = await waitFor(async () => {
        const x = await io.get(r.runId);
        return x?.status === "completed" ? x : null;
      }, 20_000);
      expect(v.output).toMatchObject({ group: r.task === "ioTask" ? "io" : "agent" });
      // 误认领会以 "not registered" 失败并退避 → attempts > 1
      expect(v.attempts).toBe(1);
    }
    expect(log.filter((l) => l.startsWith("io:")).every((l) => l === "io:ioTask")).toBe(true);
    expect(log.filter((l) => l.startsWith("agent:")).every((l) => l !== "agent:ioTask")).toBe(true);
    expect(log).toHaveLength(18);
  }, 60_000);
});