import "server-only";

import { randomUUID } from "node:crypto";
import { and, asc, desc, eq, inArray } from "drizzle-orm";
import { db } from "@/lib/db";
import {
  canvasImageTask,
  canvasLayoutReplaceBatch,
  canvasLayoutReplaceBatchItem,
} from "@/lib/db/schema";
import { deriveLayoutReplaceBatchAdvance } from "./layout-replace-batch-state";
import type { CanvasLayoutReplaceTaskRequest } from "./layout-replace-task";

export const CANVAS_LAYOUT_REPLACE_BATCH_LIMIT = 5;

export type CanvasLayoutReplaceBatchSource = {
  sourceImageUrl: string;
  sourceKind: "canvas" | "upload";
  sourceNodeId?: string;
  sourceTitle?: string;
};

export type CanvasLayoutReplaceBatchItemSnapshot = {
  id: string;
  taskId: string;
  position: number;
  sourceNodeId: string | null;
  sourceImageUrl: string;
  sourceTitle: string | null;
  sourceKind: string;
  status: string;
  taskStatus: string;
  resultImageUrl: string | null;
  resultThumbnailUrl: string | null;
  resultModelId: string | null;
  error: string | null;
  addedToCanvas: boolean;
  addedCanvasNodeId: string | null;
};

export type CanvasLayoutReplaceBatchSnapshot = {
  id: string;
  projectId: string;
  status: string;
  referenceImageUrl: string;
  supplementalPrompt: string | null;
  imageSize: string;
  totalCount: number;
  completedCount: number;
  failedCount: number;
  completedAt: Date | null;
  createdAt: Date;
  updatedAt: Date;
  items: CanvasLayoutReplaceBatchItemSnapshot[];
};

type LayoutReplaceTaskResult = {
  images?: Array<{
    originalUrl?: string;
    thumbnailUrl?: string;
  }>;
  model?: string;
  modelRef?: string;
};

function normalizeText(value: unknown, maxLength: number) {
  return String(value || "").trim().slice(0, maxLength);
}

function normalizeHttpUrl(value: unknown) {
  const url = normalizeText(value, 4096);
  return /^https?:\/\//i.test(url) ? url : "";
}

function isTerminalItemStatus(status: string) {
  return status === "succeeded" || status === "failed";
}

function getResultFields(result: LayoutReplaceTaskResult | null | undefined) {
  const image = result?.images?.[0];
  const resultImageUrl = normalizeHttpUrl(image?.originalUrl || image?.thumbnailUrl);
  return {
    resultImageUrl: resultImageUrl || null,
    resultThumbnailUrl:
      normalizeHttpUrl(image?.thumbnailUrl || image?.originalUrl) || resultImageUrl || null,
    resultModelId: normalizeText(result?.model || result?.modelRef, 160) || null,
  };
}

