import { randomUUID } from "node:crypto";
import { type Kysely, sql } from "kysely";
import { defineWorkflow, OpenWorkflow, Worker, type Workflow } from "openworkflow";
import { BackendPostgres } from "openworkflow/postgres";
import { TaskCanceledError } from "./errors.ts";
import { createLogger, type LogSink } from "./logger.ts";
import { engineVersion, parseEngineVersion, type TaskDefinition, type TaskRegistry } from "./registry.ts";
import type {
  EnqueueOptions,
  JsonValue,
  ListFilter,
  QueueCounts,
  StartWorkerOptions,
  TaskContext,
  TaskEngine,
  TaskRunStatus,
  TaskRunView,
  TaskStepApi,
  WaitForFactOptions,
  WorkerHandle,
} from "./types.ts";

// biome-ignore lint/suspicious/noExplicitAny: 本包只用原生 SQL 访问自有表
type AnyDb = Kysely<any>;
type OwWorkflowRun = NonNullable<Awaited<ReturnType<BackendPostgres["getWorkflowRun"]>>>;
type OwStep = Parameters<Parameters<typeof defineWorkflow>[1]>[0]["step"];

export const ENGINE_SCHEMA = "openworkflow";
const DEFAULT_STEP_RETRY = { maximumAttempts: 3, initialInterval: "1s", backoffCoefficient: 2, maximumInterval: "60s" };

/** 写入引擎的 run 输入:业务输入 + 追踪元数据 */
interface Envelope {
  input: JsonValue;
  traceId: string;
  priority: number;
}

export interface OpenWorkflowEngineOptions {
  /** 引擎自己的 postgres.js 连接池使用(与 Kysely 同库) */
  databaseUrl: string;
  /** 业务库(Kysely/pg);tasks_* 表所在 */
  db: AnyDb;
  registry: TaskRegistry;
  /** 固定 "openworkflow";只允许测试覆盖 */
  schema?: string;
  /** 默认 true(幂等)。部署时由 migrateEngineSchema 预先迁移,多进程启动传 false */
  runMigrations?: boolean;
  logSink?: LogSink;
  /** handler 执行期间轮询取消标志的间隔,默认 1000ms */
  cancelPollMs?: number;
}

export const signalName = (runId: string, name: string) => `umt:${runId}:${name}`;

/**
 * 只跑引擎自己的迁移(openworkflow schema)后断开。部署时由 `--role migrate` 调用,
 * 之后 api / worker 进程以 runMigrations:false 连接,避免多进程同时迁移。
 */
export async function migrateEngineSchema(databaseUrl: string, schema = ENGINE_SCHEMA): Promise<void> {
  const backend = await BackendPostgres.connect(databaseUrl, { schema, runMigrations: true });
  await backend.stop();
}

interface ClaimWrapOptions {
  /** 只认领这些 "名称@引擎版本";null = 不按任务过滤 */
  keys: string[] | null;
  /** 本 worker 正在执行的 run 数 */
  active: () => number;
  maxActive?: () => number;
  gate?: () => boolean | Promise<boolean>;
}

/**
 * 包一层 backend 的认领:
 * - keys:与 BackendPostgres.claimWorkflowRun 同一条 SQL(截止过期处理 + SKIP LOCKED 认领),candidate 多一个
 *   workflow_name/version 过滤(OpenWorkflow 0.10.1 认领不看任务名,认领到未注册的 run 会以"未注册"失败并退避);
 * - maxActive:执行中 + 认领中 ≥ 上限就不认领(同一轮 tick 的多个槽并发认领,用 reserved 计数挡住超额);
 * - gate:返回 false 不认领。
 * 其余方法经原型链委托给原 backend。
 */
