"""
Defines the economic primitives for resource accounting and budgeting.

This module models the consumption of resources as conserved physical 
quantities—Matter (tokens) and Time (latency).

It relies on `primitives.CanonicalRecord` for immutability but focuses 
strictly on:

Resource Accounting & Control
Standardized structures (`UsageStats`, `BudgetConfig`) that:
1. Model execution costs deterministically.
2. Validate mathematical consistency (e.g., input + output = total).
3. Enforce limits via hard caps or soft notifications.

This separation ensures that financial and operational limits are treated
as first-class domain concepts, decoupled from the static data types.
"""

from typing import Optional
from enum import Enum
from pydantic import Field, model_validator
import warnings

from .primitives import CanonicalRecord, FinitePositiveFloat


# ═══════════════════════════════════════════════════════════════════════════
# TOKEN USAGE TRACKING
# ═══════════════════════════════════════════════════════════════════════════

class UsageStats(CanonicalRecord):
    """Represents the tangible consumption of resources (Matter & Time) for an operation.

    This class tracks the two fundamental economic units of the system:
    1. **Information (Tokens):** The raw material processed by the model.
    2. **Duration (Latency):** The execution time consumed by the system.

    It enforces mathematical consistency for token counts and serves as the 
    standard currency for budgeting, billing, and observability.

    Cache Semantics:
        - 0 means "no cache used" or "provider doesn't support caching".
        - For observability, 0 is preferred over None to simplify aggregation.
    
    Aggregation Behavior:
        - When adding two instances (`stats1 + stats2`), token counts are summed.
        - **Latency is also summed**, assuming sequential execution. Parallel 
        orchestration logic should handle time aggregation separately if needed.

    Attributes:
        input_tokens: Number of tokens in the input prompt.
        output_tokens: Number of tokens in the generated response.
        total_tokens: Total tokens processed (must equal input + output).
        cache_read_input_tokens: Input tokens served from cache.
        cache_creation_input_tokens: Input tokens used to populate cache.
        latency_ms: Wall-clock execution time in milliseconds.

    Raises:
        ValueError: If input_tokens + output_tokens != total_tokens.
        ValueError: If cache_read_input_tokens > input_tokens.
        ValueError: If latency_ms is NaN or Infinity.
    
    Example:
        >>> stats = UsageStats(
        ...     input_tokens=100,
        ...     output_tokens=50,
        ...     total_tokens=150,
        ...     latency_ms=1200.0
        ... )
        >>> stats.total_tokens
        150
    """
    input_tokens: int = Field(
        ge=0,
        description="Number of tokens in the prompt."
    )
    
    output_tokens: int = Field(
        ge=0,
        description="Number of tokens in the generated response."
    )
    
    total_tokens: int = Field(
        ge=0,
        description="Total tokens processed (input + output)."
    )
    
    cache_read_input_tokens: int = Field(
        default=0,
        ge=0,
        description=(
            "Number of input tokens served from cache. "
            "0 if no cache hit or provider doesn't support caching."
        )
    )
    
    cache_creation_input_tokens: int = Field(
        default=0,
        ge=0,
        description=(
            "Number of input tokens used to create a new cache entry. "
            "0 if no cache created or provider doesn't support caching."
        )
    )

    latency_ms: FinitePositiveFloat = Field(
        description="Wall-clock execution time in milliseconds."
    )

    @classmethod
    def zero(cls) -> "UsageStats":
        """Initializes an empty usage record.

        Returns:
            UsageStats: A new instance with all counters set to zero.
        """
        return cls(
            input_tokens=0,
            output_tokens=0,
            total_tokens=0,
            cache_read_input_tokens=0,
            cache_creation_input_tokens=0,
            latency_ms=0.0
        )

    def __add__(self, other: "UsageStats") -> "UsageStats":
        """Aggregates two UsageStats objects.

        Allows for the summation of usage statistics across multiple calls
        or usage sessions.

        Args:
            other: The UsageStats object to add to the current instance.

        Returns:
            UsageStats: A new instance representing the sum of both records.

        Raises:
            TypeError: If the 'other' object is not an instance of UsageStats.
        """
        if not isinstance(other, UsageStats):
            return NotImplemented

        return UsageStats(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            total_tokens=self.total_tokens + other.total_tokens,
            cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens,
            cache_creation_input_tokens=self.cache_creation_input_tokens + other.cache_creation_input_tokens,
            latency_ms=self.latency_ms + other.latency_ms
        )

    @model_validator(mode='after')
    def validate_token_math(self) -> 'UsageStats':
        """Ensures the token counts are mathematically consistent.
        
        This validation catches provider bugs where reported token counts do not 
        aggregate correctly, which has been observed in production with various 
        LLM providers.

        Returns:
            UsageStats: The validated instance.

        Raises:
            ValueError: If input + output != total.
        """
        calculated = self.input_tokens + self.output_tokens
        if self.total_tokens != calculated:
            raise ValueError(
                f"Token math mismatch: "
                f"input({self.input_tokens}) + output({self.output_tokens}) "
                f"= {calculated}, but total_tokens={self.total_tokens}"
            )
        return self
    
    @model_validator(mode='after')
    def validate_cache_consistency(self) -> 'UsageStats':
        """Ensures cache read tokens do not exceed total input tokens.
        
        It is logically impossible to read more tokens from the cache than 
        were present in the input prompt.

        Returns:
            UsageStats: The validated instance.

        Raises:
            ValueError: If cache_read_input_tokens > input_tokens.
        """
        if self.cache_read_input_tokens > self.input_tokens:
            raise ValueError(
                f"Cache read tokens ({self.cache_read_input_tokens}) cannot exceed "
                f"total input tokens ({self.input_tokens})"
            )
        return self
    
    @model_validator(mode='after')
    def validate_latency_plausibility(self) -> 'UsageStats':
        """Validates physical plausibility of resource consumption.
        
        Enforces conservation of time: processing matter (tokens) requires time.
        
        Exception: 100% cache hits (cache_read_input_tokens == input_tokens) with
        zero output tokens may legitimately have near-zero latency.
        
        Note: Finite checks (NaN/Inf) are handled by the field type definition.

        Returns:
            UsageStats: The validated instance.
        """
        # Skip plausibility check if no processing occurred
        if self.total_tokens == 0:
            return self
        
        # Check if this is a full cache hit with no generation
        is_full_cache_hit = (
            self.cache_read_input_tokens == self.input_tokens 
            and self.output_tokens == 0
        )
        
        # Warn if we processed tokens but recorded no time (unless full cache hit)
        if self.latency_ms == 0 and not is_full_cache_hit:
            warnings.warn(
                f"Physical inconsistency: processed {self.total_tokens} tokens "
                f"with 0ms latency. This violates conservation of time unless "
                f"the operation was 100% cached.",
                UserWarning
            )
        
        return self
    
    @property
    def is_empty(self) -> bool:
        """Checks if this usage record represents zero consumption.
        
        Useful for logging systems to filter out noise from no-op operations.
        
        Returns:
            bool: True if both tokens and latency are zero.
        
        Example:
            >>> stats = UsageStats.zero()
            >>> stats.is_empty
            True
        """
        return self.total_tokens == 0 and self.latency_ms == 0
    
    @property
    def cache_efficiency(self) -> float:
        """Calculates the percentage of input tokens served from cache.
        
        Returns:
            float: A value between 0.0 and 1.0 representing the cache hit ratio.
                Returns 0.0 if no input tokens exist or cache was not used.
        
        Example:
            >>> stats = UsageStats(input_tokens=100, output_tokens=50, 
            ...                    total_tokens=150, cache_read_input_tokens=30)
            >>> stats.cache_efficiency
            0.3
        """
        if self.input_tokens == 0:
            return 0.0
        return self.cache_read_input_tokens / self.input_tokens
    
    @property
    def total_cache_tokens(self) -> int:
        """Calculates total cache-related tokens (both read and creation).
        
        This metric is useful for cost calculations when providers charge 
        differently for cache write vs. cache read operations.
        
        Returns:
            int: The sum of cache_read_input_tokens and cache_creation_input_tokens.
        """
        return self.cache_read_input_tokens + self.cache_creation_input_tokens