export async function createCanvasLayoutReplaceBatch(params: {
  userId: string;
  projectId: string;
  sources: CanvasLayoutReplaceBatchSource[];
  referenceImageUrl: string;
  supplementalPrompt?: string;
  imageSize?: string;
  model?: string;
}) {
  const sources = params.sources.slice(0, CANVAS_LAYOUT_REPLACE_BATCH_LIMIT);
  if (sources.length === 0) {
    throw new Error("请至少添加一张待处理图片");
  }

  const batchId = randomUUID();
  const referenceImageUrl = normalizeHttpUrl(params.referenceImageUrl);
  const supplementalPrompt = normalizeText(params.supplementalPrompt, 600);
  const imageSize = normalizeText(params.imageSize || "2K", 16) || "2K";
  const model = normalizeText(params.model, 160) || undefined;
  const now = new Date();
  const records = sources.map((source, index) => {
    const taskId = randomUUID();
    const itemId = randomUUID();
    const position = index + 1;
    const taskRequest: CanvasLayoutReplaceTaskRequest = {
      taskKind: "canvas-layout-replace",
      sourceImageUrl: normalizeHttpUrl(source.sourceImageUrl),
      referenceImageUrl,
      supplementalPrompt: supplementalPrompt || undefined,
      model,
      imageSize,
      sourceNodeId: normalizeText(source.sourceNodeId, 160) || undefined,
      referenceLabel: "画布批量排版参考图",
      batchId,
      batchItemId: itemId,
      batchPosition: position,
      batchTotal: sources.length,
      canvasProjectId: params.projectId,
      debugMeta: {
        operation: "layout-replace-batch",
        batchId,
        batchItemId: itemId,
        batchPosition: position,
        batchTotal: sources.length,
      },
    };
    return {
      itemId,
      taskId,
      position,
      source,
      taskRequest,
    };
  });

  if (!referenceImageUrl || records.some((record) => !record.taskRequest.sourceImageUrl)) {
    throw new Error("图片地址不可用，请重新上传后再试");
  }

  await db.transaction(async (tx) => {
    await tx.insert(canvasLayoutReplaceBatch).values({
      id: batchId,
      userId: params.userId,
      projectId: params.projectId,
      status: "pending",
      referenceImageUrl,
      supplementalPrompt: supplementalPrompt || null,
      imageSize,
      totalCount: records.length,
      completedCount: 0,
      failedCount: 0,
      createdAt: now,
      updatedAt: now,
    });

    await tx.insert(canvasImageTask).values(
      records.map((record, index) => ({
        id: record.taskId,
        userId: params.userId,
        status: index === 0 ? "pending" : "waiting",
        request: record.taskRequest as unknown as Record<string, unknown>,
        attempts: 0,
        createdAt: new Date(now.getTime() + index),
        updatedAt: now,
      }))
    );

    await tx.insert(canvasLayoutReplaceBatchItem).values(
      records.map((record, index) => ({
        id: record.itemId,
        batchId,
        taskId: record.taskId,
        position: record.position,
        sourceNodeId: normalizeText(record.source.sourceNodeId, 160) || null,
        sourceImageUrl: record.taskRequest.sourceImageUrl,
        sourceTitle: normalizeText(record.source.sourceTitle, 256) || null,
        sourceKind: record.source.sourceKind === "upload" ? "upload" : "canvas",
        status: index === 0 ? "pending" : "queued",
        createdAt: new Date(now.getTime() + index),
        updatedAt: now,
      }))
    );
  });

  return { batchId, firstTaskId: records[0]?.taskId || "" };
}

export async function markCanvasLayoutReplaceBatchItemRunning(
  request: CanvasLayoutReplaceTaskRequest
) {
  if (!request.batchId || !request.batchItemId) return;
  const now = new Date();
  await db.transaction(async (tx) => {
    await tx
      .update(canvasLayoutReplaceBatchItem)
      .set({ status: "running", error: null, updatedAt: now })
      .where(
        and(
          eq(canvasLayoutReplaceBatchItem.id, request.batchItemId as string),
          eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string)
        )
      );
    await tx
      .update(canvasLayoutReplaceBatch)
      .set({ status: "running", completedAt: null, updatedAt: now })
      .where(eq(canvasLayoutReplaceBatch.id, request.batchId as string));
  });
}

export async function completeCanvasLayoutReplaceBatchItem(params: {
  request: CanvasLayoutReplaceTaskRequest;
  result?: LayoutReplaceTaskResult | null;
  error?: string | null;
}) {
  const { request } = params;
  if (!request.batchId || !request.batchItemId) return null;
  const now = new Date();
  const candidateResultFields = getResultFields(params.result);
  const failed = Boolean(params.error) || !candidateResultFields.resultImageUrl;
  const resultFields = failed ? getResultFields(null) : candidateResultFields;

  return db.transaction(async (tx) => {
    const [completedItem] = await tx
      .update(canvasLayoutReplaceBatchItem)
      .set({
        status: failed ? "failed" : "succeeded",
        ...resultFields,
        error: failed
          ? normalizeText(params.error, 4000) || "排版替换未返回结果图片"
          : null,
        completedAt: now,
        updatedAt: now,
      })
      .where(
        and(
          eq(canvasLayoutReplaceBatchItem.id, request.batchItemId as string),
          eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string),
          inArray(canvasLayoutReplaceBatchItem.status, ["pending", "running"])
        )
      )
      .returning({ id: canvasLayoutReplaceBatchItem.id });

    // Completion can be replayed after a process restart. Only the first state
    // transition may promote the next item, otherwise one batch could fan out.
    if (!completedItem) return null;

    const items = await tx
      .select({
        id: canvasLayoutReplaceBatchItem.id,
        taskId: canvasLayoutReplaceBatchItem.taskId,
        status: canvasLayoutReplaceBatchItem.status,
        position: canvasLayoutReplaceBatchItem.position,
      })
      .from(canvasLayoutReplaceBatchItem)
      .where(eq(canvasLayoutReplaceBatchItem.batchId, request.batchId as string))
      .orderBy(asc(canvasLayoutReplaceBatchItem.position));

    const advance = deriveLayoutReplaceBatchAdvance(items);
    const nextItem = advance.nextTaskId
      ? items.find((item) => item.taskId === advance.nextTaskId)
      : null;

    if (nextItem && !advance.isComplete) {
      await tx
        .update(canvasLayoutReplaceBatchItem)
        .set({ status: "pending", updatedAt: now })
        .where(eq(canvasLayoutReplaceBatchItem.id, nextItem.id));
      await tx
        .update(canvasImageTask)
        .set({ status: "pending", updatedAt: now })
        .where(
          and(
            eq(canvasImageTask.id, nextItem.taskId),
            eq(canvasImageTask.status, "waiting")
          )
        );
    }

    await tx
      .update(canvasLayoutReplaceBatch)
      .set({
        status: advance.status,
        completedCount: advance.succeededCount,
        failedCount: advance.failedCount,
        completedAt: advance.isComplete ? now : null,
        updatedAt: now,
      })
      .where(eq(canvasLayoutReplaceBatch.id, request.batchId as string));

    return nextItem && !advance.isComplete ? nextItem.taskId : null;
  });
}

