import { validateGroupBatch, groupDependencyState } from "./group-execution";
import { canResumeCanvasTask, hasRecoverableCanvasTask } from "./task-recovery";
import "server-only";
import { and, asc, desc, eq, gte, lt, sql, inArray } from "drizzle-orm";
import { db } from "@/lib/db";
import { canvasCapabilityTask as table, canvasProject } from "@/lib/db/schema";
import {
  type CanvasTaskRequest,
  type CanvasGenerationTask,
} from "./task-contract";

const capabilityId = "video-canvas.generation";
const scope = (userId: string, projectId: string) =>
  and(
    eq(table.userId, userId),
    eq(table.projectId, projectId),
    eq(table.capabilityId, capabilityId)
  );
const serialize = (row: typeof table.$inferSelect): CanvasGenerationTask => ({
  id: row.id,
  requestId: row.requestId,
  status: row.status,
  createdAt: row.createdAt.toISOString(),
  updatedAt: row.updatedAt.toISOString(),
  attempts: row.attempts,
  request: row.request as CanvasTaskRequest,
  result: row.result as Record<string, any> | null,
  error: row.error,
});
export async function ownsTaskProject(userId: string, projectId: string) {
  return Boolean(
    (
      await db
        .select({ id: canvasProject.id })
        .from(canvasProject)
        .where(
          and(eq(canvasProject.id, projectId), eq(canvasProject.userId, userId))
        )
        .limit(1)
    )[0]
  );
}
export async function createCanvasGenerationTask(
  userId: string,
  request: CanvasTaskRequest,
  beforeCreate?: () => Promise<void>
) {
  return db.transaction(async (tx) => {
    await tx.execute(
      sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${request.projectId}))`
    );
    const existing = (
      await tx
        .select()
        .from(table)
        .where(
          and(
            scope(userId, request.projectId),
            eq(table.requestId, request.requestId)
          )
        )
        .limit(1)
    )[0];
    if (existing) {
      if (
        JSON.stringify(existing.request) !==
        JSON.stringify(JSON.parse(JSON.stringify(request)))
      ) {
        // JSONB property order is not stable. Compare the immutable logical request.
        const old = existing.request as CanvasTaskRequest;
        if (
          old.node.id !== request.node.id ||
          old.kind !== request.kind ||
          JSON.stringify(sortObject(old.attachment)) !==
            JSON.stringify(sortObject(request.attachment)) ||
          JSON.stringify(sortObject(old.params)) !==
            JSON.stringify(sortObject(request.params))
        )
          throw Object.assign(Error("同一任务 ID 不能用于不同请求"), {
            status: 409,
          });
      }
      return serialize(existing);
    }
    const active = await tx
      .select()
      .from(table)
      .where(
        and(
          scope(userId, request.projectId),
          sql`${table.status} in ('pending','running')`
        )
      );
    const duplicate = active.find(
      (row) => (row.request as CanvasTaskRequest).node.id === request.node.id
    );
    if (duplicate) return serialize(duplicate);
    const interrupted = await tx
      .select()
      .from(table)
      .where(
        and(scope(userId, request.projectId), eq(table.status, "interrupted"))
      );
    const recovering = interrupted.find(
      (row) =>
        (row.request as CanvasTaskRequest).node.id === request.node.id &&
        hasRecoverableCanvasTask(serialize(row))
    );
    if (recovering) return serialize(recovering);

    if (active.length >= 40)
      throw Object.assign(
        Error("当前画布排队任务已达 40 个，请等待完成后继续"),
        { status: 429 }
      );
    await beforeCreate?.();
    const [created] = await tx
      .insert(table)
      .values({
        requestId: request.requestId,
        userId,
        projectId: request.projectId,
        capabilityId,
        source: "video-canvas",
        request,
        createdAt: new Date(),
        updatedAt: new Date(),
      })
      .returning();
    return serialize(created);
  });
}
function sortObject(value: any): any {
  return Array.isArray(value)
    ? value.map(sortObject)
    : value && typeof value === "object"
      ? Object.fromEntries(
          Object.keys(value)
            .sort()
            .map((key) => [key, sortObject(value[key])])
        )
      : value;
}
export async function listCanvasGenerationTasks(
  userId: string,
  projectId: string
) {
  // A worker lost to process shutdown must never silently re-submit a paid request.
  await db
    .update(table)
    .set({
      status: "interrupted",
      error:
        "服务端执行连接中断，供应商任务可能仍在运行。请核实结果后再决定是否重新生成，系统不会自动重复提交。",
      updatedAt: new Date(),
    })
    .where(
      and(
        scope(userId, projectId),
        eq(table.status, "running"),
        lt(table.updatedAt, new Date(Date.now() - 120000))
      )
    );
  const rows = await db
    .select()
    .from(table)
    .where(scope(userId, projectId))
    .orderBy(desc(table.createdAt))
    .limit(240);
  return rows.map(serialize).reverse();
}
export async function claimCanvasGenerationTask(
  userId: string,
  projectId: string
) {
  return db.transaction(async (tx) => {
    await tx.execute(
      sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${projectId}))`
    );
    if (
      (
        await tx
          .select({ id: table.id })
          .from(table)
          .where(and(scope(userId, projectId), eq(table.status, "running")))
          .limit(1)
      ).length
    )
      return null;
    const pending = await tx
      .select()
      .from(table)
      .where(and(scope(userId, projectId), eq(table.status, "pending")))
      .orderBy(asc(table.createdAt))
      .limit(40);
    const dependencyIds = [
      ...new Set(
        pending.flatMap(
          (row) =>
            (row.request as CanvasTaskRequest).workflow?.inputs
              .map((input) => input.requestId)
              .filter((id): id is string => Boolean(id)) || []
        )
      ),
    ];
    const predecessors = dependencyIds.length
      ? await tx
          .select()
          .from(table)
          .where(
            and(
              scope(userId, projectId),
              inArray(table.requestId, dependencyIds)
            )
          )
      : [];
    const tasks = [
      ...new Map(
        [...pending, ...predecessors].map((row) => [row.id, serialize(row)])
      ).values(),
    ];
    let next: typeof table.$inferSelect | undefined;
    for (const candidate of pending) {
      const state = groupDependencyState(
        candidate.request as CanvasTaskRequest,
        tasks
      );
      if (state === "failed") {
        await tx
          .update(table)
          .set({
            status: "failed",
            error: "上游任务失败或取消，未提交此节点",
            updatedAt: new Date(),
            completedAt: new Date(),
          })
          .where(eq(table.id, candidate.id));
        tasks.find((t) => t.id === candidate.id)!.status = "failed";
      } else if (state === "ready" && !next) next = candidate;
    }
    if (!next) {
      const interrupted = await tx
        .select()
        .from(table)
        .where(
          and(
            scope(userId, projectId),
            eq(table.status, "interrupted"),
            gte(table.createdAt, new Date(Date.now() - 86400000)),
            lt(table.attempts, 20)
          )
        )
        .orderBy(asc(table.updatedAt));
      next = interrupted.find((row) => canResumeCanvasTask(serialize(row)))!;
    }
    if (!next) return null;
    const [claimed] = await tx
      .update(table)
      .set({
        status: "running",
        error: null,
        startedAt: new Date(),
        updatedAt: new Date(),
        attempts: sql`${table.attempts}+1`,
      })
      .where(and(eq(table.id, next.id), eq(table.status, next.status)))
      .returning();
    return claimed ? serialize(claimed) : null;
  });
}
export async function heartbeatCanvasGenerationTask(
  id: string,
  attempt?: number
) {
  await db
    .update(table)
    .set({ updatedAt: new Date() })
    .where(
      and(
        eq(table.id, id),
        eq(table.status, "running"),
        attempt === undefined ? undefined : eq(table.attempts, attempt)
      )
    );
}
export async function finishCanvasGenerationTask(
  id: string,
  result?: Record<string, any>,
  error?: string,
  attempt?: number
) {
  // A late real result can recover an interrupted task, but cannot replace a newer request ID.
  await db
    .update(table)
    .set({
      status: result ? "succeeded" : "failed",
      result: result || null,
      error: error || null,
      updatedAt: new Date(),
      completedAt: new Date(),
    })
    .where(
      and(
        eq(table.id, id),
        sql`${table.status} in ('running','interrupted')`,
        attempt === undefined ? undefined : eq(table.attempts, attempt)
      )
    );
}
export async function cancelCanvasGenerationTask(
  userId: string,
  projectId: string,
  id: string
) {
  const [cancelled] = await db
    .update(table)
    .set({
      status: "cancelled",
      error: "已取消排队，未提交供应商",
      updatedAt: new Date(),
      completedAt: new Date(),
    })
    .where(
      and(
        scope(userId, projectId),
        eq(table.id, id),
        eq(table.status, "pending")
      )
    )
    .returning();
  if (cancelled) return { ok: true, task: serialize(cancelled) };
  const [current] = await db
    .select()
    .from(table)
    .where(and(scope(userId, projectId), eq(table.id, id)))
    .limit(1);
  return {
    ok: current?.status === "cancelled",
    task: current ? serialize(current) : null,
  };
}
export async function hideCanvasGenerationNode(
  userId: string,
  projectId: string,
  id: string,
  nodeId: string
) {
  return db.transaction(async (tx) => {
    const [task] = await tx
      .select()
      .from(table)
      .where(and(scope(userId, projectId), eq(table.id, id)))
      .for("update");
    if (!task) return false;
    const request = task.request as CanvasTaskRequest & {
      hiddenNodeIds?: string[];
    };
    if (
      nodeId !== request.node.id &&
      !nodeId.startsWith(`${request.node.id}-result-`)
    )
      return false;
    await tx
      .update(table)
      .set({
        request: {
          ...request,
          hiddenNodeIds: [
            ...new Set([...(request.hiddenNodeIds || []), nodeId]),
          ],
        },
      })
      .where(eq(table.id, id));
    return true;
  });
}

