doczyai-pipelines/fieldExtraction/tests/test_usage_tracking.py

"""Unit tests for usage_tracking module.

Phase 1: Core Functions
- extract_usage_from_response()
- get_cost_per_token()
- track_usage()

Phase 2: Data Access Functions
- get_usage_data()
- reset_usage_data()

Phase 3: Export Functions
- get_usage_dataframes()
- _generate_per_file_per_model_df()
- _generate_batch_summary_df()
- export_usage_cost_csvs_local()

Phase 4: Utility Functions
- get_model_name()
"""

import logging
import threading
import time

import pandas as pd
import pytest

from src.utils import usage_tracking


class TestExtractUsageFromResponse:
    """Tests for extract_usage_from_response() function."""

    @pytest.mark.parametrize(
        "response_body, expected",
        [
            # Valid response with all token types
            (
                {
                    "usage": {
                        "input_tokens": 100,
                        "output_tokens": 50,
                        "cache_creation_input_tokens": 20,
                        "cache_read_input_tokens": 30,
                    }
                },
                (100, 50, 20, 30),
            ),
            # Response with only input/output (no cache tokens)
            (
                {
                    "usage": {
                        "input_tokens": 100,
                        "output_tokens": 50,
                    }
                },
                (100, 50, 0, 0),
            ),
            # Response with cache tokens only
            (
                {
                    "usage": {
                        "cache_creation_input_tokens": 20,
                        "cache_read_input_tokens": 30,
                    }
                },
                (0, 0, 20, 30),
            ),
            # Response with missing optional fields
            (
                {
                    "usage": {
                        "input_tokens": 100,
                    }
                },
                (100, 0, 0, 0),
            ),
            # Response with zero values
            (
                {
                    "usage": {
                        "input_tokens": 0,
                        "output_tokens": 0,
                        "cache_creation_input_tokens": 0,
                        "cache_read_input_tokens": 0,
                    }
                },
                (0, 0, 0, 0),
            ),
            # Response with string numbers (should convert to int)
            (
                {
                    "usage": {
                        "input_tokens": "100",
                        "output_tokens": "50",
                        "cache_creation_input_tokens": "20",
                        "cache_read_input_tokens": "30",
                    }
                },
                (100, 50, 20, 30),
            ),
        ],
    )
    def test_extract_usage_from_response_valid(self, response_body, expected):
        """Test extraction with valid response bodies."""
        result = usage_tracking.extract_usage_from_response(response_body)
        assert result == expected
        # Verify all values are integers
        assert all(isinstance(x, int) for x in result)

    @pytest.mark.parametrize(
        "response_body, expected",
        [
            # Missing "usage" key
            ({}, (0, 0, 0, 0)),
            # Empty "usage" dict
            ({"usage": {}}, (0, 0, 0, 0)),
            # None response
            (None, (0, 0, 0, 0)),
        ],
    )
    def test_extract_usage_from_response_missing_usage(self, response_body, expected, caplog):
        """Test extraction when usage data is missing."""
        with caplog.at_level(logging.DEBUG):
            result = usage_tracking.extract_usage_from_response(response_body)
        assert result == expected
        if response_body is None or "usage" not in (response_body or {}):
            assert "No usage data in API response" in caplog.text

    def test_extract_usage_from_response_nested_structure(self):
        """Test that function handles nested structure correctly."""
        response_body = {
            "other_data": "value",
            "usage": {
                "input_tokens": 100,
                "output_tokens": 50,
                "cache_creation_input_tokens": 20,
                "cache_read_input_tokens": 30,
            },
            "more_data": {"nested": "value"},
        }
        result = usage_tracking.extract_usage_from_response(response_body)
        assert result == (100, 50, 20, 30)


