"""
Copyright 2024, Zep Software, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from concurrent.futures import ThreadPoolExecutor

from graphiti_core.llm_client.token_tracker import (
    PromptTokenUsage,
    TokenUsage,
    TokenUsageTracker,
)


class TestTokenUsage:
    def test_total_tokens(self):
        """Test that total_tokens correctly sums input and output tokens."""
        usage = TokenUsage(input_tokens=100, output_tokens=50)
        assert usage.total_tokens == 150

    def test_default_values(self):
        """Test that default values are zero."""
        usage = TokenUsage()
        assert usage.input_tokens == 0
        assert usage.output_tokens == 0
        assert usage.total_tokens == 0


class TestPromptTokenUsage:
    def test_total_tokens(self):
        """Test that total_tokens correctly sums input and output tokens."""
        usage = PromptTokenUsage(
            prompt_name='test',
            call_count=5,
            total_input_tokens=1000,
            total_output_tokens=500,
        )
        assert usage.total_tokens == 1500

    def test_avg_input_tokens(self):
        """Test average input tokens calculation."""
        usage = PromptTokenUsage(
            prompt_name='test',
            call_count=4,
            total_input_tokens=1000,
            total_output_tokens=500,
        )
        assert usage.avg_input_tokens == 250.0

    def test_avg_output_tokens(self):
        """Test average output tokens calculation."""
        usage = PromptTokenUsage(
            prompt_name='test',
            call_count=4,
            total_input_tokens=1000,
            total_output_tokens=500,
        )
        assert usage.avg_output_tokens == 125.0

    def test_avg_tokens_zero_calls(self):
        """Test that average returns 0 when call_count is zero."""
        usage = PromptTokenUsage(
            prompt_name='test',
            call_count=0,
            total_input_tokens=0,
            total_output_tokens=0,
        )
        assert usage.avg_input_tokens == 0
        assert usage.avg_output_tokens == 0


class TestTokenUsageTracker:
    def test_record_new_prompt(self):
        """Test recording usage for a new prompt."""
        tracker = TokenUsageTracker()
        tracker.record('extract_nodes', 100, 50)

        usage = tracker.get_usage()
        assert 'extract_nodes' in usage
        assert usage['extract_nodes'].call_count == 1
        assert usage['extract_nodes'].total_input_tokens == 100
        assert usage['extract_nodes'].total_output_tokens == 50

    def test_record_existing_prompt(self):
        """Test that multiple calls accumulate correctly."""
        tracker = TokenUsageTracker()
        tracker.record('extract_nodes', 100, 50)
        tracker.record('extract_nodes', 200, 100)

        usage = tracker.get_usage()
        assert usage['extract_nodes'].call_count == 2
        assert usage['extract_nodes'].total_input_tokens == 300
        assert usage['extract_nodes'].total_output_tokens == 150

    def test_record_none_prompt_name(self):
        """Test that None prompt_name is recorded as 'unknown'."""
        tracker = TokenUsageTracker()
        tracker.record(None, 100, 50)

        usage = tracker.get_usage()
        assert 'unknown' in usage
        assert usage['unknown'].call_count == 1

    def test_record_multiple_prompts(self):
        """Test recording usage for multiple different prompts."""
        tracker = TokenUsageTracker()
        tracker.record('extract_nodes', 100, 50)
        tracker.record('dedupe_nodes', 200, 100)
        tracker.record('extract_edges', 150, 75)

        usage = tracker.get_usage()
        assert len(usage) == 3
        assert 'extract_nodes' in usage
        assert 'dedupe_nodes' in usage
        assert 'extract_edges' in usage

    def test_get_usage_returns_copy(self):
        """Test that get_usage returns a copy, not the internal dict."""
        tracker = TokenUsageTracker()
        tracker.record('test', 100, 50)

        usage1 = tracker.get_usage()
        usage1['test'].total_input_tokens = 9999

        usage2 = tracker.get_usage()
        assert usage2['test'].total_input_tokens == 100  # Original unchanged

    def test_get_total_usage(self):
        """Test getting total usage across all prompts."""
        tracker = TokenUsageTracker()
        tracker.record('extract_nodes', 100, 50)
        tracker.record('dedupe_nodes', 200, 100)
        tracker.record('extract_edges', 150, 75)

        total = tracker.get_total_usage()
        assert total.input_tokens == 450
        assert total.output_tokens == 225
        assert total.total_tokens == 675

    def test_get_total_usage_empty(self):
        """Test getting total usage when no records exist."""
        tracker = TokenUsageTracker()
        total = tracker.get_total_usage()
        assert total.input_tokens == 0
        assert total.output_tokens == 0

    def test_reset(self):
        """Test that reset clears all tracked usage."""
        tracker = TokenUsageTracker()
        tracker.record('extract_nodes', 100, 50)
        tracker.record('dedupe_nodes', 200, 100)

        tracker.reset()

        usage = tracker.get_usage()
        assert len(usage) == 0

        total = tracker.get_total_usage()
        assert total.total_tokens == 0

    def test_thread_safety(self):
        """Test that concurrent access from multiple threads is safe."""
        tracker = TokenUsageTracker()
        num_threads = 10
        calls_per_thread = 100

        def record_tokens(thread_id):
            for _ in range(calls_per_thread):
                tracker.record(f'prompt_{thread_id}', 10, 5)

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(record_tokens, i) for i in range(num_threads)]
            for f in futures:
                f.result()

        usage = tracker.get_usage()
        assert len(usage) == num_threads

        total = tracker.get_total_usage()
        expected_input = num_threads * calls_per_thread * 10
        expected_output = num_threads * calls_per_thread * 5
        assert total.input_tokens == expected_input
        assert total.output_tokens == expected_output

    def test_concurrent_same_prompt(self):
        """Test concurrent access to the same prompt name."""
        tracker = TokenUsageTracker()
        num_threads = 10
        calls_per_thread = 100

        def record_tokens():
            for _ in range(calls_per_thread):
                tracker.record('shared_prompt', 10, 5)

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(record_tokens) for _ in range(num_threads)]
            for f in futures:
                f.result()

        usage = tracker.get_usage()
        assert usage['shared_prompt'].call_count == num_threads * calls_per_thread
        assert usage['shared_prompt'].total_input_tokens == num_threads * calls_per_thread * 10
        assert usage['shared_prompt'].total_output_tokens == num_threads * calls_per_thread * 5