ai-chatbot/lib/video-agent/task-store.ts

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<string, any> | 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<number>`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<number>`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<any>,
    };
  }

  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,
      },
    })),
  };
}