"""Per-session and per-org cost budget domain models.

Implements a three-tier budget hierarchy:

    per-plan → per-session → per-org

Each tier independently tracks accumulated costs and enforces a
configurable maximum.  The **tightest** limit at any level wins:
if a tool invocation would exceed *any* tier's budget, the request
is denied.

Thread safety
-------------
- :class:`SessionCostBudget` is **not** thread-safe; callers must
  synchronise externally (the budget service holds a lock).
- :class:`OrgCostAccumulator` is **not** thread-safe.
- :class:`ThreadSafeOrgCostAccumulator` wraps an ``OrgCostAccumulator``
  with a :class:`threading.Lock` for concurrent access.

Based on Forgejo issue #584.
"""

from __future__ import annotations

import threading
from enum import StrEnum

from pydantic import BaseModel, ConfigDict, Field


class BudgetLevel(StrEnum):
    """Tier in the cost budget hierarchy."""

    PLAN = "plan"
    SESSION = "session"
    ORG = "org"


class BudgetCheckResult(BaseModel):
    """Outcome of a budget hierarchy check.

    Attributes:
        allowed: ``True`` when the proposed cost fits within all tiers.
        exceeded_level: The budget tier that denied the request, or
            ``None`` if allowed.
        reason: Human-readable denial reason (empty when allowed).
        warning: ``True`` when utilisation is above the warning
            threshold but the request is still allowed.
    """

    model_config = ConfigDict(frozen=True)

    allowed: bool
    exceeded_level: BudgetLevel | None = None
    reason: str = ""
    warning: bool = False


class SessionCostBudget(BaseModel):
    """Tracks accumulated cost for a single session.

    A session may span multiple plans.  Each plan's cost is recorded
    against the session total so that an operator-defined session cap
    is enforced even when individual plan budgets are generous.

    Attributes:
        max_cost_usd: Upper spending limit for this session.
            ``None`` means unlimited.
        total_cost: Running total of all recorded costs.
    """

    model_config = ConfigDict(validate_assignment=True)

    max_cost_usd: float | None = Field(
        default=None,
        ge=0.0,
        description="Maximum allowed cost in USD (None = unlimited).",
    )
    total_cost: float = Field(
        default=0.0,
        ge=0.0,
        description="Accumulated cost in USD.",
    )

    # ------------------------------------------------------------------ #
    # Queries
    # ------------------------------------------------------------------ #

    def utilization(self) -> float | None:
        """Fraction of the budget consumed (0.0-1.0+).

        Returns ``None`` when no cap is configured.
        """
        if self.max_cost_usd is None:
            return None
        if self.max_cost_usd == 0.0:
            return None
        return self.total_cost / self.max_cost_usd

    def remaining(self) -> float | None:
        """Remaining budget in USD.

        Returns ``None`` when no cap is configured.
        """
        if self.max_cost_usd is None:
            return None
        return max(0.0, self.max_cost_usd - self.total_cost)

    def is_exceeded(self) -> bool:
        """Return ``True`` if accumulated cost exceeds the cap."""
        if self.max_cost_usd is None:
            return False
        return self.total_cost > self.max_cost_usd

    def would_exceed(self, cost: float) -> bool:
        """Return ``True`` if adding *cost* would exceed the budget.

        Raises:
            TypeError: If *cost* is not numeric.
            ValueError: If *cost* is negative.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if cost < 0:
            raise ValueError("cost must be non-negative")
        if self.max_cost_usd is None:
            return False
        return self.total_cost + cost > self.max_cost_usd

    def can_afford(self, cost: float) -> bool:
        """Return ``True`` if *cost* fits within the remaining budget."""
        if self.max_cost_usd is None:
            return True
        return self.total_cost + cost <= self.max_cost_usd

    # ------------------------------------------------------------------ #
    # Mutations
    # ------------------------------------------------------------------ #

    def record_cost(self, cost: float) -> None:
        """Add *cost* to the running total.

        Raises:
            TypeError: If *cost* is not numeric.
            ValueError: If *cost* is negative.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if cost < 0:
            raise ValueError("cost must be non-negative")
        self.total_cost += cost

    def reset(self) -> None:
        """Reset the accumulated cost to zero."""
        self.total_cost = 0.0