class TestGetCostPerToken:
    """Tests for get_cost_per_token() function."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    @pytest.mark.parametrize(
        "model_id, expected_input, expected_output, expected_cache_read",
        [
            # Note: Claude 2 support has been removed - only Claude 3+ models are supported
            # Claude 3 Haiku
            (
                "anthropic.claude-3-haiku-20240307-v1:0",
                0.00025 / 1000.0,
                0.00125 / 1000.0,
                0.00003 / 1000.0,  # Corrected to match llm_utils.py (was 0.000025)
            ),
            # Claude 3.5 Sonnet
            (
                "anthropic.claude-3-5-sonnet-20240620-v1:0",
                0.003 / 1000.0,
                0.015 / 1000.0,
                0.0003 / 1000.0,
            ),
            (
                "anthropic.claude-3-5-sonnet-20241022-v2:0",
                0.003 / 1000.0,
                0.015 / 1000.0,
                0.0003 / 1000.0,
            ),
            # Claude 3.7 Sonnet
            (
                "us.anthropic.claude-3-7-sonnet-20250219-v1:0",
                0.003 / 1000.0,
                0.015 / 1000.0,
                0.0003 / 1000.0,
            ),
            # Claude 4 Sonnet
            (
                "us.anthropic.claude-sonnet-4-20250514-v1:0",
                0.003 / 1000.0,
                0.015 / 1000.0,
                0.0003 / 1000.0,
            ),
            # Claude 4.5 Sonnet
            (
                "us.anthropic.claude-sonnet-4-5-20250929-v1:0",
                0.003 / 1000.0,
                0.015 / 1000.0,
                0.0003 / 1000.0,
            ),
        ],
    )
    def test_get_cost_per_token_valid_models(
        self, model_id, expected_input, expected_output, expected_cache_read, mocker
    ):
        """Test cost calculation for valid model IDs."""
        # Mock resolve_model_id to return the model_id as-is
        mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id)

        input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id)

        assert abs(input_cost - expected_input) < 1e-10
        assert abs(output_cost - expected_output) < 1e-10
        assert abs(cache_read_cost - expected_cache_read) < 1e-10

    def test_get_cost_per_token_unknown_model(self, mocker, caplog):
        """Test cost calculation for unknown model (should use default)."""
        unknown_model = "unknown.model.id"
        # Mock resolve_model_id to return the unknown model
        mocker.patch(
            "src.utils.usage_tracking.config.resolve_model_id", return_value=unknown_model
        )

        with caplog.at_level(logging.WARNING):
            input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(
                unknown_model
            )

        # Should use default pricing (Claude 3.5 Sonnet)
        expected_input = 0.003 / 1000.0
        expected_output = 0.015 / 1000.0
        expected_cache_read = 0.0003 / 1000.0

        assert abs(input_cost - expected_input) < 1e-10
        assert abs(output_cost - expected_output) < 1e-10
        assert abs(cache_read_cost - expected_cache_read) < 1e-10
        assert "Unknown model" in caplog.text
        assert "using default pricing" in caplog.text

    def test_get_cost_per_token_model_alias(self, mocker):
        """Test that model aliases are resolved correctly."""
        alias = "sonnet_37"
        resolved_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
        # Mock resolve_model_id to resolve the alias
        mocker.patch(
            "src.utils.usage_tracking.config.resolve_model_id", return_value=resolved_id
        )

        input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(alias)

        # Should use Claude 3.7 Sonnet pricing
        expected_input = 0.003 / 1000.0
        expected_output = 0.015 / 1000.0
        expected_cache_read = 0.0003 / 1000.0

        assert abs(input_cost - expected_input) < 1e-10
        assert abs(output_cost - expected_output) < 1e-10
        assert abs(cache_read_cost - expected_cache_read) < 1e-10

    def test_get_cost_per_token_conversion(self, mocker):
        """Test that per-1k costs are correctly converted to per-token."""
        model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
        mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id)

        input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id)

        # Verify conversion: per-1k / 1000 = per-token
        assert input_cost == 0.003 / 1000.0
        assert output_cost == 0.015 / 1000.0
        assert cache_read_cost == 0.0003 / 1000.0

        # Verify all costs are floats
        assert isinstance(input_cost, float)
        assert isinstance(output_cost, float)
        assert isinstance(cache_read_cost, float)


class TestTrackUsage:
    """Tests for track_usage() function."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    @pytest.mark.parametrize(
        "filename, model_id, input_tokens, output_tokens, cache_creation, cache_read",
        [
            # All token types present
            ("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 20, 30),
            # No cache tokens
            ("file2.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 0, 0),
            # Only cache tokens
            ("file3.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 20, 30),
            # Zero tokens
            ("file4.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 0, 0),
        ],
    )
    def test_track_usage_single_call(
        self,
        filename,
        model_id,
        input_tokens,
        output_tokens,
        cache_creation,
        cache_read,
        mocker,
    ):
        """Test tracking a single API call."""
        # Mock get_cost_per_token to return known costs
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage(
            filename, model_id, input_tokens, output_tokens, cache_creation, cache_read
        )

        # Verify per-file, per-model data
        usage_data = usage_tracking.get_usage_data()
        file_data = usage_data["per_file_per_model"][filename][model_id]

        assert file_data["input_tokens"] == input_tokens
        assert file_data["output_tokens"] == output_tokens
        assert file_data["cache_creation_tokens"] == cache_creation
        assert file_data["cache_read_tokens"] == cache_read
        assert file_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read
        assert file_data["request_count"] == 1

        # Verify cost calculation
        # For Sonnet models: cache creation is 25% higher than input rate (0.00375 vs 0.003)
        # For Haiku: cache creation is 20% higher than input rate (0.0003 vs 0.00025)
        model_id_lower = model_id.lower()
        if any(
            x in model_id_lower
            for x in ["claude-3-5-sonnet", "claude-3-7-sonnet", "claude-sonnet-4", "claude-4-sonnet"]
        ):
            cache_creation_cost = input_cost * 1.25  # Sonnet: 25% higher
        elif "haiku" in model_id_lower:
            cache_creation_cost = input_cost * 1.20  # Haiku: 20% higher
        else:
            cache_creation_cost = input_cost  # Others: same as input
        
        expected_cost = (
            (input_tokens * input_cost)
            + (cache_creation * cache_creation_cost)
            + (cache_read * cache_read_cost)
            + (output_tokens * output_cost)
        )
        assert abs(file_data["total_cost"] - expected_cost) < 1e-10

        # Verify global totals
        global_data = usage_data["global"]
        assert global_data["total_fresh_input_tokens"] == input_tokens
        assert global_data["total_cache_tokens"] == cache_creation + cache_read
        assert global_data["total_input_tokens"] == input_tokens + cache_creation + cache_read
        assert global_data["total_output_tokens"] == output_tokens
        assert global_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read
        assert abs(global_data["total_cost"] - expected_cost) < 1e-10
        assert global_data["total_requests"] == 1

    def test_track_usage_multiple_calls_same_file_model(self, mocker):
        """Test tracking multiple calls for same file/model (accumulation)."""
        filename = "file1.txt"
        model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0"
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # First call
        usage_tracking.track_usage(filename, model_id, 100, 50, 10, 20)
        # Second call
        usage_tracking.track_usage(filename, model_id, 200, 100, 20, 30)

        usage_data = usage_tracking.get_usage_data()
        file_data = usage_data["per_file_per_model"][filename][model_id]

        # Verify accumulation
        assert file_data["input_tokens"] == 300  # 100 + 200
        assert file_data["output_tokens"] == 150  # 50 + 100
        assert file_data["cache_creation_tokens"] == 30  # 10 + 20
        assert file_data["cache_read_tokens"] == 50  # 20 + 30
        assert file_data["total_tokens"] == 530  # 300 + 150 + 30 + 50
        assert file_data["request_count"] == 2

        # Verify global totals
        global_data = usage_data["global"]
        assert global_data["total_fresh_input_tokens"] == 300
        assert global_data["total_cache_tokens"] == 80  # 30 + 50
        assert global_data["total_input_tokens"] == 380  # 300 + 80
        assert global_data["total_output_tokens"] == 150
        assert global_data["total_tokens"] == 530
        assert global_data["total_requests"] == 2

    def test_track_usage_multiple_files_models(self, mocker):
        """Test tracking calls for different files/models (isolation)."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # File 1, Model 1
        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        # File 1, Model 2
        usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
        # File 2, Model 1
        usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)

        usage_data = usage_tracking.get_usage_data()

        # Verify file1, model1
        file1_model1 = usage_data["per_file_per_model"]["file1.txt"]["model1"]
        assert file1_model1["input_tokens"] == 100
        assert file1_model1["request_count"] == 1

        # Verify file1, model2
        file1_model2 = usage_data["per_file_per_model"]["file1.txt"]["model2"]
        assert file1_model2["input_tokens"] == 200
        assert file1_model2["request_count"] == 1

        # Verify file2, model1
        file2_model1 = usage_data["per_file_per_model"]["file2.txt"]["model1"]
        assert file2_model1["input_tokens"] == 150
        assert file2_model1["request_count"] == 1

        # Verify global totals (sum of all)
        global_data = usage_data["global"]
        assert global_data["total_fresh_input_tokens"] == 450  # 100 + 200 + 150
        assert global_data["total_cache_tokens"] == 120  # (10+20) + (20+30) + (15+25) = 30 + 50 + 40
        assert global_data["total_requests"] == 3

    @pytest.mark.parametrize(
        "filename, input_tokens, output_tokens, cache_creation, cache_read, should_reject",
        [
            # Empty filename
            ("", 100, 50, 0, 0, True),
            (None, 100, 50, 0, 0, True),
            # Negative input tokens
            ("file.txt", -1, 50, 0, 0, True),
            # Negative output tokens
            ("file.txt", 100, -1, 0, 0, True),
            # Negative cache creation tokens
            ("file.txt", 100, 50, -1, 0, True),
            # Negative cache read tokens
            ("file.txt", 100, 50, 0, -1, True),
            # Valid (zero values are OK)
            ("file.txt", 0, 0, 0, 0, False),
        ],
    )
    def test_track_usage_input_validation(
        self,
        filename,
        input_tokens,
        output_tokens,
        cache_creation,
        cache_read,
        should_reject,
        mocker,
        caplog,
    ):
        """Test input validation (should reject invalid inputs)."""
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(0.001, 0.002, 0.0001),
        )

        initial_data = usage_tracking.get_usage_data()
        initial_requests = initial_data["global"]["total_requests"]

        with caplog.at_level(logging.WARNING):
            usage_tracking.track_usage(
                filename, "model1", input_tokens, output_tokens, cache_creation, cache_read
            )

        final_data = usage_tracking.get_usage_data()
        final_requests = final_data["global"]["total_requests"]

        if should_reject:
            # Should not have updated usage data
            assert final_requests == initial_requests
            assert "Invalid usage data" in caplog.text
        else:
            # Should have updated usage data
            assert final_requests == initial_requests + 1

    def test_track_usage_thread_safety(self, mocker):
        """Test thread safety with concurrent calls."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        num_threads = 10
        calls_per_thread = 50
        filename = "file1.txt"
        model_id = "model1"

        def track_multiple_calls():
            for _ in range(calls_per_thread):
                usage_tracking.track_usage(filename, model_id, 10, 5, 2, 3)
                time.sleep(0.001)  # Small delay to increase chance of race condition

        threads = []
        for _ in range(num_threads):
            thread = threading.Thread(target=track_multiple_calls)
            threads.append(thread)
            thread.start()

        # Wait for all threads to complete
        for thread in threads:
            thread.join()

        usage_data = usage_tracking.get_usage_data()
        file_data = usage_data["per_file_per_model"][filename][model_id]

        # Verify all calls were tracked (no data loss)
        expected_requests = num_threads * calls_per_thread
        assert file_data["request_count"] == expected_requests

        # Verify token totals are correct
        expected_input = expected_requests * 10
        expected_output = expected_requests * 5
        expected_cache_creation = expected_requests * 2
        expected_cache_read = expected_requests * 3

        assert file_data["input_tokens"] == expected_input
        assert file_data["output_tokens"] == expected_output
        assert file_data["cache_creation_tokens"] == expected_cache_creation
        assert file_data["cache_read_tokens"] == expected_cache_read

        # Verify global totals
        global_data = usage_data["global"]
        assert global_data["total_requests"] == expected_requests
        assert global_data["total_fresh_input_tokens"] == expected_input
        assert global_data["total_cache_tokens"] == expected_cache_creation + expected_cache_read

    def test_track_usage_cost_calculation(self, mocker):
        """Test that cost calculation is correct for all token types."""
        filename = "file1.txt"
        model_id = "model1"
        input_cost = 0.003 / 1000.0  # $0.003 per 1k = $0.000003 per token
        output_cost = 0.015 / 1000.0  # $0.015 per 1k = $0.000015 per token
        cache_read_cost = 0.0003 / 1000.0  # $0.0003 per 1k = $0.0000003 per token
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # Track usage with specific token counts
        input_tokens = 1000
        output_tokens = 500
        cache_creation_tokens = 200  # Charged at input rate
        cache_read_tokens = 300  # Charged at cache read rate

        usage_tracking.track_usage(
            filename, model_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens
        )

        usage_data = usage_tracking.get_usage_data()
        file_data = usage_data["per_file_per_model"][filename][model_id]

        # Calculate expected cost manually
        expected_cost = (
            (input_tokens * input_cost)
            + (cache_creation_tokens * input_cost)  # Same as input rate
            + (cache_read_tokens * cache_read_cost)  # Lower rate
            + (output_tokens * output_cost)
        )

        assert abs(file_data["total_cost"] - expected_cost) < 1e-10

        # Verify cache creation uses input rate, not cache read rate
        cache_creation_cost = cache_creation_tokens * input_cost
        cache_read_cost_calc = cache_read_tokens * cache_read_cost
        assert cache_creation_cost > cache_read_cost_calc  # Input rate > cache read rate


class TestGetUsageData:
    """Tests for get_usage_data() function - Phase 2."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_get_usage_data_empty_state(self):
        """Test getting usage data when no usage has been tracked."""
        usage_data = usage_tracking.get_usage_data()

        assert "per_file_per_model" in usage_data
        assert "global" in usage_data
        assert usage_data["per_file_per_model"] == {}
        assert usage_data["global"]["total_tokens"] == 0
        assert usage_data["global"]["total_cost"] == 0.0
        assert usage_data["global"]["total_requests"] == 0

    def test_get_usage_data_single_file_single_model(self, mocker):
        """Test getting usage data with single file and model."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_data = usage_tracking.get_usage_data()

        assert "per_file_per_model" in usage_data
        assert "global" in usage_data
        assert "file1.txt" in usage_data["per_file_per_model"]
        assert "model1" in usage_data["per_file_per_model"]["file1.txt"]
        assert usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100
        assert usage_data["global"]["total_fresh_input_tokens"] == 100

    def test_get_usage_data_multiple_files_single_model(self, mocker):
        """Test getting usage data with multiple files and single model."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0)
        usage_data = usage_tracking.get_usage_data()

        assert len(usage_data["per_file_per_model"]) == 2
        assert "file1.txt" in usage_data["per_file_per_model"]
        assert "file2.txt" in usage_data["per_file_per_model"]
        assert usage_data["global"]["total_fresh_input_tokens"] == 300

    def test_get_usage_data_single_file_multiple_models(self, mocker):
        """Test getting usage data with single file and multiple models."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0)
        usage_data = usage_tracking.get_usage_data()

        assert len(usage_data["per_file_per_model"]) == 1
        assert "file1.txt" in usage_data["per_file_per_model"]
        assert len(usage_data["per_file_per_model"]["file1.txt"]) == 2
        assert "model1" in usage_data["per_file_per_model"]["file1.txt"]
        assert "model2" in usage_data["per_file_per_model"]["file1.txt"]

    def test_get_usage_data_multiple_files_multiple_models(self, mocker):
        """Test getting usage data with multiple files and models."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
        usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)
        usage_data = usage_tracking.get_usage_data()

        assert len(usage_data["per_file_per_model"]) == 2
        assert usage_data["global"]["total_fresh_input_tokens"] == 450
        assert usage_data["global"]["total_cache_tokens"] == 120
        assert usage_data["global"]["total_requests"] == 3

    def test_get_usage_data_deep_copy(self, mocker):
        """Test that modifying returned dict doesn't affect internal state."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_data = usage_tracking.get_usage_data()

        # Modify the returned data
        usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] = 999
        usage_data["global"]["total_tokens"] = 999

        # Get fresh data and verify original values are unchanged
        fresh_data = usage_tracking.get_usage_data()
        assert fresh_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100
        assert fresh_data["global"]["total_tokens"] == 150  # 100 + 50

    def test_get_usage_data_thread_safety(self, mocker):
        """Test thread safety when getting data while other thread is updating."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        results = []
        errors = []

        def track_and_get():
            try:
                for i in range(10):
                    usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0)
                    data = usage_tracking.get_usage_data()
                    results.append(data["global"]["total_requests"])
            except Exception as e:
                errors.append(e)

        threads = []
        for _ in range(5):
            thread = threading.Thread(target=track_and_get)
            threads.append(thread)
            thread.start()

        for thread in threads:
            thread.join()

        # Verify no errors occurred
        assert len(errors) == 0

        # Verify final state is correct
        final_data = usage_tracking.get_usage_data()
        assert final_data["global"]["total_requests"] == 50  # 5 threads * 10 calls


class TestResetUsageData:
    """Tests for reset_usage_data() function - Phase 2."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_reset_usage_data_after_tracking(self, mocker):
        """Test reset after tracking some usage."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # Track some usage
        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_tracking.track_usage("file2.txt", "model2", 200, 100, 20, 30)

        # Verify data exists
        usage_data = usage_tracking.get_usage_data()
        assert usage_data["global"]["total_requests"] == 2
        assert len(usage_data["per_file_per_model"]) == 2

        # Reset
        usage_tracking.reset_usage_data()

        # Verify all data is cleared
        usage_data = usage_tracking.get_usage_data()
        assert usage_data["global"]["total_tokens"] == 0
        assert usage_data["global"]["total_cost"] == 0.0
        assert usage_data["global"]["total_requests"] == 0
        assert usage_data["global"]["total_fresh_input_tokens"] == 0
        assert usage_data["global"]["total_cache_tokens"] == 0
        assert usage_data["per_file_per_model"] == {}

    def test_reset_usage_data_when_empty(self):
        """Test reset when already empty (should not error)."""
        # Reset when already empty
        usage_tracking.reset_usage_data()

        # Verify still empty
        usage_data = usage_tracking.get_usage_data()
        assert usage_data["global"]["total_requests"] == 0
        assert usage_data["per_file_per_model"] == {}

        # Reset again (should not error)
        usage_tracking.reset_usage_data()
        usage_data = usage_tracking.get_usage_data()
        assert usage_data["global"]["total_requests"] == 0

    def test_reset_usage_data_clears_both_structures(self, mocker):
        """Test that reset clears both USAGE_DATA and GLOBAL_USAGE."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_data = usage_tracking.get_usage_data()

        # Verify both structures have data
        assert len(usage_data["per_file_per_model"]) > 0
        assert usage_data["global"]["total_requests"] > 0

        # Reset
        usage_tracking.reset_usage_data()
        usage_data = usage_tracking.get_usage_data()

        # Verify both structures are cleared
        assert usage_data["per_file_per_model"] == {}
        assert usage_data["global"]["total_requests"] == 0
        assert usage_data["global"]["total_tokens"] == 0
        assert usage_data["global"]["total_cost"] == 0.0
        assert usage_data["global"]["total_fresh_input_tokens"] == 0
        assert usage_data["global"]["total_cache_tokens"] == 0

    def test_reset_usage_data_thread_safety(self, mocker):
        """Test thread safety of reset operation."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        errors = []

        def track_and_reset():
            try:
                for i in range(10):
                    usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0)
                usage_tracking.reset_usage_data()
            except Exception as e:
                errors.append(e)

        threads = []
        for _ in range(5):
            thread = threading.Thread(target=track_and_reset)
            threads.append(thread)
            thread.start()

        for thread in threads:
            thread.join()

        # Verify no errors occurred
        assert len(errors) == 0

        # Verify final state is reset
        usage_data = usage_tracking.get_usage_data()
        assert usage_data["global"]["total_requests"] == 0
        assert usage_data["per_file_per_model"] == {}


class TestGetUsageDataframes:
    """Tests for get_usage_dataframes() function - Phase 3."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_get_usage_dataframes_empty(self):
        """Test getting DataFrames when no usage data exists."""
        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"

        per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)

        assert isinstance(per_file_df, pd.DataFrame)
        assert isinstance(summary_df, pd.DataFrame)
        assert per_file_df.empty
        assert not summary_df.empty  # Summary should have one row with zeros

    def test_get_usage_dataframes_single_file_single_model(self, mocker):
        """Test getting DataFrames with single file and model."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)

        assert isinstance(per_file_df, pd.DataFrame)
        assert isinstance(summary_df, pd.DataFrame)
        assert len(per_file_df) == 1
        assert len(summary_df) == 1

        # Verify run_id and timestamp are included
        assert per_file_df["run_id"].iloc[0] == run_id
        assert per_file_df["timestamp"].iloc[0] == timestamp
        assert summary_df["run_id"].iloc[0] == run_id
        assert summary_df["timestamp"].iloc[0] == timestamp

    def test_get_usage_dataframes_multiple_files_models(self, mocker):
        """Test getting DataFrames with multiple files and models."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30)
        usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)

        assert len(per_file_df) == 3  # 3 file/model combinations
        assert len(summary_df) == 1  # Single summary row

    def test_get_usage_dataframes_independent(self, mocker):
        """Test that modifying one DataFrame doesn't affect the other."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp)

        # Modify per_file_df
        per_file_df.loc[0, "input_tokens"] = 999

        # Verify summary_df is unchanged
        assert summary_df["total_fresh_input_tokens"].iloc[0] == 100


class TestGeneratePerFilePerModelDf:
    """Tests for _generate_per_file_per_model_df() function - Phase 3."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_generate_per_file_per_model_df_empty(self):
        """Test generating DataFrame when no usage data exists."""
        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"

        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        assert isinstance(df, pd.DataFrame)
        assert df.empty

    def test_generate_per_file_per_model_df_single_file_single_model(self, mocker):
        """Test generating DataFrame with single file and model."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 10, 20)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        assert len(df) == 1
        assert list(df.columns) == [
            "run_id",
            "timestamp",
            "file_name",
            "model_name",
            "input_tokens",
            "output_tokens",
            "cache_creation_tokens",
            "cache_read_tokens",
            "total_tokens",
            "cost_per_token",
            "total_cost",
        ]

        # Verify values
        assert df["run_id"].iloc[0] == run_id
        assert df["timestamp"].iloc[0] == timestamp
        assert df["file_name"].iloc[0] == "file1.txt"
        assert df["model_name"].iloc[0] == "Claude 3.5 Sonnet"  # Human-readable
        assert df["input_tokens"].iloc[0] == 100
        assert df["output_tokens"].iloc[0] == 50
        assert df["cache_creation_tokens"].iloc[0] == 10
        assert df["cache_read_tokens"].iloc[0] == 20
        assert df["total_tokens"].iloc[0] == 180  # 100 + 50 + 10 + 20

    def test_generate_per_file_per_model_df_multiple_files_single_model(self, mocker):
        """Test generating DataFrame with multiple files and single model."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        assert len(df) == 2
        assert set(df["file_name"].values) == {"file1.txt", "file2.txt"}

    def test_generate_per_file_per_model_df_single_file_multiple_models(self, mocker):
        """Test generating DataFrame with single file and multiple models."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        assert len(df) == 2
        assert set(df["model_name"].values) == {"Claude Unknown", "Claude Unknown"}  # Fallback for unknown models

    def test_generate_per_file_per_model_df_cost_calculation(self, mocker):
        """Test cost_per_token calculation and rounding."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 1000, 500, 200, 300)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        # Verify cost_per_token is calculated correctly (total_cost / total_tokens)
        total_tokens = 1000 + 500 + 200 + 300  # 2000
        total_cost = (
            (1000 * input_cost)
            + (200 * input_cost)  # Cache creation uses input rate
            + (300 * cache_read_cost)
            + (500 * output_cost)
        )
        expected_cost_per_token = total_cost / total_tokens

        # Allow for floating point precision (rounding happens in the function)
        assert abs(df["cost_per_token"].iloc[0] - round(expected_cost_per_token, 8)) < 1e-10
        assert df["cost_per_token"].iloc[0] == round(expected_cost_per_token, 8)
        assert df["total_cost"].iloc[0] == round(total_cost, 6)

    def test_generate_per_file_per_model_df_skips_zero_requests(self, mocker):
        """Test that files/models with request_count=0 are skipped."""
        # This is tricky - we need to create a defaultdict entry but with 0 requests
        # In practice, this shouldn't happen, but we test the guard clause
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # Track some usage
        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp)

        # Should only have the one entry we tracked
        assert len(df) == 1
        assert df["file_name"].iloc[0] == "file1.txt"


