"""Cost calculation for agent decisions.

This module provides the DecisionCostCalculator class for calculating
the cost of agent decisions without side effects (pure calculation).
"""

from pydantic import BaseModel, Field

from openclaw.core.config import CostStructure


class DecisionCostBreakdown(BaseModel):
    """Detailed breakdown of decision costs."""

    input_tokens: int = Field(..., ge=0, description="Number of input tokens")
    output_tokens: int = Field(..., ge=0, description="Number of output tokens")
    market_data_calls: int = Field(..., ge=0, description="Number of market data API calls")
    input_cost: float = Field(..., ge=0, description="Cost for input tokens")
    output_cost: float = Field(..., ge=0, description="Cost for output tokens")
    data_cost: float = Field(..., ge=0, description="Cost for market data calls")
    total_cost: float = Field(..., ge=0, description="Total cost")


class DecisionCostCalculator:
    """Calculate the cost of agent decisions.

    This calculator performs pure cost calculations without any side effects.
    It does NOT update any balances or track state - it only computes costs.

    Costs are read from EconomicTracker configuration or can be provided directly.

    Args:
        llm_input_per_1m: Cost per 1M input tokens (default: 2.5)
        llm_output_per_1m: Cost per 1M output tokens (default: 10.0)
        market_data_per_call: Cost per market data API call (default: 0.01)
    """

    def __init__(
        self,
        llm_input_per_1m: float = 2.5,
        llm_output_per_1m: float = 10.0,
        market_data_per_call: float = 0.01,
    ) -> None:
        self.llm_input_per_1m: float = llm_input_per_1m
        self.llm_output_per_1m: float = llm_output_per_1m
        self.market_data_per_call: float = market_data_per_call

    @classmethod
    def from_config(cls, config: CostStructure) -> "DecisionCostCalculator":
        """Create calculator from configuration.

        Args:
            config: CostStructure with cost parameters

        Returns:
            DecisionCostCalculator with configured rates
        """
        return cls(
            llm_input_per_1m=config.llm_input_per_1m,
            llm_output_per_1m=config.llm_output_per_1m,
            market_data_per_call=config.market_data_per_call,
        )

    def calculate_decision_cost(
        self,
        tokens_input: int,
        tokens_output: int,
        market_data_calls: int = 0,
    ) -> float:
        """Calculate the cost of a decision.

        Args:
            tokens_input: Number of input tokens to LLM
            tokens_output: Number of output tokens from LLM
            market_data_calls: Number of market data API calls

        Returns:
            Total cost of the decision (rounded to 4 decimal places)
        """
        # Calculate individual cost components
        input_cost = self._calculate_input_token_cost(tokens_input)
        output_cost = self._calculate_output_token_cost(tokens_output)
        data_cost = self._calculate_data_cost(market_data_calls)

        total_cost = input_cost + output_cost + data_cost
        return round(total_cost, 4)

    def calculate_detailed(
        self,
        tokens_input: int,
        tokens_output: int,
        market_data_calls: int = 0,
    ) -> DecisionCostBreakdown:
        """Calculate detailed cost breakdown for a decision.

        Args:
            tokens_input: Number of input tokens to LLM
            tokens_output: Number of output tokens from LLM
            market_data_calls: Number of market data API calls

        Returns:
            DecisionCostBreakdown with detailed cost breakdown
        """
        input_cost = self._calculate_input_token_cost(tokens_input)
        output_cost = self._calculate_output_token_cost(tokens_output)
        data_cost = self._calculate_data_cost(market_data_calls)

        return DecisionCostBreakdown(
            input_tokens=tokens_input,
            output_tokens=tokens_output,
            market_data_calls=market_data_calls,
            input_cost=round(input_cost, 4),
            output_cost=round(output_cost, 4),
            data_cost=round(data_cost, 4),
            total_cost=round(input_cost + output_cost + data_cost, 4),
        )

    def _calculate_input_token_cost(self, tokens: int) -> float:
        """Calculate cost for input tokens.

        Args:
            tokens: Number of input tokens

        Returns:
            Cost for input tokens
        """
        return tokens / 1e6 * self.llm_input_per_1m

    def _calculate_output_token_cost(self, tokens: int) -> float:
        """Calculate cost for output tokens.

        Args:
            tokens: Number of output tokens

        Returns:
            Cost for output tokens
        """
        return tokens / 1e6 * self.llm_output_per_1m

    def _calculate_data_cost(self, calls: int) -> float:
        """Calculate cost for market data API calls.

        Args:
            calls: Number of API calls

        Returns:
            Cost for market data calls
        """
        return calls * self.market_data_per_call

    def __repr__(self) -> str:
        """Return string representation of calculator."""
        return (
            f"DecisionCostCalculator("
            f"input=${self.llm_input_per_1m}/1M, "
            f"output=${self.llm_output_per_1m}/1M, "
            f"data=${self.market_data_per_call}/call)"
        )
