from __future__ import annotations import secrets from dataclasses import dataclass from datetime import datetime from decimal import Decimal, InvalidOperation from django.db import transaction from django.db.models import Sum from django.utils import timezone from apps.users.models import UserWallet from .models import CallRecord, PointsLedger, RechargeOrder, normalize_resolution from .pricing import quote_recharge_points class BillingOperationError(RuntimeError): code = "billing_operation_error" class InsufficientPointsError(BillingOperationError): code = "insufficient_points" def __init__(self, *, required_points: int, available_points: int): self.required_points = required_points self.available_points = available_points super().__init__("Insufficient points. Please recharge before calling this API.") class InvalidCallStateError(BillingOperationError): code = "invalid_call_state" def __init__(self, message: str): super().__init__(message) class RechargeCallbackError(BillingOperationError): code = "recharge_callback_error" class RechargeOrderNotFoundError(RechargeCallbackError): code = "order_not_found" def __init__(self, order_no: str): self.order_no = order_no super().__init__("Recharge order was not found.") class RechargeAmountMismatchError(RechargeCallbackError): code = "amount_mismatch" def __init__(self, *, order_amount: Decimal, callback_amount: Decimal): self.order_amount = order_amount self.callback_amount = callback_amount super().__init__("Payment callback amount does not match local recharge order.") class RechargePayMethodMismatchError(RechargeCallbackError): code = "bad_request" def __init__(self, *, order_pay_method: str, callback_pay_method: str): self.order_pay_method = order_pay_method self.callback_pay_method = callback_pay_method super().__init__("Payment callback method does not match local recharge order.") class InvalidRechargeOrderStateError(RechargeCallbackError): code = "bad_request" def __init__(self, *, status: str): self.status = status super().__init__("Recharge order cannot be paid from its current state.") class RechargeOrderCreateError(BillingOperationError): code = "bad_request" class InvalidRechargePayMethodError(RechargeOrderCreateError): code = "bad_request" def __init__(self, pay_method: str): self.pay_method = pay_method super().__init__("Unsupported recharge payment method.") @dataclass(frozen=True) class CallCharge: call_record: CallRecord ledger_entry: PointsLedger points_cost: int balance_after: int @dataclass(frozen=True) class RefundResult: call_record: CallRecord ledger_entry: PointsLedger | None points_refunded: int balance_after: int refunded: bool @dataclass(frozen=True) class BalanceSnapshot: points_balance: int ledger_balance: int @dataclass(frozen=True) class RechargePayment: order_no: str pay_method: str amount: Decimal transaction_id: str paid_at: datetime | None = None @dataclass(frozen=True) class RechargeResult: order: RechargeOrder ledger_entry: PointsLedger | None points_granted: int balance_after: int applied: bool def _validate_positive_points(points: int) -> int: if isinstance(points, bool) or not isinstance(points, int) or points <= 0: raise ValueError("points must be a positive integer") return points def _locked_wallet_for_user(user) -> UserWallet: wallet, _created = UserWallet.objects.select_for_update().get_or_create(user=user) return wallet def _normalize_money(value) -> Decimal: try: return Decimal(str(value)).quantize(Decimal("0.01")) except (InvalidOperation, TypeError, ValueError) as exc: raise RechargeOrderCreateError("Invalid recharge amount.") from exc def _normalize_pay_method(value: str) -> str: return str(value or "").strip().lower() def _validate_recharge_pay_method(value: str) -> str: pay_method = _normalize_pay_method(value) if pay_method not in { RechargeOrder.PayMethod.WEIXIN, RechargeOrder.PayMethod.ALIPAY, }: raise InvalidRechargePayMethodError(pay_method) return pay_method def _generate_recharge_order_no() -> str: for _attempt in range(10): timestamp = timezone.now().strftime("%Y%m%d%H%M%S") order_no = f"R{timestamp}{secrets.token_hex(4).upper()}" if not RechargeOrder.objects.filter(order_no=order_no).exists(): return order_no raise RechargeOrderCreateError("Could not generate unique recharge order number.") def create_recharge_order( *, user, amount, pay_method: str, currency: str = "CNY", payment_order_func=None, ) -> RechargeOrder: normalized_amount = _normalize_money(amount) if normalized_amount <= 0: raise RechargeOrderCreateError("Recharge amount must be greater than zero.") normalized_pay_method = _validate_recharge_pay_method(pay_method) quote = quote_recharge_points(normalized_amount, currency=currency) if quote.points_granted <= 0: raise RechargeOrderCreateError( "Recharge amount is too small for the current exchange rate." ) order = RechargeOrder.objects.create( user=user, order_no=_generate_recharge_order_no(), amount_money=quote.amount, currency=quote.currency, pay_method=normalized_pay_method, exchange_rate=quote.points_per_unit, points_granted=quote.points_granted, status=RechargeOrder.Status.PENDING, ) if payment_order_func is None: from .payment_gateways import create_payment_order payment_order_func = create_payment_order try: payment_order = payment_order_func(order) code_url = str(getattr(payment_order, "code_url", "") or "").strip() expires_at = getattr(payment_order, "expires_at", None) if not code_url: raise RechargeOrderCreateError("Payment gateway did not return a QR code URL.") except Exception: order.status = RechargeOrder.Status.FAILED order.save(update_fields=("status", "updated_at")) raise order.code_url = code_url order.expires_at = expires_at order.save(update_fields=("code_url", "expires_at", "updated_at")) return order def get_balance_snapshot(user) -> BalanceSnapshot: points_balance = ( UserWallet.objects.filter(user=user) .values_list("points_balance", flat=True) .first() ) ledger_balance = PointsLedger.objects.filter(user=user).aggregate( total=Sum("points_delta") )["total"] or 0 return BalanceSnapshot( points_balance=int(points_balance or 0), ledger_balance=int(ledger_balance), ) def apply_recharge_payment(payment: RechargePayment) -> RechargeResult: order_no = str(payment.order_no or "").strip() pay_method = _normalize_pay_method(payment.pay_method) callback_amount = _normalize_money(payment.amount) paid_at = payment.paid_at or timezone.now() with transaction.atomic(): try: order = ( RechargeOrder.objects.select_for_update() .select_related("user") .get(order_no=order_no) ) except RechargeOrder.DoesNotExist as exc: raise RechargeOrderNotFoundError(order_no) from exc if order.status == RechargeOrder.Status.PAID: ledger_entry = ( PointsLedger.objects.filter( ref_order_id=order.id, change_type=PointsLedger.ChangeType.RECHARGE, ) .order_by("id") .first() ) wallet_balance = ( UserWallet.objects.filter(user=order.user) .values_list("points_balance", flat=True) .first() ) return RechargeResult( order=order, ledger_entry=ledger_entry, points_granted=0, balance_after=int(wallet_balance or 0), applied=False, ) if order.status != RechargeOrder.Status.PENDING: raise InvalidRechargeOrderStateError(status=order.status) if _normalize_pay_method(order.pay_method) != pay_method: raise RechargePayMethodMismatchError( order_pay_method=order.pay_method, callback_pay_method=pay_method, ) order_amount = _normalize_money(order.amount_money) if order_amount != callback_amount: raise RechargeAmountMismatchError( order_amount=order_amount, callback_amount=callback_amount, ) points_granted = _validate_positive_points(order.points_granted) wallet = _locked_wallet_for_user(order.user) wallet.points_balance += points_granted wallet.save(update_fields=("points_balance", "updated_at")) order.status = RechargeOrder.Status.PAID order.payment_txn_no = str(payment.transaction_id or "").strip() order.paid_at = paid_at order.save(update_fields=("status", "payment_txn_no", "paid_at", "updated_at")) ledger_entry = PointsLedger.objects.create( user=order.user, change_type=PointsLedger.ChangeType.RECHARGE, points_delta=points_granted, balance_after=wallet.points_balance, ref_order_id=order.id, reason=f"Recharge paid via {order.pay_method}: {order.order_no}", ) return RechargeResult( order=order, ledger_entry=ledger_entry, points_granted=points_granted, balance_after=ledger_entry.balance_after, applied=True, ) def query_and_apply_recharge_payment(order_no: str, query_func) -> RechargeResult: normalized_order_no = str(order_no or "").strip() try: order = RechargeOrder.objects.get(order_no=normalized_order_no) except RechargeOrder.DoesNotExist as exc: raise RechargeOrderNotFoundError(normalized_order_no) from exc payment = query_func(order) return apply_recharge_payment(payment) def precharge_call( *, user, points_cost: int, operation_type: str, alias: str, model_used: str = "", resolution: str | None = None, api_key=None, prompt: str = "", ) -> CallCharge: points_cost = _validate_positive_points(points_cost) normalized_alias = str(alias or "").strip() normalized_resolution = normalize_resolution(resolution) with transaction.atomic(): wallet = _locked_wallet_for_user(user) if wallet.points_balance < points_cost: raise InsufficientPointsError( required_points=points_cost, available_points=wallet.points_balance, ) wallet.points_balance -= points_cost wallet.save(update_fields=("points_balance", "updated_at")) call_record = CallRecord.objects.create( user=user, api_key=api_key, operation_type=operation_type, alias=normalized_alias, model_used=str(model_used or "").strip(), resolution=normalized_resolution, prompt=str(prompt or ""), points_cost=points_cost, status=CallRecord.Status.PENDING, ) ledger_entry = PointsLedger.objects.create( user=user, change_type=PointsLedger.ChangeType.CONSUME, points_delta=-points_cost, balance_after=wallet.points_balance, ref_call=call_record, ) return CallCharge( call_record=call_record, ledger_entry=ledger_entry, points_cost=points_cost, balance_after=ledger_entry.balance_after, ) def mark_call_success( call_record: CallRecord, *, result_ref: str = "", result_summary: str = "", upstream_latency_ms: int | None = None, ) -> CallRecord: with transaction.atomic(): locked_call = CallRecord.objects.select_for_update().get(pk=call_record.pk) if locked_call.status == CallRecord.Status.FAILED: raise InvalidCallStateError("Cannot mark a failed call as successful.") locked_call.status = CallRecord.Status.SUCCESS locked_call.result_ref = str(result_ref or "") locked_call.result_summary = str(result_summary or "") locked_call.upstream_latency_ms = upstream_latency_ms locked_call.error_message = "" locked_call.save( update_fields=( "status", "result_ref", "result_summary", "upstream_latency_ms", "error_message", "updated_at", ) ) return locked_call def refund_call_points( call_record: CallRecord, *, error_message: str = "", reason: str = "", ) -> RefundResult: with transaction.atomic(): locked_call = CallRecord.objects.select_for_update().get(pk=call_record.pk) if locked_call.status == CallRecord.Status.SUCCESS: raise InvalidCallStateError("Cannot refund a successful call via failure refund.") existing_refund = ( PointsLedger.objects.filter( ref_call=locked_call, change_type=PointsLedger.ChangeType.REFUND, ) .order_by("id") .first() ) if existing_refund is not None: if locked_call.status != CallRecord.Status.FAILED: locked_call.status = CallRecord.Status.FAILED locked_call.error_message = str(error_message or locked_call.error_message or "") locked_call.save(update_fields=("status", "error_message", "updated_at")) return RefundResult( call_record=locked_call, ledger_entry=existing_refund, points_refunded=0, balance_after=existing_refund.balance_after, refunded=False, ) points_cost = _validate_positive_points(locked_call.points_cost) wallet = _locked_wallet_for_user(locked_call.user) wallet.points_balance += points_cost wallet.save(update_fields=("points_balance", "updated_at")) locked_call.status = CallRecord.Status.FAILED locked_call.error_message = str(error_message or "") locked_call.save(update_fields=("status", "error_message", "updated_at")) ledger_reason = str(reason or error_message or "Call failed; refund precharged points.").strip() ledger_entry = PointsLedger.objects.create( user=locked_call.user, change_type=PointsLedger.ChangeType.REFUND, points_delta=points_cost, balance_after=wallet.points_balance, ref_call=locked_call, reason=ledger_reason, ) return RefundResult( call_record=locked_call, ledger_entry=ledger_entry, points_refunded=points_cost, balance_after=ledger_entry.balance_after, refunded=True, )