class TestGenerateBatchSummaryDf:
    """Tests for _generate_batch_summary_df() function - Phase 3."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_generate_batch_summary_df_empty(self):
        """Test generating summary DataFrame when no usage data exists."""
        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"

        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        assert isinstance(df, pd.DataFrame)
        assert len(df) == 1  # Single row
        assert df["total_tokens"].iloc[0] == 0
        assert df["total_cost"].iloc[0] == 0.0

    def test_generate_batch_summary_df_single_file(self, mocker):
        """Test generating summary DataFrame with single file."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        assert len(df) == 1
        assert df["total_fresh_input_tokens"].iloc[0] == 100
        assert df["total_cache_tokens"].iloc[0] == 30  # 10 + 20
        assert df["total_input_tokens"].iloc[0] == 130  # 100 + 30
        assert df["total_output_tokens"].iloc[0] == 50
        assert df["total_tokens"].iloc[0] == 180  # 100 + 50 + 10 + 20

        # Verify averages (divide by 1 file)
        assert df["average_fresh_input_tokens"].iloc[0] == 100.0
        assert df["average_cache_tokens"].iloc[0] == 30.0

    def test_generate_batch_summary_df_multiple_files(self, mocker):
        """Test generating summary DataFrame with multiple files."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)
        usage_tracking.track_usage("file2.txt", "model1", 200, 100, 20, 30)
        usage_tracking.track_usage("file3.txt", "model2", 150, 75, 15, 25)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        assert len(df) == 1
        # Verify totals
        assert df["total_fresh_input_tokens"].iloc[0] == 450  # 100 + 200 + 150
        assert df["total_cache_tokens"].iloc[0] == 120  # (10+20) + (20+30) + (15+25)
        assert df["total_input_tokens"].iloc[0] == 570  # 450 + 120

        # Verify averages (divide by 3 files)
        assert df["average_fresh_input_tokens"].iloc[0] == 150.0  # 450 / 3
        assert df["average_cache_tokens"].iloc[0] == 40.0  # 120 / 3

    def test_generate_batch_summary_df_columns(self, mocker):
        """Test that summary DataFrame has all required columns."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        expected_columns = [
            "run_id",
            "timestamp",
            "average_fresh_input_tokens",
            "average_cache_tokens",
            "average_input_tokens",
            "average_output_tokens",
            "average_total_tokens",
            "average_cost",
            "total_fresh_input_tokens",
            "total_cache_tokens",
            "total_input_tokens",
            "total_output_tokens",
            "total_tokens",
            "total_cost",
            "region_mode",
        ]

        assert list(df.columns) == expected_columns

    def test_generate_batch_summary_df_region_mode_single(self, mocker):
        """Test region_mode is set correctly for single-region mode."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )
        mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", False)

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        assert df["region_mode"].iloc[0] == "single-region"

    def test_generate_batch_summary_df_region_mode_cross(self, mocker):
        """Test region_mode is set correctly for cross-region mode."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )
        mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", True)

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        assert df["region_mode"].iloc[0] == "cross-region"

    def test_generate_batch_summary_df_rounding(self, mocker):
        """Test that averages are properly rounded."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        # Create scenario where division results in many decimal places
        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)
        usage_tracking.track_usage("file2.txt", "model1", 333, 167, 0, 0)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        # Verify rounding (to 2 decimal places for averages)
        assert isinstance(df["average_fresh_input_tokens"].iloc[0], float)
        # 433 / 2 = 216.5, should be rounded to 2 decimals
        assert df["average_fresh_input_tokens"].iloc[0] == 216.5

    def test_generate_batch_summary_df_token_relationship(self, mocker):
        """Test that total_fresh_input_tokens + total_cache_tokens = total_input_tokens."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)

        run_id = "test_batch"
        timestamp = "run_20240101_12-00_test_batch"
        df = usage_tracking._generate_batch_summary_df(run_id, timestamp)

        fresh = df["total_fresh_input_tokens"].iloc[0]
        cache = df["total_cache_tokens"].iloc[0]
        total_input = df["total_input_tokens"].iloc[0]

        assert fresh + cache == total_input