export async function checkpointCanvasGenerationTask(
  id: string,
  result: Record<string, unknown>,
  attempt?: number
) {
  await db
    .update(table)
    .set({ result, updatedAt: new Date() })
    .where(
      and(
        eq(table.id, id),
        eq(table.status, "running"),
        attempt === undefined ? undefined : eq(table.attempts, attempt)
      )
    );
}
export async function interruptCanvasGenerationTask(
  id: string,
  error: string,
  attempt?: number
) {
  await db
    .update(table)
    .set({ status: "interrupted", error, updatedAt: new Date() })
    .where(
      and(
        eq(table.id, id),
        eq(table.status, "running"),
        attempt === undefined ? undefined : eq(table.attempts, attempt)
      )
    );
}

export async function ownsCanvasTaskLease(id: string, attempt: number) {
  const [row] = await db
    .select({ id: table.id })
    .from(table)
    .where(
      and(
        eq(table.id, id),
        eq(table.attempts, attempt),
        eq(table.status, "running")
      )
    )
    .limit(1);
  return Boolean(row);
}

export async function createCanvasGroupTasks(
  userId: string,
  requests: CanvasTaskRequest[],
  beforeCreate?: () => Promise<void>
) {
  validateGroupBatch(requests);
  return db.transaction(async (tx) => {
    const projectId = requests[0].projectId;
    await tx.execute(
      sql`select pg_advisory_xact_lock(hashtext(${capabilityId}), hashtext(${projectId}))`
    );
    const same = await tx
      .select()
      .from(table)
      .where(
        and(
          scope(userId, projectId),
          sql`${table.request}->'workflow'->>'runId' = ${requests[0].workflow!.runId}`
        )
      );
    if (same.length) {
      if (
        same.length !== requests.length ||
        requests.some((r) => {
          const old = same.find((t) => t.requestId === r.requestId);
          return (
            !old ||
            JSON.stringify(sortObject(old.request)) !==
              JSON.stringify(sortObject(r))
          );
        })
      )
        throw Object.assign(Error("同一编组执行编号不能用于不同参数"), {
          status: 409,
        });
      return requests.map((request) =>
        serialize(same.find((row) => row.requestId === request.requestId)!)
      );
    }
    const candidates = await tx
      .select()
      .from(table)
      .where(
        and(
          scope(userId, projectId),
          sql`(${table.status} in ('pending','running') or (${table.status}='interrupted' and ${table.createdAt} > now() - interval '24 hours'))`
        )
      );
    const active = candidates.filter(
      (t) =>
        ["pending", "running"].includes(t.status) ||
        (t.status === "interrupted" && hasRecoverableCanvasTask(serialize(t)))
    );
    if (active.length + requests.length > 40)
      throw Object.assign(Error("画布任务队列已满"), { status: 429 });
    if (
      active.some((t) =>
        requests.some(
          (r) => r.node.id === (t.request as CanvasTaskRequest).node.id
        )
      )
    )
      throw Object.assign(Error("编组中已有节点正在执行，请等待完成"), {
        status: 409,
      });
    await beforeCreate?.();
    const now = Date.now();
    const rows = await tx
      .insert(table)
      .values(
        requests.map((request, i) => ({
          requestId: request.requestId,
          userId,
          projectId,
          capabilityId,
          source: "video-canvas",
          request,
          createdAt: new Date(now + i),
          updatedAt: new Date(now),
        }))
      )
      .returning();
    return rows.map(serialize);
  });
}
export async function groupTaskInputs(
  userId: string,
  request: CanvasTaskRequest
) {
  const rows = await db
    .select()
    .from(table)
    .where(
      and(
        scope(userId, request.projectId),
        sql`${table.request}->'workflow'->>'runId' = ${request.workflow?.runId || ""}`
      )
    );
  return rows
    .map(serialize)
    .filter((t) => t.request.workflow?.runId === request.workflow?.runId);
}