function claimWrappedBackend(backend: BackendPostgres, o: ClaimWrapOptions): BackendPostgres {
  // biome-ignore lint/suspicious/noExplicitAny: 访问 BackendPostgres 的私有字段(pg / namespaceId / workflowRunsTable)
  const b = backend as any;
  const filtered = Object.create(backend) as BackendPostgres;
  let reserved = 0;
  const claim = async (params: { workerId: string; leaseDurationMs: number }) => {
    if (!o.keys) return backend.claimWorkflowRun(params);
    const table = b.workflowRunsTable();
    const [claimed] = await b.pg`
      WITH expired AS (
        UPDATE ${table}
        SET "status" = 'failed', "error" = ${b.pg.json({ message: "Workflow run deadline exceeded" })},
            "worker_id" = NULL, "available_at" = NULL, "finished_at" = NOW(), "updated_at" = NOW()
        WHERE "namespace_id" = ${b.namespaceId}
          AND "status" IN ('pending', 'running', 'sleeping')
          AND "deadline_at" IS NOT NULL
          AND "deadline_at" <= NOW()
        RETURNING "id"
      ),
      candidate AS (
        SELECT "id"
        FROM ${table}
        WHERE "namespace_id" = ${b.namespaceId}
          AND "status" IN ('pending', 'running', 'sleeping')
          AND "available_at" <= NOW()
          AND ("deadline_at" IS NULL OR "deadline_at" > NOW())
          AND ("workflow_name" || '@' || COALESCE("version", '')) = ANY(${o.keys}::text[])
        ORDER BY CASE WHEN "status" = 'pending' THEN 0 ELSE 1 END, "available_at", "created_at"
        LIMIT 1
        FOR UPDATE SKIP LOCKED
      )
      UPDATE ${table} AS wr
      SET "status" = 'running', "attempts" = "attempts" + 1, "worker_id" = ${params.workerId},
          "available_at" = NOW() + ${params.leaseDurationMs} * INTERVAL '1 millisecond',
          "started_at" = COALESCE("started_at", NOW()), "updated_at" = NOW()
      FROM candidate
      WHERE wr."id" = candidate."id" AND wr."namespace_id" = ${b.namespaceId}
      RETURNING wr.*;`;
    return claimed ?? null;
  };
  Object.defineProperty(filtered, "claimWorkflowRun", {
    value: async (params: { workerId: string; leaseDurationMs: number }) => {
      // 同步段先占位:同一轮 tick 的各槽依次进到这里,reserved 保证不超过动态上限
      if (o.maxActive && o.active() + reserved >= Math.max(0, o.maxActive())) return null;
      reserved += 1;
      try {
        if (o.gate && !(await o.gate())) return null;
        return await claim(params);
      } finally {
        // 认领到的 run 在本函数返回后同步加入 worker 的执行集合,下一轮 tick 前已计入 active
        reserved -= 1;
      }
    },
  });
  return filtered;
}

export class OpenWorkflowEngine implements TaskEngine {
  private readonly ow: OpenWorkflow;
  private readonly workers: Worker[] = [];
  private readonly log;

  private constructor(
    private readonly backend: BackendPostgres,
    private readonly opts: OpenWorkflowEngineOptions,
  ) {
    this.ow = new OpenWorkflow({ backend });
    this.log = createLogger(opts.logSink, { component: "tasks.engine" });
  }

  static async connect(opts: OpenWorkflowEngineOptions): Promise<OpenWorkflowEngine> {
    const backend = await BackendPostgres.connect(opts.databaseUrl, {
      schema: opts.schema ?? ENGINE_SCHEMA,
      runMigrations: opts.runMigrations ?? true,
    });
    return new OpenWorkflowEngine(backend, opts);
  }

  async enqueue(taskName: string, input: unknown, opts: EnqueueOptions) {
    const def = this.opts.registry.latest(taskName);
    const parsed = def.inputSchema.parse(input) as JsonValue;
    const envelope: Envelope = { input: parsed, traceId: opts.traceId, priority: opts.priority ?? 0 };
    const handle = await this.ow.runWorkflow(
      { name: def.name, version: engineVersion(def.version) },
      envelope as unknown as JsonValue,
      { idempotencyKey: opts.idempotencyKey, ...(opts.availableAt ? { availableAt: opts.availableAt } : {}) },
    );
    const run = handle.workflowRun;
    const inserted = await sql`
      INSERT INTO tasks_run_meta (run_id, task_name, task_version, idempotency_key, trace_id, priority)
      VALUES (${run.id}, ${def.name}, ${def.version}, ${opts.idempotencyKey}, ${opts.traceId}, ${envelope.priority})
      ON CONFLICT (run_id) DO NOTHING`.execute(this.opts.db);
    // 同一幂等键返回已有 run 时 meta 行已存在
    return { runId: run.id, deduplicated: Number(inserted.numAffectedRows ?? 0n) === 0 };
  }

