aiartstudio/server/utils/imageGenerationRecords.ts

import type { IUserBasicData } from "#shared/types";
import type {
  IImageGenerationStatsData,
  IImageHistoryDetail,
  IImageHistoryItem,
  IImageHistoryListData
} from "#shared/types/openai";
import { consola } from "consola";
import { ImageGenerationStatus, Prisma } from "~~/app/generated/prisma/client";
import { prisma } from "~~/server/utils/prisma";

const GLOBAL_STATS_ID = "global";
const DEFAULT_IMAGE_MODEL = "gpt-image-2";
const DEFAULT_IMAGE_SIZE = "1024x1024";

interface IFinishImageGenerationSuccessInput {
  /** NewAPI 上游返回的生成图片 URL */
  imageUrl: string;
  /** Lsky 图床归档后的图片 URL，归档失败时为空 */
  hostedImageUrl: string | null;
  /** 图片 MIME 类型，优先来自图床上传结果 */
  imageMimeType: string | null;
  /** 上游返回的修订提示词，流式生图通常为空 */
  revisedPrompt?: string | null;
  /** 完整上游响应，当前为 chat completions 流式聚合对象 */
  upstreamResponse: unknown;
  /** 完整 Lsky 上传响应，归档失败时为空 */
  hostedResponse: unknown;
  /** 成功状态下的非阻断提示，例如图床归档失败 */
  errorMessage?: string | null;
}

/** 登录或恢复登录时保存 NewAPI 用户快照 */
export const upsertUserSnapshot = (user: IUserBasicData) => {
  return prisma.user.upsert({
    where: {
      id: user.id
    },
    create: {
      id: user.id,
      username: user.username,
      displayName: user.display_name,
      group: user.group,
      role: user.role,
      status: user.status
    },
    update: {
      username: user.username,
      displayName: user.display_name,
      group: user.group,
      role: user.role,
      status: user.status
    }
  });
};

/** 仅有 userId 时兜底创建用户记录，避免旧 cookie 用户生成时外键失败 */
export const ensureUserRecord = (userId: number) => {
  return prisma.user.upsert({
    where: {
      id: userId
    },
    create: {
      id: userId
    },
    update: {}
  });
};

/** 读取用于 Lsky 文件命名的用户标识，缺失时由调用方使用兜底名 */
export const getUserArchiveIdentity = async (userId: number) => {
  const user = await prisma.user.findUnique({
    where: {
      id: userId
    },
    select: {
      username: true,
      displayName: true
    }
  });

  return {
    username: user?.username || user?.displayName || "user"
  };
};

/** 创建进行中的生图记录，并递增全局请求与进行中统计 */
export const createRunningImageGeneration = async (
  userId: number,
  prompt: string
) => {
  await ensureUserRecord(userId);

  const record = await prisma.imageGeneration.create({
    data: {
      userId,
      prompt,
      status: ImageGenerationStatus.RUNNING,
      model: DEFAULT_IMAGE_MODEL,
      size: DEFAULT_IMAGE_SIZE
    }
  });

  await safeUpsertGenerationStats({
    totalRequests: {
      increment: 1
    },
    runningRequests: {
      increment: 1
    }
  });

  return record;
};

/** 将生图记录标记为成功，并保存上游 URL、图床 URL、完整响应和耗时 */
export const finishImageGenerationSuccess = async (
  recordId: bigint,
  startedAt: Date,
  input: IFinishImageGenerationSuccessInput
) => {
  const endedAt = new Date();

  await prisma.imageGeneration.update({
    where: {
      id: recordId
    },
    data: {
      status: ImageGenerationStatus.SUCCEEDED,
      endedAt,
      durationMs: getDurationMs(startedAt, endedAt),
      imageUrl: input.imageUrl,
      hostedImageUrl: input.hostedImageUrl,
      imageMimeType: input.imageMimeType,
      revisedPrompt: input.revisedPrompt || null,
      upstreamResponse: input.upstreamResponse as Prisma.InputJsonValue,
      hostedResponse:
        input.hostedResponse === null
          ? Prisma.DbNull
          : (input.hostedResponse as Prisma.InputJsonValue),
      errorMessage: input.errorMessage || null
    }
  });

  await safeUpsertGenerationStats({
    successRequests: {
      increment: 1
    },
    runningRequests: {
      decrement: 1
    },
    totalImages: {
      increment: 1
    }
  });
};

/** 将生图记录标记为失败，并保存失败原因和耗时 */
export const finishImageGenerationFailed = async (
  recordId: bigint,
  startedAt: Date,
  error: unknown
) => {
  const endedAt = new Date();

  await prisma.imageGeneration.update({
    where: {
      id: recordId
    },
    data: {
      status: ImageGenerationStatus.FAILED,
      endedAt,
      durationMs: getDurationMs(startedAt, endedAt),
      errorMessage: getSafeErrorMessage(error)
    }
  });

  await safeUpsertGenerationStats({
    failedRequests: {
      increment: 1
    },
    runningRequests: {
      decrement: 1
    }
  });
};

