"""Cost metadata model for plan execution tracking.

Tracks token usage, monetary costs, and budget exhaustion events
attached to plan execution lifecycle.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator


class BudgetExhaustionEvent(BaseModel):
    """Record of a budget limit being reached during plan execution.

    Persisted in plan metadata for auditability (issue #324).
    """

    timestamp: datetime = Field(
        default_factory=datetime.now,
        description="When the budget threshold was reached",
    )
    budget_type: str = Field(
        ...,
        min_length=1,
        description="Type of budget exhausted: 'plan' or 'daily'",
    )
    limit: float = Field(
        ...,
        ge=0.0,
        description="The budget limit that was reached (USD)",
    )
    used: float = Field(
        ...,
        ge=0.0,
        description="Amount used at time of exhaustion (USD)",
    )
    provider: str = Field(
        default="",
        description="Provider that triggered the exhaustion",
    )
    model: str = Field(
        default="",
        description="Model that triggered the exhaustion",
    )

    @field_validator("budget_type")
    @classmethod
    def validate_budget_type(cls: type[BudgetExhaustionEvent], v: str) -> str:
        allowed = {"plan", "daily"}
        if v not in allowed:
            raise ValueError(f"budget_type must be one of {allowed}, got {v!r}")
        return v

    model_config = ConfigDict(validate_assignment=True)


class CostMetadata(BaseModel):
    """Cost and token tracking metadata for a plan execution.

    Attached to plan execution to provide cost visibility and
    auditability. Surfaced in ``plan status`` CLI output.
    """

    total_tokens: int = Field(
        default=0,
        ge=0,
        description="Total tokens consumed (input + output)",
    )
    input_tokens: int = Field(
        default=0,
        ge=0,
        description="Input/prompt tokens consumed",
    )
    output_tokens: int = Field(
        default=0,
        ge=0,
        description="Output/completion tokens consumed",
    )
    total_cost: float = Field(
        default=0.0,
        ge=0.0,
        description="Total estimated cost in USD",
    )
    budget_remaining: float | None = Field(
        default=None,
        description="Remaining budget in USD (None if unlimited)",
    )
    provider_costs: dict[str, float] = Field(
        default_factory=dict,
        description="Cost breakdown by provider name",
    )
    budget_exhaustion_events: list[BudgetExhaustionEvent] = Field(
        default_factory=list,
        description="Audit log of budget limit events",
    )

    def record_usage(
        self,
        *,
        input_tokens: int,
        output_tokens: int,
        cost: float,
        provider: str,
    ) -> None:
        """Record token usage and cost from a provider call.

        Args:
            input_tokens: Number of input tokens consumed.
            output_tokens: Number of output tokens consumed.
            cost: Estimated cost in USD.
            provider: Provider name for cost breakdown.

        Raises:
            ValueError: If any numeric argument is negative.
        """
        if input_tokens < 0:
            raise ValueError(f"input_tokens must be >= 0, got {input_tokens}")
        if output_tokens < 0:
            raise ValueError(f"output_tokens must be >= 0, got {output_tokens}")
        if cost < 0.0:
            raise ValueError(f"cost must be >= 0.0, got {cost}")
        if not provider:
            raise ValueError("provider must be a non-empty string")

        self.input_tokens += input_tokens
        self.output_tokens += output_tokens
        self.total_tokens += input_tokens + output_tokens
        self.total_cost += cost
        self.provider_costs[provider] = self.provider_costs.get(provider, 0.0) + cost

    def as_display_dict(self) -> dict[str, Any]:
        """Return a dictionary suitable for CLI display."""
        result: dict[str, Any] = {
            "total_tokens": self.total_tokens,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_cost_usd": round(self.total_cost, 6),
        }
        if self.budget_remaining is not None:
            result["budget_remaining_usd"] = round(self.budget_remaining, 6)
        if self.provider_costs:
            result["provider_costs"] = {
                k: round(v, 6) for k, v in self.provider_costs.items()
            }
        if self.budget_exhaustion_events:
            result["budget_exhaustion_events"] = [
                {
                    "timestamp": evt.timestamp.isoformat(),
                    "budget_type": evt.budget_type,
                    "limit": evt.limit,
                    "used": evt.used,
                    "provider": evt.provider,
                    "model": evt.model,
                }
                for evt in self.budget_exhaustion_events
            ]
        return result

    model_config = ConfigDict(validate_assignment=True)
