"""Cost tracking and budget enforcement service.

Implements cost tracking, budget enforcement, and provider fallback controls
for AI provider usage. Tracks token usage, monetary costs, and enforces
configurable budget limits at plan, session, and organization levels.
"""

from __future__ import annotations

from typing import Any

from cleveragents.core.exceptions import BudgetExceededError
from cleveragents.domain.models.core.cost_budget import (
    BudgetCheckResult,
    BudgetLevel,
    SessionCostBudget,
    ThreadSafeOrgCostAccumulator,
)
from cleveragents.domain.models.core.cost_metadata import (
    BudgetExhaustionEvent,
    CostMetadata,
)

WARNING_THRESHOLD = 0.9


class CostTrackingService:
    """Service for tracking costs and enforcing budget limits.

    Manages cost tracking at plan, session, and organization levels.
    Enforces budget limits and raises BudgetExceededError when limits
    are exceeded.

    Thread safety:
        - Session budgets are protected by an external lock (caller's responsibility)
        - Org accumulators use ThreadSafeOrgCostAccumulator for thread safety
    """

    def __init__(self) -> None:
        """Initialize the cost tracking service."""

    def check_budget(
        self,
        session_budget: SessionCostBudget | None,
        org_accumulator: ThreadSafeOrgCostAccumulator | None,
        proposed_cost: float,
    ) -> BudgetCheckResult:
        """Check if a proposed cost fits within all budget tiers.

        Args:
            session_budget: Session-level budget (None = unlimited).
            org_accumulator: Organization-level accumulator (None = unlimited).
            proposed_cost: Cost to check in USD.

        Returns:
            BudgetCheckResult indicating whether the cost is allowed.

        Raises:
            ValueError: If proposed_cost is negative.
        """
        if proposed_cost < 0:
            raise ValueError("proposed_cost must be non-negative")

        # Check session budget
        if session_budget is not None and session_budget.would_exceed(proposed_cost):
            utilization = session_budget.utilization()
            warning = False
            if (
                not session_budget.is_exceeded()
                and utilization is not None
                and utilization >= WARNING_THRESHOLD
            ):
                warning = True
            return BudgetCheckResult(
                allowed=False,
                exceeded_level=BudgetLevel.SESSION,
                reason=(
                    f"Session budget of ${session_budget.max_cost_usd:.2f} "
                    f"would be exceeded (current: "
                    f"${session_budget.total_cost:.2f}, "
                    f"proposed: ${proposed_cost:.2f})"
                ),
                warning=warning,
            )

        # Check org budget
        if org_accumulator is not None and org_accumulator.would_exceed(proposed_cost):
            utilization = org_accumulator.utilization()
            warning = False
            if (
                not org_accumulator.is_exceeded()
                and utilization is not None
                and utilization >= WARNING_THRESHOLD
            ):
                warning = True
            return BudgetCheckResult(
                allowed=False,
                exceeded_level=BudgetLevel.ORG,
                reason=(
                    f"Organization budget of ${org_accumulator.max_cost_usd:.2f} "
                    f"would be exceeded (current: "
                    f"${org_accumulator.total_cost:.2f}, "
                    f"proposed: ${proposed_cost:.2f})"
                ),
                warning=warning,
            )

        # Check for warning thresholds
        warning = False
        if session_budget is not None:
            utilization = session_budget.utilization()
            if utilization is not None and utilization >= WARNING_THRESHOLD:
                warning = True

        if org_accumulator is not None:
            utilization = org_accumulator.utilization()
            if utilization is not None and utilization >= WARNING_THRESHOLD:
                warning = True

        return BudgetCheckResult(allowed=True, warning=warning)

    def record_cost(
        self,
        cost_metadata: CostMetadata,
        session_budget: SessionCostBudget | None,
        org_accumulator: ThreadSafeOrgCostAccumulator | None,
        *,
        input_tokens: int,
        output_tokens: int,
        cost: float,
        provider: str,
        model: str = "",
    ) -> None:
        """Record cost usage and update budgets.

        Args:
            cost_metadata: Cost metadata to update.
            session_budget: Session budget to update (None = skip).
            org_accumulator: Org accumulator to update (None = skip).
            input_tokens: Input tokens consumed.
            output_tokens: Output tokens consumed.
            cost: Cost in USD.
            provider: Provider name.
            model: Model name (optional).

        Raises:
            ValueError: If any numeric argument is invalid.
        """
        # Record in cost metadata
        cost_metadata.record_usage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            provider=provider,
        )

        # Update session budget
        if session_budget is not None:
            session_budget.record_cost(cost)
            cost_metadata.budget_remaining = session_budget.remaining()

        # Update org accumulator
        if org_accumulator is not None:
            org_accumulator.record_cost(cost)

    def enforce_budget(
        self,
        session_budget: SessionCostBudget | None,
        org_accumulator: ThreadSafeOrgCostAccumulator | None,
        cost_metadata: CostMetadata,
        *,
        provider: str,
        model: str = "",
    ) -> None:
        """Enforce budget limits and raise error if exceeded.

        Args:
            session_budget: Session budget to check.
            org_accumulator: Org accumulator to check.
            cost_metadata: Cost metadata for recording exhaustion events.
            provider: Provider name for error reporting.
            model: Model name for error reporting.

        Raises:
            BudgetExceededError: If any budget limit is exceeded.
        """
        # Check session budget
        if session_budget is not None and session_budget.is_exceeded():
            event = BudgetExhaustionEvent(
                budget_type=BudgetLevel.SESSION.value,
                limit=session_budget.max_cost_usd or 0.0,
                used=session_budget.total_cost,
                provider=provider,
                model=model,
            )
            cost_metadata.budget_exhaustion_events.append(event)
            raise BudgetExceededError(
                budget_type=BudgetLevel.SESSION.value,
                limit=session_budget.max_cost_usd or 0.0,
                used=session_budget.total_cost,
                provider=provider,
                model=model,
            )

        # Check org budget
        if org_accumulator is not None and org_accumulator.is_exceeded():
            event = BudgetExhaustionEvent(
                budget_type=BudgetLevel.ORG.value,
                limit=org_accumulator.max_cost_usd or 0.0,
                used=org_accumulator.total_cost,
                provider=provider,
                model=model,
            )
            cost_metadata.budget_exhaustion_events.append(event)
            raise BudgetExceededError(
                budget_type=BudgetLevel.ORG.value,
                limit=org_accumulator.max_cost_usd or 0.0,
                used=org_accumulator.total_cost,
                provider=provider,
                model=model,
            )

    def get_budget_status(
        self,
        session_budget: SessionCostBudget | None,
        org_accumulator: ThreadSafeOrgCostAccumulator | None,
    ) -> dict[str, Any]:
        """Get current budget status.

        Args:
            session_budget: Session budget to query.
            org_accumulator: Org accumulator to query.

        Returns:
            Dictionary with budget status information.
        """
        status: dict[str, Any] = {}

        if session_budget is not None:
            status["session"] = {
                "limit": session_budget.max_cost_usd,
                "used": session_budget.total_cost,
                "remaining": session_budget.remaining(),
                "utilization": session_budget.utilization(),
            }

        if org_accumulator is not None:
            status["organization"] = {
                "limit": org_accumulator.max_cost_usd,
                "used": org_accumulator.total_cost,
                "remaining": org_accumulator.remaining(),
                "utilization": org_accumulator.utilization(),
            }

        return status


__all__ = [
    "CostTrackingService",
]