  async signal(runId: string, name: string, payload: JsonValue = null) {
    await sql`INSERT INTO tasks_fact (run_id, name, payload) VALUES (${runId}, ${name}, ${JSON.stringify(payload)}::jsonb)`.execute(
      this.opts.db,
    );
    const r = await this.ow.sendSignal({ signal: signalName(runId, name), data: payload });
    return { delivered: r.workflowRunIds.length };
  }

  async cancel(runId: string, reason?: string) {
    // 先写业务侧标志:运行中的 handler 通过它触发 abortSignal,步骤间检查也读它
    await sql`
      INSERT INTO tasks_run_meta (run_id, task_name, task_version, idempotency_key, cancel_requested_at, cancel_reason)
      VALUES (${runId}, '', 0, '', now(), ${reason ?? null})
      ON CONFLICT (run_id) DO UPDATE SET cancel_requested_at = COALESCE(tasks_run_meta.cancel_requested_at, now()),
        cancel_reason = COALESCE(tasks_run_meta.cancel_reason, EXCLUDED.cancel_reason)`.execute(this.opts.db);
    try {
      await this.ow.cancelWorkflowRun(runId);
      return true;
    } catch (err) {
      const run = await this.backend.getWorkflowRun({ workflowRunId: runId });
      if (run?.status === "canceled") return true;
      if (run) return false;
      throw err;
    }
  }

  async get(runId: string): Promise<TaskRunView | null> {
    const run = await this.backend.getWorkflowRun({ workflowRunId: runId });
    if (!run) return null;
    const metas = await this.loadMeta([run.id]);
    return toView(run, metas.get(run.id));
  }

  async list(filter: ListFilter = {}): Promise<TaskRunView[]> {
    const res = await this.backend.listWorkflowRuns({
      limit: filter.limit ?? 50,
      ...(filter.taskName ? { workflowName: filter.taskName } : {}),
      ...(filter.status ? { status: filter.status } : {}),
    });
    const metas = await this.loadMeta(res.data.map((r) => r.id));
    return res.data.map((r) => toView(r, metas.get(r.id)));
  }

  async startWorker(opts: StartWorkerOptions = {}): Promise<WorkerHandle> {
    const defs = this.opts.registry.select(opts.versions);
    const workflows = defs.map((d) => this.toWorkflow(d)) as Workflow<unknown, unknown, unknown>[];
    // 指定了 versions 的 worker 只认领这些任务:OpenWorkflow 0.10.1 认领不看任务名,认领到未注册的 run 会
    // 以"未注册"失败并指数退避(5s→5m),分组 worker(各道 / agent)会互相拖慢对方的任务。
    const keys = opts.versions?.length ? [...new Set(defs.map((d) => `${d.name}@${engineVersion(d.version)}`))] : null;
    let worker: Worker | null = null;
    // biome-ignore lint/suspicious/noExplicitAny: OpenWorkflow Worker 的执行集合是私有字段,只读取其大小
    const active = () => ((worker as any)?.activeExecutions?.size as number | undefined) ?? 0;
    const backend =
      keys || opts.maxActive || opts.claimGate
        ? claimWrappedBackend(this.backend, {
            keys,
            active,
            ...(opts.maxActive ? { maxActive: opts.maxActive } : {}),
            ...(opts.claimGate ? { gate: opts.claimGate } : {}),
          })
        : this.backend;
    const concurrency = opts.concurrency ?? 4;
    worker = new Worker({ backend, workflows, concurrency });
    await worker.start();
    this.workers.push(worker);
    this.log.info("worker started", { tasks: defs.map((d) => `${d.name}@${d.version}`) });
    return { concurrency, active };
  }

