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.") class WalletAdjustmentError(BillingOperationError): code = "invalid_wallet_adjustment" class WalletAdjustmentWouldOverdrawError(WalletAdjustmentError): code = "adjustment_would_overdraw" def __init__(self, *, points_delta: int, available_points: int): self.points_delta = points_delta self.available_points = available_points super().__init__("Wallet adjustment would make the balance negative.") @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 @dataclass(frozen=True) class WalletAdjustment: wallet: UserWallet ledger_entry: PointsLedger points_delta: int balance_after: int 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 _validate_adjustment_delta(points_delta: int) -> int: if ( isinstance(points_delta, bool) or not isinstance(points_delta, int) or points_delta == 0 ): raise WalletAdjustmentError("points_delta must be a non-zero integer.") return points_delta def _normalize_adjustment_reason(reason: str) -> str: normalized_reason = str(reason or "").strip() if not normalized_reason: raise WalletAdjustmentError("Adjustment reason is required.") return normalized_reason def _format_operator(actor) -> str: if actor is None: return "" username_getter = getattr(actor, "get_username", None) username = ( username_getter() if callable(username_getter) else getattr(actor, "username", "") ) actor_pk = getattr(actor, "pk", None) if username and actor_pk: return f"{username}#{actor_pk}" if username: return str(username) if actor_pk: return f"id#{actor_pk}" return str(actor) 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 adjust_wallet_points( *, user, points_delta: int, reason: str, actor=None, ) -> WalletAdjustment: points_delta = _validate_adjustment_delta(points_delta) normalized_reason = _normalize_adjustment_reason(reason) operator = _format_operator(actor) ledger_reason = ( f"{normalized_reason}\n\noperator={operator}" if operator else normalized_reason ) with transaction.atomic(): wallet = _locked_wallet_for_user(user) balance_after = wallet.points_balance + points_delta if balance_after < 0: raise WalletAdjustmentWouldOverdrawError( points_delta=points_delta, available_points=wallet.points_balance, ) wallet.points_balance = balance_after wallet.save(update_fields=("points_balance", "updated_at")) ledger_entry = PointsLedger.objects.create( user=user, change_type=PointsLedger.ChangeType.ADJUST, points_delta=points_delta, balance_after=wallet.points_balance, reason=ledger_reason, ) return WalletAdjustment( wallet=wallet, ledger_entry=ledger_entry, points_delta=points_delta, balance_after=ledger_entry.balance_after, ) 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, )