"""Ledger-backed credit reservations for work that is billed after execution."""
from __future__ import annotations

from dataclasses import dataclass
from uuid import UUID

from src.domain.billing.balance_service import WorkspaceUsageBalanceService
from src.domain.billing.repository import (
    WorkspaceCreditLedgerRepository,
    WorkspaceUsageAllowanceRepository,
)

RESERVATION_SOURCE_TYPE = "credit_reservation"
SETTLEMENT_SOURCE_TYPE = "credit_settlement"


@dataclass(frozen=True)
class CreditReservation:
    reference: str
    reserved_units: int
    already_reserved: bool = False


@dataclass(frozen=True)
class CreditSettlement:
    reference: str
    reserved_units: int
    charged_units: int
    released_units: int
    already_settled: bool = False


class CreditReservationService:
    """Reserve the estimate first, then release the unused portion after execution.

    The workspace credit ledger remains the financial audit trail.  A reservation is
    represented by normal negative consumption entries, which immediately removes
    those credits from the available balance and prevents concurrent over-spend.
    """

    def __init__(
        self,
        *,
        balance_service: WorkspaceUsageBalanceService,
        ledger_repo: WorkspaceCreditLedgerRepository,
        allowance_repo: WorkspaceUsageAllowanceRepository,
    ) -> None:
        self._balance = balance_service
        self._ledger = ledger_repo
        self._allowance = allowance_repo

    async def reserve(
        self,
        *,
        workspace_id: UUID,
        units: int,
        reference: str,
        project_id: UUID | None = None,
        codegen_job_id: UUID | None = None,
        description: str | None = None,
    ) -> CreditReservation | None:
        """Reserve units once.  ``None`` means the workspace cannot afford them.

        The caller must commit its enclosing database transaction after a successful
        reservation.  The advisory lock in the balance service serializes competing
        requests for the same workspace.
        """
        if units <= 0:
            raise ValueError("Reservation units must be positive")

        await self._ledger.acquire_consumption_lock(workspace_id)
        existing = await self._ledger.list_by_source(
            workspace_id, source_type=RESERVATION_SOURCE_TYPE, source_id=reference
        )
        if existing:
            reserved_units = sum(
                -entry.units_delta for entry in existing if entry.units_delta < 0
            )
            return CreditReservation(
                reference=reference,
                reserved_units=reserved_units,
                already_reserved=True,
            )

        consumed = await self._balance.consume_units(
            workspace_id,
            units,
            source_type=RESERVATION_SOURCE_TYPE,
            source_id=reference,
            project_id=project_id,
            codegen_job_id=codegen_job_id,
            description=description or "Reserved credits for AI work",
            metadata={"reservation_reference": reference, "reserved_units": units},
        )
        if not consumed:
            return None
        return CreditReservation(reference=reference, reserved_units=units)

    async def settle(
        self,
        *,
        workspace_id: UUID,
        reference: str,
        charged_units: int,
        description: str | None = None,
    ) -> CreditSettlement:
        """Finalize a reservation and refund the unused credits exactly once."""
        if charged_units < 0:
            raise ValueError("Charged units cannot be negative")

        await self._ledger.acquire_consumption_lock(workspace_id)
        settlements = await self._ledger.list_by_source(
            workspace_id, source_type=SETTLEMENT_SOURCE_TYPE, source_id=reference
        )
        if settlements:
            marker = settlements[0]
            metadata = marker.metadata_json or {}
            return CreditSettlement(
                reference=reference,
                reserved_units=int(metadata.get("reserved_units", 0)),
                charged_units=int(metadata.get("charged_units", 0)),
                released_units=int(metadata.get("released_units", 0)),
                already_settled=True,
            )

        reservations = await self._ledger.list_by_source(
            workspace_id, source_type=RESERVATION_SOURCE_TYPE, source_id=reference
        )
        reserved = sum(-entry.units_delta for entry in reservations if entry.units_delta < 0)
        if not reservations:
            raise ValueError("Credit reservation was not found")
        if charged_units > reserved:
            raise ValueError("Charged units exceed the reserved credit cap")

        remaining_release = reserved - charged_units
        released = 0
        # Reverse the latest reservation portions first, leaving the original
        # allowance-first consumption order intact for the final charged amount.
        for entry in reversed(reservations):
            if remaining_release <= 0 or entry.units_delta >= 0:
                continue
            portion = min(remaining_release, -entry.units_delta)
            if entry.source_bucket == "subscription_allowance":
                allowance_id = (entry.metadata_json or {}).get("allowance_id")
                if allowance_id:
                    restored = await self._allowance.restore_units(UUID(allowance_id), portion)
                    if restored is None:
                        raise ValueError("Reservation allowance was not found")
                else:
                    raise ValueError("Reservation allowance reference is missing")

            await self._ledger.add_entry(
                workspace_id=workspace_id,
                entry_type="refund",
                units_delta=portion,
                source_bucket=entry.source_bucket,
                source_type=SETTLEMENT_SOURCE_TYPE,
                source_id=reference,
                subscription_id=entry.subscription_id,
                project_id=entry.project_id,
                codegen_job_id=entry.codegen_job_id,
                description=description or "Released unused AI credits",
                metadata={
                    "reservation_reference": reference,
                    "released_from_ledger_id": str(entry.id),
                },
            )
            remaining_release -= portion
            released += portion

        await self._ledger.add_entry(
            workspace_id=workspace_id,
            entry_type="adjustment",
            units_delta=0,
            source_type=SETTLEMENT_SOURCE_TYPE,
            source_id=reference,
            description=description or "Settled AI credit reservation",
            metadata={
                "reservation_reference": reference,
                "reserved_units": reserved,
                "charged_units": charged_units,
                "released_units": released,
            },
            allow_zero=True,
        )
        return CreditSettlement(reference, reserved, charged_units, released)