class TestExportUsageCostCsvsLocal:
    """Tests for export_usage_cost_csvs_local() function - Phase 3."""

    @pytest.fixture(autouse=True)
    def reset_before_test(self):
        """Reset usage data before each test."""
        usage_tracking.reset_usage_data()
        yield
        usage_tracking.reset_usage_data()

    def test_export_usage_cost_csvs_local_with_data(self, mocker, tmp_path, caplog):
        """Test exporting CSV files with valid usage data."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )
        mocker.patch(
            "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20)

        run_id = "test_batch"
        run_timestamp = "run_20240101_12-00_test_batch"

        with caplog.at_level(logging.INFO):
            usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)

        # Verify directory was created
        output_dir = tmp_path / run_timestamp
        assert output_dir.exists()

        # Verify both files were created
        usage_file = output_dir / f"{run_id}-USAGE.csv"
        summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv"

        assert usage_file.exists()
        assert summary_file.exists()

        # Verify file contents can be read back
        usage_df = pd.read_csv(usage_file)
        summary_df = pd.read_csv(summary_file)

        assert len(usage_df) == 1
        assert len(summary_df) == 1
        assert usage_df["file_name"].iloc[0] == "file1.txt"

        # Verify logging
        assert "Saved usage data to:" in caplog.text
        assert "Saved usage summary to:" in caplog.text

    def test_export_usage_cost_csvs_local_empty_data(self, mocker, tmp_path, caplog):
        """Test exporting CSV files with empty usage data."""
        mocker.patch(
            "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
        )

        run_id = "test_batch"
        run_timestamp = "run_20240101_12-00_test_batch"

        with caplog.at_level(logging.WARNING):
            usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)

        # Verify directory was created
        output_dir = tmp_path / run_timestamp
        assert output_dir.exists()

        # Verify files were NOT created (empty data)
        usage_file = output_dir / f"{run_id}-USAGE.csv"
        summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv"

        assert not usage_file.exists()
        assert summary_file.exists()  # Summary should still exist (with zeros)

        # Verify logging
        assert "No usage data to export" in caplog.text

    def test_export_usage_cost_csvs_local_directory_creation(self, mocker, tmp_path):
        """Test that directory is created if it doesn't exist."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )
        mocker.patch(
            "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        run_timestamp = "run_20240101_12-00_test_batch"

        # Directory shouldn't exist yet
        output_dir = tmp_path / run_timestamp
        assert not output_dir.exists()

        usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)

        # Directory should now exist
        assert output_dir.exists()

    def test_export_usage_cost_csvs_local_file_paths(self, mocker, tmp_path):
        """Test that file paths are correct."""
        input_cost = 0.003 / 1000.0
        output_cost = 0.015 / 1000.0
        cache_read_cost = 0.0003 / 1000.0
        mocker.patch(
            "src.utils.usage_tracking.get_cost_per_token",
            return_value=(input_cost, output_cost, cache_read_cost),
        )
        mocker.patch(
            "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)
        )

        usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0)

        run_id = "test_batch"
        run_timestamp = "run_20240101_12-00_test_batch"

        usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp)

        # Verify exact file paths
        expected_usage_path = tmp_path / run_timestamp / f"{run_id}-USAGE.csv"
        expected_summary_path = tmp_path / run_timestamp / f"{run_id}-USAGE-SUMMARY.csv"

        assert expected_usage_path.exists()
        assert expected_summary_path.exists()