  /**
   * 按任务名分组统计排队 / 运行数(全库):
   * queued = 已到 available_at 可被认领(pending、等待后到点、租约过期待接手);
   * running = worker 持有未过期租约。等信号 / sleep 未到点的不计。
   */
  async queueCounts(groups: readonly (readonly string[])[]): Promise<QueueCounts[]> {
    // biome-ignore lint/suspicious/noExplicitAny: 同 claimWrappedBackend
    const b = this.backend as any;
    const table = b.workflowRunsTable();
    const all = [...new Set(groups.flat())];
    if (all.length === 0) return groups.map(() => ({ queued: 0, running: 0 }));
    const rows: { name: string; queued: string | number; running: string | number }[] = await b.pg`
      SELECT "workflow_name" AS name,
        COUNT(*) FILTER (WHERE "available_at" <= NOW()) AS queued,
        COUNT(*) FILTER (WHERE "status" = 'running' AND "worker_id" IS NOT NULL AND "available_at" > NOW()) AS running
      FROM ${table}
      WHERE "namespace_id" = ${b.namespaceId}
        AND "status" IN ('pending', 'running', 'sleeping')
        AND ("deadline_at" IS NULL OR "deadline_at" > NOW())
        AND "workflow_name" = ANY(${all}::text[])
      GROUP BY "workflow_name"`;
    const by = new Map(rows.map((r) => [r.name, { queued: Number(r.queued), running: Number(r.running) }]));
    return groups.map((names) =>
      names.reduce(
        (acc, n) => {
          const c = by.get(n);
          return c ? { queued: acc.queued + c.queued, running: acc.running + c.running } : acc;
        },
        { queued: 0, running: 0 },
      ),
    );
  }

  /** 等进行中的 handler 执行完再断开连接池 */
  async stop() {
    await Promise.all(this.workers.splice(0).map((w) => w.stop()));
    await this.backend.stop();
  }

  private async loadMeta(runIds: string[]) {
    const out = new Map<string, MetaRow>();
    if (runIds.length === 0) return out;
    const r = await sql<MetaRow>`SELECT * FROM tasks_run_meta WHERE run_id = ANY(${runIds}::text[])`.execute(
      this.opts.db,
    );
    for (const row of r.rows) out.set(row.run_id, row);
    return out;
  }

  private async cancelState(runId: string): Promise<{ at: Date | null; reason: string | null }> {
    const r = await sql<{ cancel_requested_at: Date | null; cancel_reason: string | null }>`
      SELECT cancel_requested_at, cancel_reason FROM tasks_run_meta WHERE run_id = ${runId}`.execute(this.opts.db);
    return { at: r.rows[0]?.cancel_requested_at ?? null, reason: r.rows[0]?.cancel_reason ?? null };
  }

  private toWorkflow(def: TaskDefinition) {
    const { db } = this.opts;
    const pollMs = this.opts.cancelPollMs ?? 1000;
    const stepRetry = { ...DEFAULT_STEP_RETRY, ...def.retryPolicy };
    return defineWorkflow(
      { name: def.name, version: engineVersion(def.version), retryPolicy: { maximumAttempts: 1 } },
      async ({ input, step, run }) => {
        const env = input as unknown as Envelope;
        const parsed = def.inputSchema.parse(env.input);
        const attemptToken = randomUUID();
        const logger = createLogger(this.opts.logSink, {
          task: def.name,
          version: def.version,
          runId: run.id,
          traceId: env.traceId,
        });
        const abort = new AbortController();
        let cancelReason: string | null = null;
        const poll = async () => {
          const s = await this.cancelState(run.id);
          if (s.at && !abort.signal.aborted) {
            cancelReason = s.reason;
            abort.abort(new TaskCanceledError(run.id, s.reason));
          }
          return abort.signal.aborted;
        };
        const timer = setInterval(() => {
          poll().catch((err) => logger.warn("cancel poll failed", { err }));
        }, pollMs);
        const throwIfCanceled = async () => {
          if (await poll()) throw new TaskCanceledError(run.id, cancelReason);
        };

        const stepApi = makeStepApi(step as OwStep, run.id, db, stepRetry, throwIfCanceled);
        const ctx: TaskContext<unknown> = {
          input: parsed,
          runId: run.id,
          taskName: def.name,
          version: def.version,
          traceId: env.traceId,
          priority: env.priority,
          attemptToken,
          abortSignal: abort.signal,
          logger,
          step: stepApi,
          throwIfCanceled,
        };
        try {
          await throwIfCanceled();
          const out = await def.handler(ctx);
          return (out ?? null) as JsonValue;
        } finally {
          clearInterval(timer);
        }
      },
    );
  }
}

