"""Cost budget service for per-session and per-org budget enforcement.

Provides a centralised service that manages budget state for sessions
and organisations, enforces the three-tier hierarchy (plan → session →
org), and emits :class:`~cleveragents.infrastructure.events.types.EventType`
events when budgets approach or exceed their limits.

Thread safety: all public methods acquire the internal ``_lock`` so the
service can safely be shared across threads as a DI Singleton.

Based on Forgejo issue #584.
"""

from __future__ import annotations

import logging
import threading
from typing import Protocol

from cleveragents.domain.models.core.cost_budget import (
    BudgetCheckResult,
    BudgetLevel,
    OrgCostAccumulator,
    SessionCostBudget,
    ThreadSafeOrgCostAccumulator,
)
from cleveragents.infrastructure.events.models import DomainEvent
from cleveragents.infrastructure.events.types import EventType

logger = logging.getLogger(__name__)


class _EventBusProtocol(Protocol):
    """Minimal protocol that the event bus must satisfy."""

    def emit(self, event: DomainEvent) -> None: ...


class CostBudgetService:
    """Manages per-session and per-org cost budgets.

    Args:
        event_bus: Optional event bus for BUDGET_WARNING / BUDGET_EXCEEDED.
        warning_threshold: Utilisation ratio (0.0-1.0) at which a warning
            event is emitted.  Defaults to ``0.8`` (80 %).

    Raises:
        TypeError: If *warning_threshold* is not numeric.
        ValueError: If *warning_threshold* is outside ``[0.0, 1.0]``.
    """

    def __init__(
        self,
        event_bus: _EventBusProtocol | None = None,
        warning_threshold: float = 0.8,
    ) -> None:
        if not isinstance(warning_threshold, (int, float)):
            raise TypeError("warning_threshold must be a number")
        if warning_threshold < 0.0 or warning_threshold > 1.0:
            raise ValueError("warning_threshold must be between 0.0 and 1.0")

        self._event_bus = event_bus
        self._warning_threshold = float(warning_threshold)

        self._sessions: dict[str, SessionCostBudget] = {}
        self._orgs: dict[str, ThreadSafeOrgCostAccumulator] = {}
        self._session_org: dict[str, str] = {}
        self._warned_sessions: set[str] = set()
        self._lock = threading.RLock()

    # ------------------------------------------------------------------ #
    # Session budget management
    # ------------------------------------------------------------------ #

    def configure_session_budget(
        self,
        session_id: str,
        *,
        max_cost_usd: float | None = None,
        org_id: str | None = None,
    ) -> None:
        """Create or update the budget for a session.

        Args:
            session_id: Unique session identifier (must not be empty).
            max_cost_usd: Maximum cost in USD (``None`` = unlimited).
            org_id: Optional org to associate this session with.

        Raises:
            ValueError: If *session_id* is empty, *max_cost_usd* is negative,
                or *org_id* is an empty string.
        """
        if not session_id:
            raise ValueError("session_id must not be empty")
        if max_cost_usd is not None and max_cost_usd < 0:
            raise ValueError("max_cost_usd must be non-negative")
        if org_id is not None and not org_id:
            raise ValueError("org_id must not be empty when provided")

        with self._lock:
            existing = self._sessions.get(session_id)
            total = existing.total_cost if existing else 0.0
            self._sessions[session_id] = SessionCostBudget(
                max_cost_usd=max_cost_usd,
                total_cost=total,
            )
            if org_id is not None:
                self._session_org[session_id] = org_id

    def get_session_budget(self, session_id: str) -> SessionCostBudget | None:
        """Return the budget for *session_id*, or ``None``.

        Raises:
            ValueError: If *session_id* is empty.
        """
        if not session_id:
            raise ValueError("session_id must not be empty")
        with self._lock:
            return self._sessions.get(session_id)

    def remove_session(self, session_id: str) -> None:
        """Remove all budget state for a session."""
        with self._lock:
            self._sessions.pop(session_id, None)
            self._session_org.pop(session_id, None)
            self._warned_sessions.discard(session_id)

    # ------------------------------------------------------------------ #
    # Org budget management
    # ------------------------------------------------------------------ #

    def configure_org_budget(
        self,
        org_id: str,
        *,
        max_cost_usd: float | None = None,
    ) -> None:
        """Create or update the budget for an organisation.

        Args:
            org_id: Organisation identifier (must not be empty).
            max_cost_usd: Maximum cost in USD (``None`` = unlimited).

        Raises:
            ValueError: If *org_id* is empty or *max_cost_usd* is negative.
        """
        if not org_id:
            raise ValueError("org_id must not be empty")
        if max_cost_usd is not None and max_cost_usd < 0:
            raise ValueError("max_cost_usd must be non-negative")

        with self._lock:
            existing = self._orgs.get(org_id)
            if existing is not None:
                existing.max_cost_usd = max_cost_usd
            else:
                inner = OrgCostAccumulator(org_id=org_id, max_cost_usd=max_cost_usd)
                self._orgs[org_id] = ThreadSafeOrgCostAccumulator(inner)

    def get_org_accumulator(self, org_id: str) -> ThreadSafeOrgCostAccumulator | None:
        """Return the accumulator for *org_id*, or ``None``.

        Raises:
            ValueError: If *org_id* is empty.
        """
        if not org_id:
            raise ValueError("org_id must not be empty")
        with self._lock:
            return self._orgs.get(org_id)

    def remove_org(self, org_id: str) -> None:
        """Remove all budget state for an organisation."""
        with self._lock:
            self._orgs.pop(org_id, None)

    # ------------------------------------------------------------------ #
    # Budget hierarchy enforcement
    # ------------------------------------------------------------------ #

    def check_budget_hierarchy(
        self,
        session_id: str,
        plan_cost: float,
        *,
        plan_budget: float | None = None,
        plan_spent: float | None = None,
    ) -> BudgetCheckResult:
        """Check whether *plan_cost* is allowed by the full hierarchy.

        Evaluation order: plan → session → org.  The first tier that
        would be exceeded causes an immediate denial.

        Args:
            session_id: Session identifier.
            plan_cost: Proposed cost for this plan step.
            plan_budget: Optional per-plan budget cap.
            plan_spent: Optional per-plan accumulated spend.

        Returns:
            :class:`BudgetCheckResult` with the enforcement outcome.

        Raises:
            ValueError: On empty *session_id*, negative *plan_cost*, or
                negative *plan_spent*.
            TypeError: If *plan_cost* is not numeric.
        """
        if not isinstance(plan_cost, (int, float)):
            raise TypeError("plan_cost must be a number")
        if not session_id:
            raise ValueError("session_id must not be empty")
        if plan_cost < 0:
            raise ValueError("plan_cost must be non-negative")
        if plan_spent is not None and plan_spent < 0:
            raise ValueError("plan_spent must be non-negative")

        with self._lock:
            # --- Plan-level check ---
            if plan_budget is not None:
                spent = plan_spent if plan_spent is not None else 0.0
                if spent + plan_cost > plan_budget:
                    return BudgetCheckResult(
                        allowed=False,
                        exceeded_level=BudgetLevel.PLAN,
                        reason=(
                            f"Plan budget exceeded: {spent + plan_cost:.2f}"
                            f" > {plan_budget:.2f}"
                        ),
                    )

            # --- Session-level check ---
            session_budget = self._sessions.get(session_id)
            if (
                session_budget is not None
                and session_budget.max_cost_usd is not None
                and session_budget.would_exceed(plan_cost)
            ):
                self._emit_exceeded(session_id)
                return BudgetCheckResult(
                    allowed=False,
                    exceeded_level=BudgetLevel.SESSION,
                    reason=(
                        f"Session budget exceeded: "
                        f"{session_budget.total_cost + plan_cost:.2f}"
                        f" > {session_budget.max_cost_usd:.2f}"
                    ),
                )

            # --- Org-level check ---
            org_id = self._session_org.get(session_id)
            if org_id is not None:
                org_acc = self._orgs.get(org_id)
                if (
                    org_acc is not None
                    and org_acc.max_cost_usd is not None
                    and org_acc.would_exceed(plan_cost)
                ):
                    self._emit_exceeded(session_id)
                    return BudgetCheckResult(
                        allowed=False,
                        exceeded_level=BudgetLevel.ORG,
                        reason=(
                            f"Org budget exceeded: "
                            f"{org_acc.total_cost + plan_cost:.2f}"
                            f" > {org_acc.max_cost_usd:.2f}"
                        ),
                    )

            # --- Warning check ---
            warning = self._check_warning(session_id)

            return BudgetCheckResult(allowed=True, warning=warning)

    def record_plan_cost(self, session_id: str, cost: float) -> None:
        """Record a cost against the session (and its org if linked).

        Args:
            session_id: Session identifier.
            cost: Cost to record.

        Raises:
            ValueError: On empty *session_id* or negative *cost*.
            TypeError: If *cost* is not numeric.
        """
        if not isinstance(cost, (int, float)):
            raise TypeError("cost must be a number")
        if not session_id:
            raise ValueError("session_id must not be empty")
        if cost < 0:
            raise ValueError("cost must be non-negative")

        with self._lock:
            session_budget = self._sessions.get(session_id)
            if session_budget is not None:
                session_budget.record_cost(cost)

            org_id = self._session_org.get(session_id)
            if org_id is not None:
                org_acc = self._orgs.get(org_id)
                if org_acc is not None:
                    org_acc.record_cost(cost)

    # ------------------------------------------------------------------ #
    # Private helpers
    # ------------------------------------------------------------------ #

    def _check_warning(self, session_id: str) -> bool:
        """Return True if the session is above warning threshold."""
        if session_id in self._warned_sessions:
            return False
        session_budget = self._sessions.get(session_id)
        if session_budget is None:
            return False
        util = session_budget.utilization()
        if util is None:
            return False
        if util >= self._warning_threshold:
            self._warned_sessions.add(session_id)
            self._emit_warning(session_id)
            return True
        return False

    def _emit_warning(self, session_id: str) -> None:
        """Emit a BUDGET_WARNING event if an event bus is configured."""
        if self._event_bus is None:
            return
        event = DomainEvent(
            event_type=EventType.BUDGET_WARNING,
            session_id=session_id,
            details={"session_id": session_id, "level": "session"},
        )
        self._event_bus.emit(event)

    def _emit_exceeded(self, session_id: str) -> None:
        """Emit a BUDGET_EXCEEDED event if an event bus is configured."""
        if self._event_bus is None:
            return
        event = DomainEvent(
            event_type=EventType.BUDGET_EXCEEDED,
            session_id=session_id,
            details={"session_id": session_id},
        )
        self._event_bus.emit(event)


__all__ = ["CostBudgetService"]