async function buildBatchSnapshot(batch: typeof canvasLayoutReplaceBatch.$inferSelect) {
  const items = await db
    .select()
    .from(canvasLayoutReplaceBatchItem)
    .where(eq(canvasLayoutReplaceBatchItem.batchId, batch.id))
    .orderBy(asc(canvasLayoutReplaceBatchItem.position));
  const taskIds = items.map((item) => item.taskId);
  const tasks = taskIds.length
    ? await db
        .select({
          id: canvasImageTask.id,
          status: canvasImageTask.status,
          error: canvasImageTask.error,
        })
        .from(canvasImageTask)
        .where(inArray(canvasImageTask.id, taskIds))
    : [];
  const taskById = new Map(tasks.map((task) => [task.id, task]));

  return {
    id: batch.id,
    projectId: batch.projectId,
    status: batch.status,
    referenceImageUrl: batch.referenceImageUrl,
    supplementalPrompt: batch.supplementalPrompt,
    imageSize: batch.imageSize,
    totalCount: batch.totalCount,
    completedCount: batch.completedCount,
    failedCount: batch.failedCount,
    completedAt: batch.completedAt,
    createdAt: batch.createdAt,
    updatedAt: batch.updatedAt,
    items: items.map((item) => {
      const task = taskById.get(item.taskId);
      return {
        id: item.id,
        taskId: item.taskId,
        position: item.position,
        sourceNodeId: item.sourceNodeId,
        sourceImageUrl: item.sourceImageUrl,
        sourceTitle: item.sourceTitle,
        sourceKind: item.sourceKind,
        status: item.status,
        taskStatus: task?.status || item.status,
        resultImageUrl: item.resultImageUrl,
        resultThumbnailUrl: item.resultThumbnailUrl,
        resultModelId: item.resultModelId,
        error: item.error || task?.error || null,
        addedToCanvas: item.addedToCanvas,
        addedCanvasNodeId: item.addedCanvasNodeId,
      } satisfies CanvasLayoutReplaceBatchItemSnapshot;
    }),
  } satisfies CanvasLayoutReplaceBatchSnapshot;
}

export async function getCanvasLayoutReplaceBatchForUser(params: {
  batchId: string;
  userId: string;
}) {
  const [batch] = await db
    .select()
    .from(canvasLayoutReplaceBatch)
    .where(
      and(
        eq(canvasLayoutReplaceBatch.id, params.batchId),
        eq(canvasLayoutReplaceBatch.userId, params.userId)
      )
    )
    .limit(1);
  return batch ? buildBatchSnapshot(batch) : null;
}

export async function getLatestCanvasLayoutReplaceBatchForUser(params: {
  projectId: string;
  userId: string;
  sourceNodeId?: string;
}) {
  const batches = await db
    .select()
    .from(canvasLayoutReplaceBatch)
    .where(
      and(
        eq(canvasLayoutReplaceBatch.projectId, params.projectId),
        eq(canvasLayoutReplaceBatch.userId, params.userId)
      )
    )
    .orderBy(desc(canvasLayoutReplaceBatch.createdAt))
    .limit(12);

  for (const batch of batches) {
    const snapshot = await buildBatchSnapshot(batch);
    if (
      !params.sourceNodeId ||
      snapshot.items.some((item) => item.sourceNodeId === params.sourceNodeId)
    ) {
      return snapshot;
    }
  }
  return null;
}