class TestGetModelName:
    """Tests for get_model_name() function - Phase 4."""

    @pytest.mark.parametrize(
        "model_id, expected_name",
        [
            # Note: Claude 2 support has been removed - only Claude 3+ models are supported
            # Claude 3 Haiku
            ("anthropic.claude-3-haiku-20240307-v1:0", "Claude 3 Haiku"),
            ("anthropic.claude-3.5-haiku-20240307-v1:0", "Claude 3 Haiku"),
            ("ANTHROPIC.CLAUDE-3-HAIKU-20240307-V1:0", "Claude 3 Haiku"),  # Uppercase
            # Claude 3.5 Sonnet
            ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
            ("anthropic.claude-3-5-sonnet-20241022-v2:0", "Claude 3.5 Sonnet"),
            ("anthropic.claude-3.5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
            ("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"),  # Mixed case
            # Claude 3.7 Sonnet
            ("us.anthropic.claude-3-7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"),
            ("us.anthropic.claude-3.7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"),
            ("US.ANTHROPIC.CLAUDE-3-7-SONNET-20250219-V1:0", "Claude 3.7 Sonnet"),  # Uppercase
            # Claude 4 Sonnet
            ("us.anthropic.claude-sonnet-4-20250514-v1:0", "Claude 4 Sonnet"),
            ("us.anthropic.claude-4-sonnet-20250514-v1:0", "Claude 4 Sonnet"),
            ("US.ANTHROPIC.CLAUDE-SONNET-4-20250514-V1:0", "Claude 4 Sonnet"),  # Uppercase
            # Claude 4.5 Sonnet
            ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"),
            ("us.anthropic.claude-sonnet-4.5-20250929-v1:0", "Claude 4.5 Sonnet"),
            # Note: claude-4-5-sonnet and claude-4.5-sonnet don't match the pattern
            # They need claude-sonnet-4 or claude-4-sonnet structure
            ("US.ANTHROPIC.CLAUDE-SONNET-4-5-20250929-V1:0", "Claude 4.5 Sonnet"),  # Uppercase
        ],
    )
    def test_get_model_name_valid_models(self, model_id, expected_name):
        """Test model name conversion for all supported models."""
        result = usage_tracking.get_model_name(model_id)
        assert result == expected_name

    @pytest.mark.parametrize(
        "model_id, expected_pattern",
        [
            # Unknown models should use fallback
            ("unknown.model.id", "Claude "),  # Should extract something or "Unknown"
            ("anthropic.claude-999", "Claude "),  # Unknown version
            ("some.random.model", "Claude "),  # Completely unknown
            ("", "Claude Unknown"),  # Empty string
        ],
    )
    def test_get_model_name_unknown_models(self, model_id, expected_pattern):
        """Test fallback behavior for unknown model IDs."""
        result = usage_tracking.get_model_name(model_id)
        assert result.startswith("Claude ")
        # Should either extract a version or say "Unknown"
        assert "Unknown" in result or result != "Claude "

    def test_get_model_name_case_insensitivity(self):
        """Test that function handles case variations correctly."""
        # Test various case combinations
        test_cases = [
            ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
            ("ANTHROPIC.CLAUDE-3-5-SONNET-20240620-V1:0", "Claude 3.5 Sonnet"),
            ("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"),
            ("AnThRoPiC.cLaUdE-3-5-sOnNeT-20240620-v1:0", "Claude 3.5 Sonnet"),
        ]

        for model_id, expected in test_cases:
            result = usage_tracking.get_model_name(model_id)
            assert result == expected

    def test_get_model_name_malformed_ids(self):
        """Test handling of malformed model IDs."""
        # Test various malformed inputs
        malformed_ids = [
            "claude",  # Too short
            ".",  # Just a dot
            "anthropic.",  # Ends with dot
            ".claude",  # Starts with dot
            "anthropic.claude-3-5",  # Missing version info
        ]

        for model_id in malformed_ids:
            result = usage_tracking.get_model_name(model_id)
            # Should not crash and should return something starting with "Claude"
            assert isinstance(result, str)
            assert result.startswith("Claude")

    def test_get_model_name_special_characters(self):
        """Test handling of special characters in model IDs."""
        # Test with special characters that might appear
        special_cases = [
            ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"),
            ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"),
        ]

        for model_id, expected in special_cases:
            result = usage_tracking.get_model_name(model_id)
            assert result == expected

    def test_get_model_name_fallback_extraction(self):
        """Test fallback extraction logic for unknown models."""
        # Test models that should trigger fallback extraction
        # Fallback extracts: model_id.split('.')[-1].split('-')[0]
        test_cases = [
            ("anthropic.claude-5-sonnet-v1:0", "Claude claude"),  # Extracts "claude" from "claude-5-sonnet-v1:0"
            ("anthropic.claude-10-v1:0", "Claude claude"),  # Extracts "claude" from "claude-10-v1:0"
            ("some.prefix.model-7-v1:0", "Claude model"),  # Extracts "model" from "model-7-v1:0"
        ]

        for model_id, expected in test_cases:
            result = usage_tracking.get_model_name(model_id)
            assert result == expected

    def test_get_model_name_no_dot_fallback(self):
        """Test fallback when model ID has no dot."""
        # Model IDs without dots should return "Claude Unknown"
        # BUT: if they match known patterns (like "claude-3-5-sonnet"), they still work
        test_cases = [
            ("claude-3-5-sonnet", "Claude 3.5 Sonnet"),  # Matches pattern even without dot
            ("claude-instant-v1", "Claude Unknown"),  # Claude 2 support removed
            ("anthropic-claude-3", "Claude Unknown"),  # No dot, doesn't match pattern
            ("just-text", "Claude Unknown"),  # No dot, doesn't match pattern
        ]

        for model_id, expected in test_cases:
            result = usage_tracking.get_model_name(model_id)
            assert result == expected

    def test_get_model_name_all_supported_models(self):
        """Test all officially supported model IDs."""
        # Note: Claude 2 support has been removed - only Claude 3+ models are supported
        supported_models = {
            "anthropic.claude-3-haiku-20240307-v1:0": "Claude 3 Haiku",
            "anthropic.claude-3-5-sonnet-20240620-v1:0": "Claude 3.5 Sonnet",
            "anthropic.claude-3-5-sonnet-20241022-v2:0": "Claude 3.5 Sonnet",
            "us.anthropic.claude-3-7-sonnet-20250219-v1:0": "Claude 3.7 Sonnet",
            "us.anthropic.claude-sonnet-4-20250514-v1:0": "Claude 4 Sonnet",
            "us.anthropic.claude-sonnet-4-5-20250929-v1:0": "Claude 4.5 Sonnet",
        }

        for model_id, expected_name in supported_models.items():
            result = usage_tracking.get_model_name(model_id)
            assert result == expected_name, f"Failed for {model_id}: got {result}, expected {expected_name}"