# ═══════════════════════════════════════════════════════════════════════════
# ECONOMIC PRIMITIVES (BUDGETING)
# ═══════════════════════════════════════════════════════════════════════════

class BudgetExceededError(RuntimeError):
    """Raised when a strict execution limit (HARD_CAP) is breached.
    
    This exception is part of the contract defined by BudgetStrategy.HARD_CAP.
    It indicates that the operation was halted to preserve resources.
    """
    def __init__(self, message: str, current_usage: float, limit: float):
        self.current_usage = current_usage
        self.limit = limit
        super().__init__(f"{message} (Usage: {current_usage} > Limit: {limit})")


class BudgetStrategy(str, Enum):
    """Defines the behavior policy when a resource limit is reached.
    
    This enumeration dictates whether the system should forcefully abort operations
    or simply flag them for review when a budget threshold is crossed.

    Attributes:
        HARD_CAP: Immediately terminates the execution path (raises BudgetExceededError).
                Typically used for autonomous agents (scrapers, workers) to prevent loops.
        SOFT_NOTIFY: Allows execution to continue but flags the run (e.g., logs a warning).
                    Typically used for critical paths or VIP requests where completion is priority.
    """
    HARD_CAP = "hard_cap"
    SOFT_NOTIFY = "soft_notify"


