Jyotisha/frontend/src/lib/consult/settlement.ts

import "server-only";

import { NextResponse } from "next/server";
import { z } from "zod";

import { mergePromptCacheUsage, promptCacheUsage } from "@/lib/agent-generation-settings";
import { logAgentObservability, type AgentSettlementResult } from "@/lib/agent-observability";
import { CreditRpcError } from "@/lib/consultation-billing";
import { consultationUsageMetadata } from "@/lib/consultation-step-timing";
import type { SmalltalkUsage } from "@/lib/consultation-smalltalk";
import { resolveFeaturePricing } from "@/lib/feature-pricing";

type RpcClient = {
  rpc(name: string, args: Record<string, unknown>): Promise<{
    data: unknown;
    error: { message: string } | null;
  }>;
};

type LatencyView = {
  modelStepTimings?: Parameters<typeof consultationUsageMetadata>[0]["modelSteps"];
  answerReasoningMs?: number | null;
} | null;

export type Usage = {
  inputTokens?: number;
  outputTokens?: number;
  cache?: ReturnType<typeof promptCacheUsage>;
};

export type ConsultSettlementContext = {
  accounting: RpcClient;
  userId: string;
  requestId: string;
  sessionId: string;
  visibleQuestion: string;
  getUsageStartedAt: () => number;
  getSelectedModel: () => {
    id: string;
    configVersion?: number;
    inputCostMicrousdPerMillion?: number;
    outputCostMicrousdPerMillion?: number;
  };
  getClassification: () => {
    usage: SmalltalkUsage;
    outcome: string | undefined;
    durationMs: number | null;
  };
  getLatencyState: () => LatencyView;
  logPrep?: () => void;
  noteAppendPhase?: (status: "completed" | "failed", durationMs: number) => void;
};

export const consultationReservationSchema = 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(),
});

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

export const consultationQuestionAppendSchema = z.object({
  success: z.boolean(),
  error_code: z.string().nullable().optional(),
});

export function first<T>(value: T | T[]): T {
  return Array.isArray(value) ? value[0] : value;
}

export const detachedSettlementAttempts = 3;
// ponytail: Staging MVP ceiling—without a queue/worker, detached settlement gets only three short retries.
export async function retryDetachedSettlement<T>(action: () => Promise<T>): Promise<T> {
  let lastError: unknown;
  for (let attempt = 1; attempt <= detachedSettlementAttempts; attempt += 1) {
    try {
      return await action();
    } catch (error) {
      lastError = error;
      if (attempt < detachedSettlementAttempts) {
        await new Promise((resolve) => setTimeout(resolve, attempt * 150));
      }
    }
  }
  throw lastError;
}

export function mergeUsage(usages: Promise<Usage>[]): Promise<Usage> {
  return Promise.all(usages).then((items) => items.reduce<Usage>((total, item) => {
    const usage = item && typeof item === "object" ? item as Record<string, unknown> : {};
    const cache = item.cache ?? promptCacheUsage(usage);
    return {
      inputTokens: (total.inputTokens ?? 0) + (typeof usage.inputTokens === "number" ? usage.inputTokens : 0),
      outputTokens: (total.outputTokens ?? 0) + (typeof usage.outputTokens === "number" ? usage.outputTokens : 0),
      cache: mergePromptCacheUsage([total.cache, cache]),
    };
  }, {} as Usage));
}

