omniai-server/src/enterpriseVideoBilling.js

"use strict";

const ENTERPRISE_VIDEO_ALLOWED_MODELS = new Set([
  "happyhorse-1.0",
  "happyhorse-1.0-t2v",
  "happyhorse-1.0-i2v",
  "happyhorse-1.0-r2v",
  "wan2.7-i2v",
  "wan2.2-animate-mix",
  "wan2.2-s2v",
  "kling-3.0-dashscope",
  "kling-v3-omni-dashscope",
  "kling/kling-v3-omni-video-generation",
]);

function normalizeModel(value) {
  return String(value || "").trim().toLowerCase();
}

function normalizeEnterpriseVideoResolution(value) {
  return String(value || "").trim().toUpperCase() === "720P" ? "720P" : "1080P";
}

function normalizeEnterpriseVideoDuration(value) {
  const numeric = Number(value);
  if (!Number.isFinite(numeric)) return 1;
  return Math.max(1, Math.ceil(numeric));
}

function isEnterpriseVideoBillingUser(user) {
  return Boolean(user?.enterpriseId);
}

function isEnterpriseVideoModelAllowed(providerConfig, model) {
  const requestedModel = normalizeModel(model || providerConfig?.requestedModel || providerConfig?.model);
  const resolvedModel = normalizeModel(providerConfig?.model);
  const protocol = String(providerConfig?.protocol || "").toLowerCase();

  if (ENTERPRISE_VIDEO_ALLOWED_MODELS.has(requestedModel)) return true;
  if (ENTERPRISE_VIDEO_ALLOWED_MODELS.has(resolvedModel)) return true;
  if (protocol.startsWith("happyhorse-")) return true;
  if (protocol === "wan-i2v") return true;
  if (protocol === "wan-animate-mix") return true;
  if (protocol === "wan-s2v") return true;
  if (protocol === "kling-dashscope") return true;
  return false;
}

function assertEnterpriseVideoModelAllowed(providerConfig, model) {
  if (isEnterpriseVideoModelAllowed(providerConfig, model)) return;
  const error = new Error("该企业账号不可使用当前视频模型");
  error.status = 403;
  error.code = "ENTERPRISE_VIDEO_MODEL_NOT_ALLOWED";
  throw error;
}

function getEnterpriseVideoCreditRate(input) {
  const resolution = normalizeEnterpriseVideoResolution(input.resolution || input.quality);
  const model = normalizeModel(input.model || input.requestedModel);

  if (model.includes("happyhorse")) {
    return resolution === "720P" ? 0.72 : 1.28;
  }

  if (model.includes("wan2.7-i2v") || model.includes("wanxiang")) {
    return resolution === "720P" ? 0.6 : 1;
  }

  if (model.includes("animate-mix") || model.includes("s2v")) {
    return resolution === "720P" ? 0.6 : 1;
  }

  if (model.includes("kling")) {
    if (input.muted) {
      if (input.hasReferenceVideo) return resolution === "720P" ? 0.9 : 1.2;
      return resolution === "720P" ? 0.6 : 0.8;
    }
    return resolution === "720P" ? 0.9 : 1.2;
  }

  const error = new Error(`Unsupported enterprise video model: ${input.model || input.requestedModel}`);
  error.status = 403;
  error.code = "ENTERPRISE_VIDEO_MODEL_NOT_ALLOWED";
  throw error;
}

function calculateEnterpriseVideoCredits(input) {
  const duration = normalizeEnterpriseVideoDuration(input.durationSeconds || input.duration);
  return Number((getEnterpriseVideoCreditRate(input) * duration).toFixed(2));
}

function calculateEnterpriseVideoCost(input) {
  const resolution = normalizeEnterpriseVideoResolution(input.resolution || input.quality);
  const durationSeconds = normalizeEnterpriseVideoDuration(input.durationSeconds || input.duration);
  const rateCreditsPerSecond = getEnterpriseVideoCreditRate({
    ...input,
    resolution,
    durationSeconds,
  });
  const rateCentsPerSecond = Math.round(rateCreditsPerSecond * 100);
  return {
    resolution,
    durationSeconds,
    rateCentsPerSecond,
    amountCents: rateCentsPerSecond * durationSeconds,
  };
}

function prepareEnterpriseVideoBilling({ user, providerConfig, params }) {
  if (!isEnterpriseVideoBillingUser(user)) return null;
  assertEnterpriseVideoModelAllowed(providerConfig, params.requestedModel || params.model);
  const hasReferenceVideo = Boolean(params.hasReferenceVideo);
  const muted = Boolean(params.muted);
  const pricing = calculateEnterpriseVideoCost({
    model: params.model,
    requestedModel: params.requestedModel,
    resolution: params.resolution || params.quality,
    duration: params.duration,
    muted,
    hasReferenceVideo,
  });

  return {
    enterpriseId: user.enterpriseId,
    userId: user.id,
    model: params.model,
    requestedModel: params.requestedModel,
    muted,
    hasReferenceVideo,
    ...pricing,
  };
}

async function reserveEnterpriseVideoCredits(client, billing) {
  if (!billing || billing.amountCents <= 0) return null;

  const { rows: balanceRows } = await client.query(
    "UPDATE enterprises SET balance_cents = balance_cents - $1, updated_at = NOW() WHERE id = $2 AND enabled = 1 AND balance_cents >= $1 RETURNING balance_cents",
    [billing.amountCents, billing.enterpriseId],
  );
  if (balanceRows.length === 0) {
    const error = new Error(
      `企业积分不足，预计需要 ${Number((billing.amountCents / 100).toFixed(2))} 积分`,
    );
    error.status = 402;
    error.code = "INSUFFICIENT_ENTERPRISE_BALANCE";
    throw error;
  }

  const {
    rows: [ledger],
  } = await client.query(
    `
    INSERT INTO credit_ledger (
      enterprise_id,
      user_id,
      task_id,
      model,
      task_type,
      resolution,
      duration_seconds,
      rate_cents_per_second,
      amount_cents,
      status
    )
    VALUES ($1, $2, $3, $4, 'video', $5, $6, $7, $8, 'reserved')
    RETURNING id
  `,
    [
      billing.enterpriseId,
      billing.userId,
      billing.taskId || null,
      billing.requestedModel || billing.model,
      billing.resolution,
      billing.durationSeconds,
      billing.rateCentsPerSecond,
      billing.amountCents,
    ],
  );

  return {
    ...billing,
    creditLedgerId: ledger.id,
    enterpriseBalanceCents: Number(balanceRows[0].balance_cents || 0),
  };
}

async function markEnterpriseVideoCreditsAccepted(clientOrPool, creditLedgerId) {
  if (!creditLedgerId) return false;
  const { rowCount } = await clientOrPool.query(
    "UPDATE credit_ledger SET status = 'charged', updated_at = NOW() WHERE id = $1 AND status = 'reserved'",
    [creditLedgerId],
  );
  return rowCount > 0;
}

module.exports = {
  ENTERPRISE_VIDEO_ALLOWED_MODELS,
  assertEnterpriseVideoModelAllowed,
  calculateEnterpriseVideoCost,
  calculateEnterpriseVideoCredits,
  getEnterpriseVideoCreditRate,
  isEnterpriseVideoBillingUser,
  isEnterpriseVideoModelAllowed,
  markEnterpriseVideoCreditsAccepted,
  normalizeEnterpriseVideoDuration,
  normalizeEnterpriseVideoResolution,
  prepareEnterpriseVideoBilling,
  reserveEnterpriseVideoCredits,
};