import { sql } from "kysely";
import { z } from "zod";
import { defineTask, TaskRegistry } from "../../src";
import type { TestDb } from "../helpers";

const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

/** 测试任务;副作用写 side_effect 表,用来判断步骤是否重跑 */
export function createTestRegistry(db: TestDb) {
  const mark = (runId: string, step: string, note: string | null = null) =>
    sql`INSERT INTO side_effect (run_id, step, note, pid) VALUES (${runId}, ${step}, ${note}, ${process.pid})`.execute(
      db,
    );

  const threeSteps = defineTask({
    name: "threeSteps",
    version: 1,
    inputSchema: z.object({ label: z.string() }),
    handler: async (ctx) => {
      const a = await ctx.step.run("step1", async () => {
        await mark(ctx.runId, "step1");
        return 1;
      });
      const b = await ctx.step.run("step2", async () => {
        await mark(ctx.runId, "step2");
        return 2;
      });
      const c = await ctx.step.run("step3", async () => {
        await mark(ctx.runId, "step3", "start");
        // 崩溃测试:第一个 worker 卡在这里被 kill -9
        if (process.env.HANG_STEP3 === "1") await sleep(10 * 60_000);
        await mark(ctx.runId, "step3", "end");
        return 3;
      });
      return { label: ctx.input.label, sum: a + b + c };
    },
  });

  const cancellable = defineTask({
    name: "cancellable",
    version: 1,
    inputSchema: z.object({}),
    handler: async (ctx) => {
      await ctx.step.run("first", () => mark(ctx.runId, "first").then(() => 1));
      await ctx.step.run("long", async () => {
        await mark(ctx.runId, "long", "start");
        for (let i = 0; i < 300; i++) {
          if (ctx.abortSignal.aborted) {
            await mark(ctx.runId, "long", "aborted");
            throw ctx.abortSignal.reason;
          }
          await sleep(100);
        }
        return 1;
      });
      await ctx.step.run("after", () => mark(ctx.runId, "after").then(() => 1));
      return "done";
    },
  });

  const waitFacts = defineTask({
    name: "waitFacts",
    version: 1,
    inputSchema: z.object({}),
    handler: async (ctx) => {
      await ctx.step.run("before", async () => {
        await mark(ctx.runId, "before");
        await sleep(1500);
        return 1;
      });
      const a = await ctx.step.waitForFact("a", { timeout: "60s" });
      await ctx.step.run("got-a", () => mark(ctx.runId, "got-a").then(() => 1));
      const b = await ctx.step.waitForFact("b", { timeout: "60s" });
      return { a, b };
    },
  });

  const echo = defineTask({
    name: "echo",
    version: 1,
    inputSchema: z.object({ n: z.number() }),
    handler: async (ctx) => ctx.step.run("echo", async () => ({ n: ctx.input.n, trace: ctx.traceId })),
  });

  return new TaskRegistry().register(threeSteps, cancellable, waitFacts, echo);
}