"""
Token usage tracking for AI model interactions in the Agentle framework.

This module provides the Usage model for tracking and calculating token consumption
in AI model interactions. It captures both prompt tokens (input) and completion tokens
(output), allowing for precise monitoring of token usage across different providers.

The Usage model supports aggregation through arithmetic operations, making it easy
to calculate cumulative usage across multiple requests or sessions. This is
particularly important for monitoring costs and optimizing applications that
make heavy use of AI models.
"""

from __future__ import annotations

from rsb.decorators.value_objects import valueobject
from rsb.models.base_model import BaseModel
from rsb.models.field import Field


@valueobject
class Usage(BaseModel):
    """
    Tracks and calculates token consumption in AI model interactions.

    This class encapsulates metrics about token usage in AI generation operations,
    tracking both tokens sent to the model (prompt) and tokens generated by the
    model (completion). These metrics are crucial for monitoring resource usage,
    estimating costs, and optimizing interactions with AI models.

    The class provides utility methods for calculating total token usage and
    supports aggregation of usage data across multiple operations through
    arithmetic operations.

    Attributes:
        prompt_tokens: Number of tokens in the prompt sent to the model.
        completion_tokens: Number of tokens generated by the model in its response.

    Methods:
        total_tokens: Calculates the total number of tokens used (prompt + completion).
        __add__: Enables adding two Usage objects together, aggregating their token counts.
        __radd__: Supports adding Usage objects in sum() operations, starting from 0.
        zero: Creates an empty Usage object with zero tokens.
    """

    prompt_tokens: int = Field(
        default=0,
        description="Number of tokens in the prompt sent to the AI model. Represents the input token count that affects context window usage and factors into pricing calculations.",
        ge=0,
        examples=[10, 500, 2000],
    )

    completion_tokens: int = Field(
        default=0,
        description="Number of tokens generated by the AI model in its response. Represents the output token count that factors into pricing calculations and affects generation time.",
        ge=0,
        examples=[5, 250, 1000],
    )

    @property
    def total_tokens(self) -> int:
        """
        Calculate the total number of tokens used in the operation.

        Returns:
            The sum of prompt_tokens and completion_tokens.
        """
        return self.prompt_tokens + self.completion_tokens

    def __add__(self, other: Usage) -> Usage:
        """
        Add two Usage objects together, aggregating their token counts.

        Args:
            other: Another Usage object to add to this one.

        Returns:
            A new Usage object with summed token counts.
        """
        return Usage(
            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
            completion_tokens=self.completion_tokens + other.completion_tokens,
        )

    def __radd__(self, other: int | Usage) -> Usage:
        """
        Support adding Usage objects in sum() operations, starting from 0.

        Args:
            other: Typically 0 when used in sum().

        Returns:
            This Usage object if other is 0, otherwise performs normal addition.
        """
        if other == 0:
            return self
        return self.__add__(other)  # type: ignore

    @classmethod
    def zero(cls) -> Usage:
        """
        Create an empty Usage object with zero tokens.

        Returns:
            A new Usage object with all token counts set to zero.
        """
        return cls(prompt_tokens=0, completion_tokens=0)