export function createConsultSettlement(ctx: ConsultSettlementContext) {
  const {
    accounting,
    userId,
    requestId,
    sessionId,
    visibleQuestion,
    getUsageStartedAt,
    getSelectedModel,
    getClassification,
    getLatencyState,
    logPrep,
    noteAppendPhase,
  } = ctx;

  async function reserveCredit(model: { id: string }) {
    const pricing = await resolveFeaturePricing(accounting, "chat.standard", model.id);
    const { data, error } = await accounting.rpc("reserve_consultation_usage", {
      p_user_id: userId,
      p_request_id: requestId,
      p_session_id: sessionId,
      p_requested_model_id: model.id,
      p_credit_cost: pricing.credit_cost,
    });
    const reservation = consultationReservationSchema.safeParse(first(data ?? []));
    if (error || !reservation.success) {
      throw new CreditRpcError(error?.message || "invalid_reservation_response");
    }
    return {
      success: reservation.data.success,
      credits: reservation.data.credits,
      error_code: reservation.data.reason,
    };
  }

  async function cancel(): Promise<AgentSettlementResult> {
    try {
      const result = await retryDetachedSettlement(async () => {
        const { data, error } = await accounting.rpc("cancel_consultation_credit", {
          p_user_id: userId,
          p_request_id: requestId,
        });
        const parsedCancellation = consultationCompletionSchema.safeParse(first(data ?? []));
        if (error || !parsedCancellation.success) {
          throw new CreditRpcError(error?.message || "invalid_cancellation_response");
        }
        return parsedCancellation.data;
      });
      if (!result.success && result.error_code !== "request_completed") {
        throw new CreditRpcError(result.error_code || "cancellation_rejected");
      }
      return "cancelled";
    } catch (error) {
      const reason = error instanceof Error ? error.name : "UnknownError";
      console.error(
        `[billing] cancellation failed request=${requestId} reason=${reason}`,
      );
      return "failed";
    }
  }

  async function appendQuestion(): Promise<Response | null> {
    let appendedQuestion: { success: boolean; error_code?: string | null };
    const appendStarted = Date.now();
    try {
      appendedQuestion = await retryDetachedSettlement(async () => {
        const { data, error } = await accounting.rpc("append_consultation_question", {
          p_user_id: userId,
          p_request_id: requestId,
          p_session_id: sessionId,
          p_question_message: { role: "user", text: visibleQuestion },
        });
        const parsedAppend = consultationQuestionAppendSchema.safeParse(first(data ?? []));
        if (error || !parsedAppend.success) {
          throw new CreditRpcError(error?.message || "invalid_question_append_response");
        }
        return parsedAppend.data;
      });
      noteAppendPhase?.("completed", Math.max(0, Date.now() - appendStarted));
    } catch {
      noteAppendPhase?.("failed", Math.max(0, Date.now() - appendStarted));
      logPrep?.();
      await cancel();
      return NextResponse.json(
        { error: "暂时无法保存问题", message: "请稍后重试。" },
        { status: 503 },
      );
    }
    if (!appendedQuestion.success) {
      await cancel();
      if (appendedQuestion.error_code === "session_full") {
        return NextResponse.json(
          {
            error: "这段对话已写满",
            message: "这段对话已写满，开个新对话继续吧",
            code: "session_full",
          },
          { status: 409 },
        );
      }
      return NextResponse.json(
        { error: "暂时无法保存问题", message: "请稍后重试。" },
        { status: 503 },
      );
    }
    return null;
  }

  async function completeSmalltalk(reply: string): Promise<void> {
    const actualUsage = await usagePayload(Promise.resolve({}));
    const completion = await retryDetachedSettlement(async () => {
      const { data, error } = await accounting.rpc("complete_consultation_free", {
        p_user_id: userId,
        p_request_id: requestId,
        p_session_id: sessionId,
        p_response_message: { role: "assistant", text: reply, responseKind: "smalltalk" },
        p_actual_usage: actualUsage,
      });
      const result = consultationCompletionSchema.safeParse(first(data ?? []));
      if (error || !result.success || !result.data.success) {
        throw new CreditRpcError(error?.message || (result.success ? result.data.error_code : "invalid_free_completion") || "free_completion_rejected");
      }
      return result.data;
    });
    logAgentObservability({ requestId, sessionId, agentVersion: "consultation-smalltalk-v1",
      billingSettlementResult: completion.success ? "completed" : "failed" });
  }

  async function usagePayload(usage: Promise<{ inputTokens?: number; outputTokens?: number }>) {
    const selectedModel = getSelectedModel();
    const classification = getClassification();
    const classificationUsage = classification.usage;
    const classificationOutcome = classification.outcome;
    const classificationDurationMs = classification.durationMs;
    const latencyState = getLatencyState();
    const resolved = await mergeUsage([usage, Promise.resolve(classificationUsage)]);
    const usageRecord = resolved as Record<string, unknown>;
    const cache = resolved.cache ?? promptCacheUsage(usageRecord);
    const inputTokens = Math.max(0, Math.trunc(typeof usageRecord.inputTokens === "number" ? usageRecord.inputTokens : 0));
    const outputTokens = Math.max(0, Math.trunc(typeof usageRecord.outputTokens === "number" ? usageRecord.outputTokens : 0));
    const usageStartedAt = getUsageStartedAt();
    return {
      eventKey: requestId,
      actualModelId: selectedModel.id,
      modelConfigVersion: selectedModel.configVersion,
      inputTokens,
      outputTokens,
      costMicrousd: Math.round((
        inputTokens * (selectedModel.inputCostMicrousdPerMillion ?? 0)
        + outputTokens * (selectedModel.outputCostMicrousdPerMillion ?? 0)
      ) / 1_000_000),
      durationMs: Date.now() - usageStartedAt,
      metadata: {
        ...(cache ? { cache: { ...cache, hit: cache.readTokens > 0 } } : {}),
        ...(classificationOutcome || classificationDurationMs !== null ? { classification: {
          ...(classificationOutcome ? { outcome: classificationOutcome } : {}),
          usageKnown: classificationUsage.inputTokens !== undefined || classificationUsage.outputTokens !== undefined,
          inputTokens: classificationUsage.inputTokens ?? null,
          outputTokens: classificationUsage.outputTokens ?? null,
          durationMs: classificationDurationMs,
        } } : {}),
        ...(latencyState?.modelStepTimings === undefined ? {} : consultationUsageMetadata({
          modelSteps: latencyState.modelStepTimings,
          answerReasoningMs: latencyState.answerReasoningMs,
        })),
      },
    };
  }

  return { reserveCredit, cancel, appendQuestion, completeSmalltalk, usagePayload };
}