/** 查询当前用户未删除的生图历史列表，不返回完整上游/图床响应 */
export const listImageGenerationHistory = async ({
  userId,
  page,
  pageSize
}: {
  userId: number;
  page: number;
  pageSize: number;
}): Promise<IImageHistoryListData> => {
  const where = {
    userId,
    deletedAt: null
  };

  const [total, records] = await prisma.$transaction([
    prisma.imageGeneration.count({ where }),
    prisma.imageGeneration.findMany({
      where,
      orderBy: {
        createdAt: "desc"
      },
      skip: (page - 1) * pageSize,
      take: pageSize
    })
  ]);

  return {
    page,
    pageSize,
    total,
    items: records.map(mapImageGenerationItem)
  };
};

/** 查询当前用户单条生图历史详情，包含完整上游与图床响应 */
export const getImageGenerationDetail = async (
  userId: number,
  recordId: bigint
): Promise<IImageHistoryDetail | null> => {
  const record = await prisma.imageGeneration.findFirst({
    where: {
      id: recordId,
      userId,
      deletedAt: null
    }
  });

  if (!record) return null;

  return {
    ...mapImageGenerationItem(record),
    imageMimeType: record.imageMimeType,
    upstreamResponse: record.upstreamResponse,
    hostedResponse: record.hostedResponse
  };
};

/** 软删除当前用户单条生图历史 */
export const softDeleteImageGeneration = async (
  userId: number,
  recordId: bigint
): Promise<boolean> => {
  const result = await prisma.imageGeneration.updateMany({
    where: {
      id: recordId,
      userId,
      deletedAt: null
    },
    data: {
      deletedAt: new Date()
    }
  });

  return result.count > 0;
};

/** 读取全局生图统计，没有记录时返回全 0 */
export const getGenerationStats =
  async (): Promise<IImageGenerationStatsData> => {
    const stats = await prisma.generationStats.findUnique({
      where: {
        id: GLOBAL_STATS_ID
      }
    });

    return {
      totalRequests: stats?.totalRequests ?? 0,
      successRequests: stats?.successRequests ?? 0,
      failedRequests: stats?.failedRequests ?? 0,
      queuedRequests: stats?.queuedRequests ?? 0,
      runningRequests: stats?.runningRequests ?? 0,
      totalImages: stats?.totalImages ?? 0
    };
  };

/** 统一创建或更新全局统计行 */
const upsertGenerationStats = (update: Prisma.GenerationStatsUpdateInput) => {
  return prisma.generationStats.upsert({
    where: {
      id: GLOBAL_STATS_ID
    },
    create: {
      id: GLOBAL_STATS_ID,
      ...buildStatsCreateInput(update)
    },
    update
  });
};

const safeUpsertGenerationStats = async (
  update: Prisma.GenerationStatsUpdateInput
) => {
  try {
    await upsertGenerationStats(update);
  } catch (error) {
    consola.error("[imageGenerationRecords] 更新生图统计失败", {
      message: error instanceof Error ? error.message : String(error)
    });
  }
};

const buildStatsCreateInput = (update: Prisma.GenerationStatsUpdateInput) => {
  return {
    totalRequests: getIncrementValue(update.totalRequests),
    successRequests: getIncrementValue(update.successRequests),
    failedRequests: getIncrementValue(update.failedRequests),
    queuedRequests: getIncrementValue(update.queuedRequests),
    runningRequests: getIncrementValue(update.runningRequests),
    totalImages: getIncrementValue(update.totalImages)
  };
};

const getIncrementValue = (value: unknown): number => {
  if (
    value &&
    typeof value === "object" &&
    "increment" in value &&
    typeof value.increment === "number"
  ) {
    return value.increment;
  }

  return 0;
};

const mapImageGenerationItem = (record: {
  id: bigint;
  userId: number;
  prompt: string;
  status: ImageGenerationStatus;
  model: string;
  size: string;
  startedAt: Date;
  endedAt: Date | null;
  durationMs: number | null;
  imageUrl: string | null;
  hostedImageUrl: string | null;
  revisedPrompt: string | null;
  errorMessage: string | null;
  createdAt: Date;
}): IImageHistoryItem => {
  return {
    id: record.id.toString(),
    userId: record.userId,
    prompt: record.prompt,
    status: record.status,
    model: record.model,
    size: record.size,
    startedAt: record.startedAt.toISOString(),
    endedAt: record.endedAt?.toISOString() ?? null,
    durationMs: record.durationMs,
    imageUrl: record.imageUrl,
    hostedImageUrl: record.hostedImageUrl,
    revisedPrompt: record.revisedPrompt,
    errorMessage: record.errorMessage,
    createdAt: record.createdAt.toISOString()
  };
};

const getDurationMs = (startedAt: Date, endedAt: Date) => {
  return Math.max(0, endedAt.getTime() - startedAt.getTime());
};

const getSafeErrorMessage = (error: unknown): string => {
  if (error instanceof Error) return error.message;
  if (typeof error === "string") return error;
  return "图片生成失败";
};