function makeStepApi(
  step: OwStep,
  runId: string,
  db: AnyDb,
  stepRetry: typeof DEFAULT_STEP_RETRY,
  throwIfCanceled: () => Promise<void>,
): TaskStepApi {
  const api: TaskStepApi = {
    async run(name, fn, opts) {
      await throwIfCanceled();
      // biome-ignore lint/suspicious/noExplicitAny: 引擎的 DurationString 是模板字面量类型
      const retryPolicy = { ...stepRetry, ...opts?.retryPolicy } as any;
      return (await step.run({ name, retryPolicy }, async () => (await fn()) as never)) as never;
    },
    async sleep(name, duration) {
      await throwIfCanceled();
      // biome-ignore lint/suspicious/noExplicitAny: 同上
      await step.sleep(name, duration as any);
    },
    async waitForFact<T>(name: string, opts: WaitForFactOptions<T> = {}) {
      const check =
        opts.check ??
        (async () => {
          const r = await sql<{ payload: T }>`
            SELECT payload FROM tasks_fact WHERE run_id = ${runId} AND name = ${name} ORDER BY id DESC LIMIT 1`.execute(
            db,
          );
          return r.rows[0] ? r.rows[0].payload : undefined;
        });
      const maxRounds = opts.maxRounds ?? 100;
      for (let i = 0; ; i++) {
        const got = await api.run(`fact:${name}:check:${i}`, async () => {
          const v = await check();
          return v === undefined || v === null ? { found: false } : { found: true, value: v as unknown };
        });
        if (got.found) return (got as { value: T }).value;
        if (i >= maxRounds) return null;
        await throwIfCanceled();
        await step.waitForSignal({
          name: `fact:${name}:wait:${i}`,
          signal: signalName(runId, name),
          timeout: opts.timeout ?? "5m",
        });
      }
    },
  };
  return api;
}

interface MetaRow {
  run_id: string;
  task_name: string;
  task_version: number;
  idempotency_key: string;
  trace_id: string | null;
  priority: number;
  cancel_requested_at: Date | null;
  cancel_reason: string | null;
}

function toStatus(run: OwWorkflowRun): TaskRunStatus {
  switch (run.status) {
    case "succeeded":
    case "completed":
      return "completed";
    case "sleeping":
      return "waiting";
    case "running":
      // 等信号 / sleep 期间 status 仍是 running,但没有 worker 持有
      return run.workerId ? "running" : "waiting";
    default:
      return run.status;
  }
}

function toView(run: OwWorkflowRun, meta: MetaRow | undefined): TaskRunView {
  const env = (run.input ?? {}) as Partial<Envelope>;
  return {
    runId: run.id,
    taskName: run.workflowName,
    version: parseEngineVersion(run.version),
    status: toStatus(run),
    attempts: run.attempts,
    idempotencyKey: run.idempotencyKey,
    traceId: meta?.trace_id ?? env.traceId ?? null,
    priority: meta?.priority ?? env.priority ?? 0,
    input: env.input ?? null,
    output: run.output,
    error: run.error ? { name: run.error.name, message: run.error.message } : null,
    cancelRequestedAt: meta?.cancel_requested_at ?? null,
    cancelReason: meta?.cancel_reason ?? null,
    availableAt: run.availableAt,
    createdAt: run.createdAt,
    startedAt: run.startedAt,
    finishedAt: run.finishedAt,
  };
}