class OrgCostAccumulator(BaseModel):
    """Tracks accumulated cost across an entire organisation.

    Multiple sessions (and their plans) contribute to a single org
    accumulator so that organisation-wide spending limits can be
    enforced.

    Attributes:
        org_id: Organisation identifier.
        max_cost_usd: Organisation spending cap.  ``None`` = unlimited.
        total_cost: Running total across all sessions/plans.
    """

    model_config = ConfigDict(validate_assignment=True)

    org_id: str = Field(
        ...,
        min_length=1,
        description="Organisation identifier.",
    )
    max_cost_usd: float | None = Field(
        default=None,
        ge=0.0,
        description="Maximum allowed cost in USD (None = unlimited).",
    )
    total_cost: float = Field(
        default=0.0,
        ge=0.0,
        description="Accumulated cost in USD.",
    )

    # ------------------------------------------------------------------ #
    # Queries
    # ------------------------------------------------------------------ #

    def utilization(self) -> float | None:
        """Fraction of the org budget consumed (0.0-1.0+).

        Returns ``None`` when no cap is configured.
        """
        if self.max_cost_usd is None:
            return None
        if self.max_cost_usd == 0.0:
            return None
        return self.total_cost / self.max_cost_usd

    def remaining(self) -> float | None:
        """Remaining org budget in USD.

        Returns ``None`` when no cap is configured.
        """
        if self.max_cost_usd is None:
            return None
        return max(0.0, self.max_cost_usd - self.total_cost)

    def is_exceeded(self) -> bool:
        """Return ``True`` if accumulated cost exceeds the cap."""
        if self.max_cost_usd is None:
            return False
        return self.total_cost > self.max_cost_usd

    def would_exceed(self, cost: float) -> bool:
        """Return ``True`` if adding *cost* would exceed the budget.

        Raises:
            TypeError: If *cost* is not numeric.
            ValueError: If *cost* is negative.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if cost < 0:
            raise ValueError("cost must be non-negative")
        if self.max_cost_usd is None:
            return False
        return self.total_cost + cost > self.max_cost_usd

    def can_afford(self, cost: float) -> bool:
        """Return ``True`` if *cost* fits within the remaining budget."""
        if self.max_cost_usd is None:
            return True
        return self.total_cost + cost <= self.max_cost_usd

    # ------------------------------------------------------------------ #
    # Mutations
    # ------------------------------------------------------------------ #

    def record_cost(self, cost: float) -> None:
        """Add *cost* to the running total.

        Raises:
            TypeError: If *cost* is not numeric.
            ValueError: If *cost* is negative.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if cost < 0:
            raise ValueError("cost must be non-negative")
        self.total_cost += cost

    def reset(self) -> None:
        """Reset the accumulated cost to zero."""
        self.total_cost = 0.0


class ThreadSafeOrgCostAccumulator:
    """Thread-safe wrapper around :class:`OrgCostAccumulator`.

    All attribute access and mutations are serialised through a
    :class:`threading.Lock`.
    """

    def __init__(
        self,
        inner: OrgCostAccumulator,
    ) -> None:
        if not isinstance(inner, OrgCostAccumulator):
            raise TypeError(
                f"inner must be an OrgCostAccumulator, got {type(inner).__name__}"
            )
        self._inner = inner
        self._lock = threading.Lock()

    # ------------------------------------------------------------------ #
    # Queries (delegated)
    # ------------------------------------------------------------------ #

    @property
    def max_cost_usd(self) -> float | None:
        with self._lock:
            return self._inner.max_cost_usd

    @max_cost_usd.setter
    def max_cost_usd(self, value: float | None) -> None:
        with self._lock:
            self._inner.max_cost_usd = value

    @property
    def total_cost(self) -> float:
        with self._lock:
            return self._inner.total_cost

    @property
    def org_id(self) -> str:
        with self._lock:
            return self._inner.org_id

    def utilization(self) -> float | None:
        with self._lock:
            return self._inner.utilization()

    def remaining(self) -> float | None:
        with self._lock:
            return self._inner.remaining()

    def is_exceeded(self) -> bool:
        with self._lock:
            return self._inner.is_exceeded()

    def would_exceed(self, cost: float) -> bool:
        with self._lock:
            return self._inner.would_exceed(cost)

    def can_afford(self, cost: float) -> bool:
        with self._lock:
            return self._inner.can_afford(cost)

    # ------------------------------------------------------------------ #
    # Mutations (delegated)
    # ------------------------------------------------------------------ #

    def record_cost(self, cost: float) -> None:
        with self._lock:
            self._inner.record_cost(cost)

    def reset(self) -> None:
        with self._lock:
            self._inner.reset()

    def snapshot(self) -> OrgCostAccumulator:
        """Return a snapshot copy of the inner accumulator."""
        with self._lock:
            return self._inner.model_copy(deep=True)


__all__ = [
    "BudgetCheckResult",
    "BudgetLevel",
    "OrgCostAccumulator",
    "SessionCostBudget",
    "ThreadSafeOrgCostAccumulator",
]
