import "server-only"; import { and, asc, desc, eq, inArray, sql, type SQL, } from "drizzle-orm"; import { db } from "@/lib/db"; import { videoStoryboardShotTask, videoStoryboardTask, } from "@/lib/db/schema"; import type { PreparedStoryboardShot } from "@/lib/video-agent/types"; export type VideoStoryboardTaskStatus = | "pending" | "running" | "partial_success" | "succeeded" | "failed" | "cancelled"; export type VideoStoryboardShotStatus = | "pending" | "running" | "completed" | "failed" | "cancelled"; function toConfidenceInt(confidence: number) { if (!Number.isFinite(confidence)) return 0; return Math.max(0, Math.min(10000, Math.round(confidence * 10000))); } export async function createVideoStoryboardTask(params: { userId: string; chatId?: string; providerType: string; modelRef?: string; sourceImageUrl: string; isNineGrid: boolean; detectionConfidence: number; normalizedScript?: Record | null; totalShots: number; status?: VideoStoryboardTaskStatus; }) { const now = new Date(); const [task] = await db .insert(videoStoryboardTask) .values({ userId: params.userId, chatId: params.chatId || null, providerType: params.providerType, modelRef: params.modelRef || null, sourceImageUrl: params.sourceImageUrl, isNineGrid: params.isNineGrid, detectionConfidence: toConfidenceInt(params.detectionConfidence), normalizedScript: params.normalizedScript || null, status: params.status || "running", currentShot: 0, totalShots: Math.max(0, Math.round(params.totalShots)), createdAt: now, updatedAt: now, startedAt: now, }) .returning(); return task; } export async function createVideoStoryboardShotTasks(params: { parentTaskId: string; shots: PreparedStoryboardShot[]; fallbackChain: string[]; }) { const now = new Date(); if (!params.shots.length) return []; const inserted = await db .insert(videoStoryboardShotTask) .values( params.shots.map((shot) => ({ parentTaskId: params.parentTaskId, shotIndex: shot.shotIndex, tileIndex: shot.tileIndex, tileImageUrl: shot.tileImageUrl, upscaledImageUrl: shot.upscaled ? shot.inputImageUrl : null, prompt: shot.prompt, modelId: null, fallbackChain: params.fallbackChain, attemptCount: 0, status: "pending", createdAt: now, updatedAt: now, })) ) .returning(); return inserted; } export async function attachVideoStoryboardTaskMessage(params: { taskId: string; messageId: string; }) { await db .update(videoStoryboardTask) .set({ messageId: params.messageId, updatedAt: new Date(), }) .where(eq(videoStoryboardTask.id, params.taskId)); } export async function markVideoStoryboardShotRunning(params: { taskId: string; shotIndex: number; modelId: string; }) { const now = new Date(); await db .update(videoStoryboardShotTask) .set({ status: "running", modelId: params.modelId, startedAt: now, completedAt: null, errorCode: null, errorMessage: null, updatedAt: now, attemptCount: sql`${videoStoryboardShotTask.attemptCount} + 1`, }) .where( and( eq(videoStoryboardShotTask.parentTaskId, params.taskId), eq(videoStoryboardShotTask.shotIndex, params.shotIndex) ) ); await db .update(videoStoryboardTask) .set({ status: "running", currentShot: params.shotIndex, updatedAt: now, }) .where(eq(videoStoryboardTask.id, params.taskId)); } export async function markVideoStoryboardShotRetry(params: { taskId: string; shotIndex: number; errorMessage: string; }) { await db .update(videoStoryboardShotTask) .set({ status: "running", errorMessage: params.errorMessage.slice(0, 4000), updatedAt: new Date(), }) .where( and( eq(videoStoryboardShotTask.parentTaskId, params.taskId), eq(videoStoryboardShotTask.shotIndex, params.shotIndex) ) ); } export async function markVideoStoryboardShotCompleted(params: { taskId: string; shotIndex: number; modelId: string; videoUrl: string; }) { const now = new Date(); await db .update(videoStoryboardShotTask) .set({ status: "completed", modelId: params.modelId, videoUrl: params.videoUrl, errorCode: null, errorMessage: null, completedAt: now, updatedAt: now, }) .where( and( eq(videoStoryboardShotTask.parentTaskId, params.taskId), eq(videoStoryboardShotTask.shotIndex, params.shotIndex) ) ); } export async function markVideoStoryboardShotFailed(params: { taskId: string; shotIndex: number; errorMessage: string; }) { const now = new Date(); await db .update(videoStoryboardShotTask) .set({ status: "failed", errorCode: "shot_failed", errorMessage: params.errorMessage.slice(0, 4000), completedAt: now, updatedAt: now, }) .where( and( eq(videoStoryboardShotTask.parentTaskId, params.taskId), eq(videoStoryboardShotTask.shotIndex, params.shotIndex) ) ); } export async function finalizeVideoStoryboardTask(params: { taskId: string; status: VideoStoryboardTaskStatus; error?: string; }) { const now = new Date(); await db .update(videoStoryboardTask) .set({ status: params.status, error: params.error ? params.error.slice(0, 4000) : null, completedAt: params.status === "succeeded" || params.status === "partial_success" || params.status === "failed" || params.status === "cancelled" ? now : null, updatedAt: now, }) .where(eq(videoStoryboardTask.id, params.taskId)); } export async function syncVideoStoryboardTaskStatus(taskId: string) { const rows = await db .select({ status: videoStoryboardShotTask.status, }) .from(videoStoryboardShotTask) .where(eq(videoStoryboardShotTask.parentTaskId, taskId)); const total = rows.length; const success = rows.filter((row) => row.status === "completed").length; const failed = rows.filter((row) => row.status === "failed").length; const running = rows.filter((row) => row.status === "running").length; const pending = rows.filter((row) => row.status === "pending").length; let status: VideoStoryboardTaskStatus = "pending"; if (running > 0) { status = "running"; } else if (success === total && total > 0) { status = "succeeded"; } else if (failed === total && total > 0) { status = "failed"; } else if (success > 0 && failed > 0) { status = "partial_success"; } else if (success > 0 && pending > 0) { status = "running"; } else if (failed > 0 && pending > 0) { status = "running"; } const currentShot = await db .select({ maxShot: sql`max(${videoStoryboardShotTask.shotIndex})::int`, }) .from(videoStoryboardShotTask) .where( and( eq(videoStoryboardShotTask.parentTaskId, taskId), inArray(videoStoryboardShotTask.status, ["running", "completed", "failed"]) ) ); await db .update(videoStoryboardTask) .set({ status, currentShot: Number(currentShot[0]?.maxShot || 0), completedAt: status === "succeeded" || status === "partial_success" || status === "failed" ? new Date() : null, updatedAt: new Date(), }) .where(eq(videoStoryboardTask.id, taskId)); return { total, success, failed, status }; } export async function getVideoStoryboardTaskDetail(taskId: string) { const taskRows = await db .select() .from(videoStoryboardTask) .where(eq(videoStoryboardTask.id, taskId)) .limit(1); const task = taskRows[0]; if (!task) return null; const shots = await db .select() .from(videoStoryboardShotTask) .where(eq(videoStoryboardShotTask.parentTaskId, taskId)) .orderBy(asc(videoStoryboardShotTask.shotIndex)); return { task, shots }; } export async function listVideoStoryboardTasks(params: { page: number; pageSize: number; status?: string; providerType?: string; chatId?: string; userId?: string; }) { const conditions: SQL[] = []; if (params.status) { conditions.push(eq(videoStoryboardTask.status, params.status)); } if (params.providerType) { conditions.push(eq(videoStoryboardTask.providerType, params.providerType)); } if (params.chatId) { conditions.push(eq(videoStoryboardTask.chatId, params.chatId)); } if (params.userId) { conditions.push(eq(videoStoryboardTask.userId, params.userId)); } const whereClause = conditions.length > 0 ? and(...conditions) : undefined; const page = Math.max(1, Math.round(params.page || 1)); const pageSize = Math.max(1, Math.min(200, Math.round(params.pageSize || 20))); const offset = (page - 1) * pageSize; const [countRow] = await db .select({ count: sql`count(*)::int` }) .from(videoStoryboardTask) .where(whereClause); const tasks = await db .select() .from(videoStoryboardTask) .where(whereClause) .orderBy(desc(videoStoryboardTask.createdAt)) .limit(pageSize) .offset(offset); if (tasks.length === 0) { return { total: Number(countRow?.count || 0), page, pageSize, tasks: [] as Array, }; } const taskIds = tasks.map((item) => item.id); const shotRows = await db .select({ parentTaskId: videoStoryboardShotTask.parentTaskId, status: videoStoryboardShotTask.status, }) .from(videoStoryboardShotTask) .where(inArray(videoStoryboardShotTask.parentTaskId, taskIds)); const summaryMap = new Map< string, { totalShots: number; completedShots: number; failedShots: number; runningShots: number } >(); for (const row of shotRows) { const summary = summaryMap.get(row.parentTaskId) || { totalShots: 0, completedShots: 0, failedShots: 0, runningShots: 0 }; summary.totalShots += 1; if (row.status === "completed") summary.completedShots += 1; if (row.status === "failed") summary.failedShots += 1; if (row.status === "running") summary.runningShots += 1; summaryMap.set(row.parentTaskId, summary); } return { total: Number(countRow?.count || 0), page, pageSize, tasks: tasks.map((task) => ({ ...task, summary: summaryMap.get(task.id) || { totalShots: 0, completedShots: 0, failedShots: 0, runningShots: 0, }, })), }; }