class BudgetConfig(CanonicalRecord):
    """Defines the economic constraints (Resource Limits) for an execution run.

    This class represents the "A Priori" contract (Input Constraints) enforced by 
    the system, as opposed to `UsageStats` which represents the "A Posteriori" 
    report (Output Costs).

    It decouples policy from identity, allowing dynamic budget assignment per request. 
    Limits can be set on **Matter** (Tokens) and/or **Time** (Latency).

    Attributes:
        token_limit: Maximum allowed total tokens (input + output). If None, 
                    unbounded (provider limits still apply).
        time_limit_ms: Maximum allowed execution duration in milliseconds. If None,
                    unbounded (though server timeouts may apply).
        strategy: Enforcement policy (Strict vs. Passive) triggered when EITHER limit is hit.

    Example:
        >>> budget = BudgetConfig(
        ...     token_limit=1000,
        ...     time_limit_ms=5000,
        ...     strategy=BudgetStrategy.HARD_CAP
        ... )
    """
    token_limit: Optional[int] = Field(
        default=None,
        gt=0,
        description="Hard limit on total tokens. If None, no limit is enforced."
    )

    time_limit_ms: Optional[FinitePositiveFloat] = Field(
        default=None,
        description="Hard limit on total execution time (latency). If None, no limit is enforced."
    )

    strategy: BudgetStrategy = Field(
        default=BudgetStrategy.HARD_CAP,
        description="Policy to apply when the limit is reached."
    )

    @model_validator(mode='after')
    def validate_strategy_semantics(self) -> 'BudgetConfig':
        """Prevents "toothless watchdog" configurations.
        
        A HARD_CAP strategy without any defined limits is logically meaningless
        and likely indicates a configuration error. This validator enforces that
        at least one resource constraint exists when hard enforcement is requested.
        
        For SOFT_NOTIFY without limits, emits a warning since this creates a
        passive observer that tracks but never enforces any constraints.

        Returns:
            BudgetConfig: The validated instance.

        Raises:
            ValueError: If strategy is HARD_CAP but both limits are None.
        
        Warns:
            UserWarning: If strategy is SOFT_NOTIFY but both limits are None.
        """
        if self.strategy == BudgetStrategy.HARD_CAP:
            if self.token_limit is None and self.time_limit_ms is None:
                raise ValueError(
                    "BudgetStrategy.HARD_CAP requires at least one limit "
                    "(token_limit or time_limit_ms) to be defined. "
                    "An unbounded hard cap is semantically meaningless."
                )
        
        if self.strategy == BudgetStrategy.SOFT_NOTIFY:
            if self.is_unbounded:
                warnings.warn(
                    "BudgetConfig with SOFT_NOTIFY and no limits acts as a passive "
                    "observer. It will track resource consumption but never enforce "
                    "constraints. If this is intentional (e.g., for analytics), you "
                    "can safely ignore this warning.",
                    UserWarning,
                    stacklevel=2
                )
        
        return self

    @model_validator(mode='after')
    def validate_positive_time_limit(self) -> 'BudgetConfig':
        """Ensures time limit is greater than zero when specified.
        
        A zero time limit would be logically impossible to satisfy, as any operation
        requires non-zero time to execute.
        
        Returns:
            BudgetConfig: The validated instance.
        
        Raises:
            ValueError: If time_limit_ms is exactly 0.
        """
        if self.time_limit_ms is not None and self.time_limit_ms == 0:
            raise ValueError("time_limit_ms must be greater than 0")
        return self
    
    @property
    def is_unbounded(self) -> bool:
        """Checks if this budget imposes no resource constraints.
        
        Useful for the Kernel to optimize away budget checking logic when
        no limits are actually enforced.
        
        Returns:
            bool: True if both token_limit and time_limit_ms are None.
        
        Example:
            >>> budget = BudgetConfig(strategy=BudgetStrategy.SOFT_NOTIFY)
            >>> budget.is_unbounded
            True
        """
        return self.token_limit is None and self.time_limit_ms is None