cmhub/apps/billing/services.py

from __future__ import annotations

import secrets
from dataclasses import dataclass
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,
    SignupBonusGrant,
    normalize_resolution,
)
from .payment_gateways import PaymentReceipt
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


RechargePayment = PaymentReceipt


@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


@dataclass(frozen=True)
class SignupBonusGrantResult:
    wallet: UserWallet
    grant: SignupBonusGrant
    ledger_entry: PointsLedger | None
    points_granted: int
    balance_after: int
    granted: 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 _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 grant_signup_bonus(
    *,
    user,
    points: int = 10,
    reason: str = "new_user_registration",
) -> SignupBonusGrantResult:
    points = _validate_positive_points(points)
    ledger_reason = str(reason or "new_user_registration").strip()

    with transaction.atomic():
        wallet = _locked_wallet_for_user(user)
        grant, created = SignupBonusGrant.objects.get_or_create(
            user=user,
            defaults={"points_granted": points},
        )
        if not created:
            existing_ledger = (
                PointsLedger.objects.filter(
                    user=user,
                    change_type=PointsLedger.ChangeType.SIGNUP_BONUS,
                )
                .order_by("id")
                .first()
            )
            return SignupBonusGrantResult(
                wallet=wallet,
                grant=grant,
                ledger_entry=existing_ledger,
                points_granted=0,
                balance_after=wallet.points_balance,
                granted=False,
            )

        wallet.points_balance += points
        wallet.save(update_fields=("points_balance", "updated_at"))

        ledger_entry = PointsLedger.objects.create(
            user=user,
            change_type=PointsLedger.ChangeType.SIGNUP_BONUS,
            points_delta=points,
            balance_after=wallet.points_balance,
            reason=ledger_reason,
        )

    return SignupBonusGrantResult(
        wallet=wallet,
        grant=grant,
        ledger_entry=ledger_entry,
        points_granted=points,
        balance_after=ledger_entry.balance_after,
        granted=True,
    )


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,
    client_device=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,
            client_device=client_device,
            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,
    )