Jyotisha/frontend/src/lib/consultation-billing.ts

import { z } from "zod";

const creditResultSchema = z.object({
  success: z.boolean(),
  credits: z.number().int().nullable(),
  error_code: z.string().nullable().optional(),
});

const authorizationSchema = z.object({
  success: z.boolean(),
  reservation_id: z.string().uuid().nullable(),
  source: z.enum(["subscription", "credits"]).nullable(),
  credits: z.number().int().nullable(),
  subscription_id: z.string().uuid().nullable(),
  reason: z.string().nullable(),
  retry_after_seconds: z.number().int().nullable(),
});

const settlementSchema = z.object({
  success: z.boolean(),
  reservation_id: z.string().uuid().nullable(),
  credits: z.number().int().nullable(),
  error_code: z.string().nullable(),
});

type CreditRpcName = "begin_consultation_credit" | "complete_consultation_credit" | "cancel_consultation_credit";
type AccountingClient = {
  rpc(rpcName: string, args: Record<string, unknown>): PromiseLike<{
    data: unknown;
    error: { message: string } | null;
  }>;
};

export type CreditResult = z.infer<typeof creditResultSchema>;
export type UsageAuthorization = z.infer<typeof authorizationSchema>;
export type UsageSettlement = z.infer<typeof settlementSchema>;

export type ActualUsage = {
  eventKey: string;
  actualModelId: string;
  modelConfigVersion?: number;
  inputTokens: number;
  outputTokens: number;
  costMicrousd: number;
  durationMs: number;
};

export class CreditRpcError extends Error {
  readonly code: string;

  constructor(code: string) {
    super(`Credit operation failed: ${code}`);
    this.name = "CreditRpcError";
    this.code = code;
  }
}

function first(value: unknown) {
  return Array.isArray(value) ? value[0] : value;
}

async function runRpc<T>(
  accounting: AccountingClient,
  rpcName: string,
  args: Record<string, unknown>,
  schema: z.ZodType<T>,
): Promise<T> {
  let lastError = "unknown_billing_error";
  for (let attempt = 1; attempt <= 3; attempt += 1) {
    try {
      const { data, error } = await accounting.rpc(rpcName, args);
      const parsed = schema.safeParse(first(data));
      if (!error && parsed.success) return parsed.data;
      lastError = error?.message || "invalid_billing_response";
    } catch (error) {
      lastError = error instanceof Error ? error.message : "billing_request_failed";
    }
    if (attempt < 3) await new Promise((resolve) => setTimeout(resolve, attempt * 150));
  }
  throw new CreditRpcError(lastError);
}

export async function runCreditRpc(
  accounting: AccountingClient,
  rpcName: CreditRpcName,
  userId: string,
  requestId: string,
): Promise<CreditResult> {
  return runRpc(accounting, rpcName, {
    p_user_id: userId,
    p_request_id: requestId,
  }, creditResultSchema);
}

export async function authorizeUsage(
  accounting: AccountingClient,
  input: {
    userId: string;
    requestId: string;
    featureKey: "chat.standard" | "chat.premium" | "rectification" | "report.full" | "report.export";
    requestedModelId: string;
    creditCost: number;
  },
): Promise<UsageAuthorization> {
  return runRpc(accounting, "authorize_usage", {
    p_user_id: input.userId,
    p_feature_key: input.featureKey,
    p_requested_model_id: input.requestedModelId,
    p_request_id: input.requestId,
    p_credit_cost: input.creditCost,
  }, authorizationSchema);
}

export async function completeUsage(
  accounting: AccountingClient,
  userId: string,
  requestId: string,
  usage: ActualUsage,
): Promise<UsageSettlement> {
  return runRpc(accounting, "complete_usage", {
    p_user_id: userId,
    p_request_id: requestId,
    p_actual_usage: {
      eventKey: usage.eventKey,
      actualModelId: usage.actualModelId,
      modelConfigVersion: usage.modelConfigVersion,
      inputTokens: Math.max(0, Math.trunc(usage.inputTokens)),
      outputTokens: Math.max(0, Math.trunc(usage.outputTokens)),
      costMicrousd: Math.max(0, Math.trunc(usage.costMicrousd)),
      durationMs: Math.max(0, Math.trunc(usage.durationMs)),
    },
  }, settlementSchema);
}

export async function releaseUsage(
  accounting: AccountingClient,
  userId: string,
  requestId: string,
  reason: string,
): Promise<UsageSettlement> {
  return runRpc(accounting, "release_usage", {
    p_user_id: userId,
    p_request_id: requestId,
    p_reason: reason,
  }, settlementSchema);
}