"""Reserve / settle / release millicredits on the runtime AI ledger only."""
from __future__ import annotations

from dataclasses import dataclass
from uuid import UUID

from src.domain.runtime_ai.repository import RuntimeAiLedgerRepository

RESERVE_SOURCE = "runtime_ai_reservation"
SETTLE_SOURCE = "runtime_ai_settlement"


@dataclass(frozen=True)
class RuntimeReservation:
    reference: str
    millicredits: int
    already_reserved: bool = False


@dataclass(frozen=True)
class RuntimeSettlement:
    reference: str
    reserved: int
    charged: int
    released: int
    already_settled: bool = False


class RuntimeAiWallet:
    def __init__(self, ledger: RuntimeAiLedgerRepository) -> None:
        self._ledger = ledger

    async def balance_millicredits(self, workspace_id: UUID) -> int:
        return await self._ledger.sum_balance(workspace_id)

    async def credit_purchase(
        self,
        *,
        workspace_id: UUID,
        millicredits: int,
        purchase_id: UUID,
        source_id: str,
        description: str,
    ) -> None:
        await self._ledger.acquire_lock(workspace_id)
        existing = await self._ledger.list_by_source(workspace_id, "stripe_checkout", source_id)
        if existing:
            return
        await self._ledger.add_entry(
            workspace_id=workspace_id,
            entry_type="purchase",
            millicredits_delta=millicredits,
            source_type="stripe_checkout",
            source_id=source_id,
            purchase_id=purchase_id,
            description=description,
        )

    async def reserve(
        self,
        *,
        workspace_id: UUID,
        millicredits: int,
        reference: str,
        usage_product: str,
        project_id: UUID | None = None,
    ) -> RuntimeReservation | None:
        if millicredits <= 0:
            raise ValueError("Reservation millicredits must be positive")
        await self._ledger.acquire_lock(workspace_id)
        existing = await self._ledger.list_by_source(workspace_id, RESERVE_SOURCE, reference)
        if existing:
            reserved = sum(-e.millicredits_delta for e in existing if e.millicredits_delta < 0)
            return RuntimeReservation(reference, reserved, already_reserved=True)
        available = await self._ledger.sum_balance(workspace_id)
        if available < millicredits:
            return None
        await self._ledger.add_entry(
            workspace_id=workspace_id,
            entry_type="reservation",
            millicredits_delta=-millicredits,
            usage_product=usage_product,
            source_type=RESERVE_SOURCE,
            source_id=reference,
            project_id=project_id,
            description=f"Reserved runtime AI credits for {usage_product}",
            metadata={"reserved_millicredits": millicredits},
        )
        return RuntimeReservation(reference, millicredits)

    async def settle(
        self,
        *,
        workspace_id: UUID,
        reference: str,
        charged_millicredits: int,
        usage_product: str,
    ) -> RuntimeSettlement:
        if charged_millicredits < 0:
            raise ValueError("Charged millicredits cannot be negative")
        await self._ledger.acquire_lock(workspace_id)
        settlements = await self._ledger.list_by_source(workspace_id, SETTLE_SOURCE, reference)
        if settlements:
            marker = next((e for e in settlements if e.entry_type == "adjustment"), settlements[0])
            meta = marker.metadata_json or {}
            return RuntimeSettlement(
                reference,
                int(meta.get("reserved_millicredits", 0)),
                int(meta.get("charged_millicredits", 0)),
                int(meta.get("released_millicredits", 0)),
                already_settled=True,
            )
        reservations = await self._ledger.list_by_source(workspace_id, RESERVE_SOURCE, reference)
        reserved = sum(-e.millicredits_delta for e in reservations if e.millicredits_delta < 0)
        if not reservations:
            raise ValueError("Runtime AI reservation was not found")
        charged = min(charged_millicredits, reserved) if reserved else 0
        released = max(0, reserved - charged)
        if released:
            await self._ledger.add_entry(
                workspace_id=workspace_id,
                entry_type="refund",
                millicredits_delta=released,
                usage_product=usage_product,
                source_type=SETTLE_SOURCE,
                source_id=reference,
                description="Released unused runtime AI credits",
            )
        extra = charged_millicredits - reserved
        if extra > 0:
            available = await self._ledger.sum_balance(workspace_id)
            take = min(extra, available)
            if take:
                await self._ledger.add_entry(
                    workspace_id=workspace_id,
                    entry_type="consumption",
                    millicredits_delta=-take,
                    usage_product=usage_product,
                    source_type=SETTLE_SOURCE,
                    source_id=reference,
                    description="Additional runtime AI usage beyond reservation",
                )
                charged = reserved + take
        await self._ledger.add_entry(
            workspace_id=workspace_id,
            entry_type="adjustment",
            millicredits_delta=0,
            usage_product=usage_product,
            source_type=SETTLE_SOURCE,
            source_id=reference,
            description="Settled runtime AI reservation",
            metadata={
                "reserved_millicredits": reserved,
                "charged_millicredits": charged,
                "released_millicredits": released,
            },
            allow_zero=True,
        )
        return RuntimeSettlement(reference, reserved, charged, released)

    async def release(self, *, workspace_id: UUID, reference: str, usage_product: str) -> RuntimeSettlement:
        return await self.settle(
            workspace_id=workspace_id,
            reference=reference,
            charged_millicredits=0,
            usage_product=usage_product,
        )