"""Tests for cost tracking module."""

import pytest

from agenteval.evaluation.cost_tracking import (
    CostBreakdown,
    ModelPricing,
    TokenUsage,
    aggregate_token_usage,
    calculate_case_cost,
    calculate_cost,
    calculate_run_cost,
    calculate_turn_cost,
)


def test_model_pricing_model():
    """ModelPricing model should work correctly."""
    pricing = ModelPricing(
        model_id="gpt-4o",
        prompt_cost_per_1m=5.0,
        completion_cost_per_1m=15.0,
    )
    assert pricing.model_id == "gpt-4o"
    assert pricing.prompt_cost_per_1m == 5.0


def test_token_usage_model():
    """TokenUsage model should work correctly."""
    usage = TokenUsage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
    assert usage.prompt_tokens == 100
    assert usage.completion_tokens == 50
    assert usage.total_tokens == 150


def test_calculate_cost_gpt4o_mini():
    """Cost calculation for gpt-4o-mini should be correct."""
    pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
    # 1000 prompt tokens + 500 completion tokens
    # Cost = (1000/1M * 0.15) + (500/1M * 0.60) = 0.00015 + 0.0003 = 0.00045
    cost = calculate_cost(1000, 500, pricing)
    assert abs(cost - 0.00045) < 0.00001


def test_calculate_cost_zero_tokens():
    """Zero tokens should result in zero cost."""
    pricing = ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0)
    cost = calculate_cost(0, 0, pricing)
    assert cost == 0.0


def test_aggregate_token_usage():
    """Aggregate token usage from multiple turns."""
    turns = [
        {"prompt_tokens": 100, "completion_tokens": 50},
        {"prompt_tokens": 200, "completion_tokens": 100},
        {"prompt_tokens": None, "completion_tokens": None},  # Missing data
    ]
    usage = aggregate_token_usage(turns)
    assert usage.prompt_tokens == 300
    assert usage.completion_tokens == 150
    assert usage.total_tokens == 450


def test_aggregate_token_usage_empty():
    """Empty turns should result in zero usage."""
    usage = aggregate_token_usage([])
    assert usage.prompt_tokens == 0
    assert usage.completion_tokens == 0
    assert usage.total_tokens == 0


def test_calculate_turn_cost():
    """Calculate cost for a single turn."""
    pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
    turn = {"prompt_tokens": 1000, "completion_tokens": 500}
    breakdown = calculate_turn_cost(turn, pricing)
    assert breakdown.prompt_tokens == 1000
    assert breakdown.completion_tokens == 500
    assert breakdown.total_tokens == 1500
    assert abs(breakdown.cost_usd - 0.00045) < 0.00001


def test_calculate_turn_cost_missing_tokens():
    """Turn with missing token data should have zero cost."""
    pricing = ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0)
    turn = {}  # No token data
    breakdown = calculate_turn_cost(turn, pricing)
    assert breakdown.prompt_tokens == 0
    assert breakdown.completion_tokens == 0
    assert breakdown.cost_usd == 0.0


def test_calculate_case_cost():
    """Calculate cost for a case with multiple turns."""
    pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
    turns = [
        {"prompt_tokens": 1000, "completion_tokens": 500},
        {"prompt_tokens": 2000, "completion_tokens": 1000},
    ]
    breakdown = calculate_case_cost(turns, pricing)
    assert breakdown.prompt_tokens == 3000
    assert breakdown.completion_tokens == 1500
    assert breakdown.total_tokens == 4500
    # Cost = (3000/1M * 0.15) + (1500/1M * 0.60) = 0.00045 + 0.0009 = 0.00135
    assert abs(breakdown.cost_usd - 0.00135) < 0.00001


def test_calculate_run_cost():
    """Calculate total cost for a run with multiple cases."""
    pricing = ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60)
    cases = [
        {
            "case_id": "c1",
            "turns": [
                {"prompt_tokens": 1000, "completion_tokens": 500},
            ],
        },
        {
            "case_id": "c2",
            "turns": [
                {"prompt_tokens": 2000, "completion_tokens": 1000},
            ],
        },
    ]
    breakdown = calculate_run_cost(cases, pricing)
    assert breakdown.prompt_tokens == 3000
    assert breakdown.completion_tokens == 1500
    assert breakdown.total_tokens == 4500


def test_cost_breakdown_model():
    """CostBreakdown model should work correctly."""
    breakdown = CostBreakdown(
        prompt_tokens=1000,
        completion_tokens=500,
        total_tokens=1500,
        cost_usd=0.001,
    )
    assert breakdown.prompt_tokens == 1000
    assert breakdown.cost_usd == 0.001