export async function markCanvasLayoutReplaceBatchItemsAdded(params: {
  batchId: string;
  userId: string;
  items: Array<{ itemId: string; canvasNodeId: string }>;
}) {
  const [ownedBatch] = await db
    .select({ id: canvasLayoutReplaceBatch.id })
    .from(canvasLayoutReplaceBatch)
    .where(
      and(
        eq(canvasLayoutReplaceBatch.id, params.batchId),
        eq(canvasLayoutReplaceBatch.userId, params.userId)
      )
    )
    .limit(1);
  if (!ownedBatch) return false;

  const now = new Date();
  await db.transaction(async (tx) => {
    for (const item of params.items.slice(0, CANVAS_LAYOUT_REPLACE_BATCH_LIMIT)) {
      await tx
        .update(canvasLayoutReplaceBatchItem)
        .set({
          addedToCanvas: true,
          addedCanvasNodeId: normalizeText(item.canvasNodeId, 160) || null,
          updatedAt: now,
        })
        .where(
          and(
            eq(canvasLayoutReplaceBatchItem.id, item.itemId),
            eq(canvasLayoutReplaceBatchItem.batchId, params.batchId),
            eq(canvasLayoutReplaceBatchItem.status, "succeeded")
          )
        );
    }
  });
  return true;
}

export async function retryCanvasLayoutReplaceBatchItem(params: {
  batchId: string;
  itemId: string;
  userId: string;
}) {
  const [ownedBatch] = await db
    .select({
      id: canvasLayoutReplaceBatch.id,
      status: canvasLayoutReplaceBatch.status,
    })
    .from(canvasLayoutReplaceBatch)
    .where(
      and(
        eq(canvasLayoutReplaceBatch.id, params.batchId),
        eq(canvasLayoutReplaceBatch.userId, params.userId)
      )
    )
    .limit(1);
  if (!ownedBatch || !isCanvasLayoutReplaceBatchTerminal(ownedBatch.status)) {
    return null;
  }

  return db.transaction(async (tx) => {
    const now = new Date();
    const [retriedItem] = await tx
      .update(canvasLayoutReplaceBatchItem)
      .set({
        status: "pending",
        resultImageUrl: null,
        resultThumbnailUrl: null,
        resultModelId: null,
        error: null,
        completedAt: null,
        updatedAt: now,
      })
      .where(
        and(
          eq(canvasLayoutReplaceBatchItem.id, params.itemId),
          eq(canvasLayoutReplaceBatchItem.batchId, params.batchId),
          eq(canvasLayoutReplaceBatchItem.status, "failed")
        )
      )
      .returning({ taskId: canvasLayoutReplaceBatchItem.taskId });
    if (!retriedItem) return null;

    await tx
      .update(canvasImageTask)
      .set({
        status: "pending",
        result: null,
        error: null,
        startedAt: null,
        completedAt: null,
        updatedAt: now,
      })
      .where(eq(canvasImageTask.id, retriedItem.taskId));

    const remainingItems = await tx
      .select({ status: canvasLayoutReplaceBatchItem.status })
      .from(canvasLayoutReplaceBatchItem)
      .where(eq(canvasLayoutReplaceBatchItem.batchId, params.batchId));
    await tx
      .update(canvasLayoutReplaceBatch)
      .set({
        status: "running",
        completedCount: remainingItems.filter((item) => item.status === "succeeded")
          .length,
        failedCount: remainingItems.filter((item) => item.status === "failed").length,
        completedAt: null,
        updatedAt: now,
      })
      .where(eq(canvasLayoutReplaceBatch.id, params.batchId));

    return retriedItem.taskId;
  });
}

export function getCanvasLayoutReplaceBatchActiveTaskId(
  snapshot: CanvasLayoutReplaceBatchSnapshot
) {
  if (["succeeded", "failed", "partial_failed"].includes(snapshot.status)) {
    return null;
  }
  const active = snapshot.items.find((item) =>
    ["pending", "running"].includes(item.status)
  );
  return active?.taskId || null;
}

export function isCanvasLayoutReplaceBatchTerminal(status: string) {
  return ["succeeded", "failed", "partial_failed"].includes(status);
}