"""Tests for token usage tracking in analyzer.

Tests verify that PaperAnalyzer correctly extracts token usage from
Pydantic AI AgentRunResult and reports it in AnalysisResult.
"""

from unittest.mock import MagicMock, patch

from pydantic_ai import AgentRunResult
from pydantic_ai.models.test import TestModel
import pytest

from information_composer.pubmed.analyzer.agents import (
    analyze_domain,
    analyze_summary,
)
from information_composer.pubmed.analyzer.analyzer import PaperAnalyzer
from information_composer.pubmed.analyzer.dependencies import AnalysisDeps
from information_composer.pubmed.analyzer.models import (
    AnalysisConfig,
    DomainResult,
    PaperInput,
    SummaryResult,
)


@pytest.fixture
def analysis_config() -> AnalysisConfig:
    """Create test analysis config with Pydantic AI enabled."""
    return AnalysisConfig(
        analysis_types=["summary", "domain"],
        llm_provider="openai",
        llm_model="gpt-4o-mini",
        domain_list=["Test Domain 1", "Test Domain 2"],
        use_pydantic_ai=True,
    )


@pytest.fixture
def paper_input() -> PaperInput:
    """Create test paper input."""
    return PaperInput(
        pmid="12345678",
        title="Test Paper Title",
        abstract="This is a test abstract for testing purposes.",
    )


@pytest.fixture
def mock_summary_result() -> SummaryResult:
    """Create a mock SummaryResult."""
    return SummaryResult(
        main_findings=["Finding 1", "Finding 2"],
        innovations=["Innovation 1"],
        conclusions="Test conclusions",
    )


@pytest.fixture
def mock_domain_result() -> DomainResult:
    """Create a mock DomainResult."""
    return DomainResult(
        relevant_domains=["Test Domain 1"],
        domain_scores={"Test Domain 1": 0.9},
        primary_domain="Test Domain 1",
        reasoning="Test reasoning",
    )


@pytest.fixture
def mock_agent_result_with_tokens(
    mock_summary_result: SummaryResult,
) -> MagicMock:
    """Create a mock AgentRunResult with token usage."""
    mock_result = MagicMock(spec=AgentRunResult)
    mock_result.output = mock_summary_result
    mock_usage = MagicMock()
    mock_usage.total_tokens = 150
    mock_result.usage.return_value = mock_usage
    return mock_result


@pytest.fixture
def mock_agent_result_no_tokens(
    mock_summary_result: SummaryResult,
) -> MagicMock:
    """Create a mock AgentRunResult with no token usage (e.g., TestModel)."""
    mock_result = MagicMock(spec=AgentRunResult)
    mock_result.output = mock_summary_result
    mock_result.usage.return_value = None
    return mock_result


class TestAnalyzeSummaryTokenUsage:
    """Tests for token usage tracking in _analyze_summary_pydantic."""

    @pytest.mark.asyncio
    async def test_analyze_summary_token_usage_with_mock(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
        mock_agent_result_with_tokens: MagicMock,
    ) -> None:
        """Test that token usage is extracted from AgentRunResult."""
        analyzer = PaperAnalyzer(analysis_config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_summary",
            return_value=mock_agent_result_with_tokens,
        ):
            result, tokens = await analyzer._analyze_summary_pydantic(paper_input)

        assert isinstance(result, SummaryResult)
        assert tokens == 150

    @pytest.mark.asyncio
    async def test_analyze_summary_token_usage_with_test_model(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
    ) -> None:
        """Test that TestModel returns token usage info."""
        import asyncio

        from information_composer.pubmed.analyzer import agents
        from information_composer.pubmed.analyzer.cache import AnalysisCache

        deps = AnalysisDeps(
            config=analysis_config,
            cache=AnalysisCache(enabled=False),
            semaphore=asyncio.Semaphore(1),
        )

        test_model = TestModel()

        # Call agent directly with TestModel
        agent_result = await agents.analyze_summary(paper_input, deps, model=test_model)

        # TestModel provides usage info with total_tokens
        assert agent_result.usage() is not None
        assert agent_result.usage().total_tokens > 0


class TestAnalyzeDomainTokenUsage:
    """Tests for token usage tracking in _analyze_domain_pydantic."""

    @pytest.mark.asyncio
    async def test_analyze_domain_token_usage_with_mock(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
        mock_domain_result: DomainResult,
    ) -> None:
        """Test that token usage is extracted from AgentRunResult."""
        mock_result = MagicMock(spec=AgentRunResult)
        mock_result.output = mock_domain_result
        mock_usage = MagicMock()
        mock_usage.total_tokens = 200
        mock_result.usage.return_value = mock_usage

        analyzer = PaperAnalyzer(analysis_config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_domain",
            return_value=mock_result,
        ):
            result, tokens = await analyzer._analyze_domain_pydantic(paper_input)

        assert isinstance(result, DomainResult)
        assert tokens == 200

    @pytest.mark.asyncio
    async def test_analyze_domain_token_usage_zero_when_no_usage(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
        mock_domain_result: DomainResult,
    ) -> None:
        """Test that tokens=0 when usage() returns None."""
        mock_result = MagicMock(spec=AgentRunResult)
        mock_result.output = mock_domain_result
        mock_result.usage.return_value = None

        analyzer = PaperAnalyzer(analysis_config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_domain",
            return_value=mock_result,
        ):
            result, tokens = await analyzer._analyze_domain_pydantic(paper_input)

        assert isinstance(result, DomainResult)
        assert tokens == 0


class TestAnalyzePaperTokenUsage:
    """Tests for token usage in full analyze_paper flow."""

    @pytest.mark.asyncio
    async def test_analyze_paper_sums_tokens_from_both_analyses(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
    ) -> None:
        """Test that analyze_paper() correctly sums tokens from both analyses."""
        mock_summary = MagicMock(spec=AgentRunResult)
        mock_summary.output = SummaryResult(
            main_findings=["Finding 1"],
            innovations=["Innovation 1"],
            conclusions="Test",
        )
        mock_summary_usage = MagicMock()
        mock_summary_usage.total_tokens = 100
        mock_summary.usage.return_value = mock_summary_usage

        mock_domain = MagicMock(spec=AgentRunResult)
        mock_domain.output = DomainResult(
            relevant_domains=["Domain 1"],
            domain_scores={"Domain 1": 0.8},
            primary_domain="Domain 1",
            reasoning="Test",
        )
        mock_domain_usage = MagicMock()
        mock_domain_usage.total_tokens = 150
        mock_domain.usage.return_value = mock_domain_usage

        analyzer = PaperAnalyzer(analysis_config)

        with (
            patch(
                "information_composer.pubmed.analyzer.analyzer.agents.analyze_summary",
                return_value=mock_summary,
            ),
            patch(
                "information_composer.pubmed.analyzer.analyzer.agents.analyze_domain",
                return_value=mock_domain,
            ),
        ):
            result = await analyzer.analyze_paper(paper_input)

        assert result.processing_metadata.tokens_used == 250  # 100 + 150

    @pytest.mark.asyncio
    async def test_analyze_paper_token_usage_in_metadata(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
    ) -> None:
        """Test that tokens_used appears in ProcessingMetadata."""
        mock_summary = MagicMock(spec=AgentRunResult)
        mock_summary.output = SummaryResult(
            main_findings=["Finding 1"],
            innovations=["Innovation 1"],
            conclusions="Test",
        )
        mock_summary_usage = MagicMock()
        mock_summary_usage.total_tokens = 75
        mock_summary.usage.return_value = mock_summary_usage

        mock_domain = MagicMock(spec=AgentRunResult)
        mock_domain.output = DomainResult(
            relevant_domains=["Domain 1"],
            domain_scores={"Domain 1": 0.9},
            primary_domain="Domain 1",
            reasoning="Test reasoning",
        )
        mock_domain_usage = MagicMock()
        mock_domain_usage.total_tokens = 125
        mock_domain.usage.return_value = mock_domain_usage

        analyzer = PaperAnalyzer(analysis_config)

        with (
            patch(
                "information_composer.pubmed.analyzer.analyzer.agents.analyze_summary",
                return_value=mock_summary,
            ),
            patch(
                "information_composer.pubmed.analyzer.analyzer.agents.analyze_domain",
                return_value=mock_domain,
            ),
        ):
            result = await analyzer.analyze_paper(paper_input)

        assert result.processing_metadata.tokens_used is not None
        assert result.processing_metadata.tokens_used == 200  # 75 + 125


class TestTokenUsageEdgeCases:
    """Tests for edge cases in token usage tracking."""

    @pytest.mark.asyncio
    async def test_token_usage_handles_missing_usage_attribute(
        self,
        analysis_config: AnalysisConfig,
        paper_input: PaperInput,
        mock_summary_result: SummaryResult,
    ) -> None:
        """Test graceful handling when result has no usage attribute."""
        mock_result = MagicMock()
        mock_result.output = mock_summary_result
        # No usage attribute at all
        del mock_result.usage

        analyzer = PaperAnalyzer(analysis_config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_summary",
            return_value=mock_result,
        ):
            result, tokens = await analyzer._analyze_summary_pydantic(paper_input)

        assert isinstance(result, SummaryResult)
        assert tokens == 0  # Should default to 0

    @pytest.mark.asyncio
    async def test_token_usage_with_only_summary_analysis(
        self,
        paper_input: PaperInput,
    ) -> None:
        """Test token tracking when only summary analysis is requested."""
        config = AnalysisConfig(
            analysis_types=["summary"],
            llm_provider="openai",
            llm_model="gpt-4o-mini",
            use_pydantic_ai=True,
        )

        mock_summary = MagicMock(spec=AgentRunResult)
        mock_summary.output = SummaryResult(
            main_findings=["Finding 1"],
            innovations=["Innovation 1"],
            conclusions="Test",
        )
        mock_summary_usage = MagicMock()
        mock_summary_usage.total_tokens = 100
        mock_summary.usage.return_value = mock_summary_usage

        analyzer = PaperAnalyzer(config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_summary",
            return_value=mock_summary,
        ):
            result = await analyzer.analyze_paper(paper_input)

        assert result.processing_metadata.tokens_used == 100

    @pytest.mark.asyncio
    async def test_token_usage_with_only_domain_analysis(
        self,
        paper_input: PaperInput,
    ) -> None:
        """Test token tracking when only domain analysis is requested."""
        config = AnalysisConfig(
            analysis_types=["domain"],
            llm_provider="openai",
            llm_model="gpt-4o-mini",
            domain_list=["Domain 1", "Domain 2"],
            use_pydantic_ai=True,
        )

        mock_domain = MagicMock(spec=AgentRunResult)
        mock_domain.output = DomainResult(
            relevant_domains=["Domain 1"],
            domain_scores={"Domain 1": 0.9},
            primary_domain="Domain 1",
            reasoning="Test reasoning",
        )
        mock_domain_usage = MagicMock()
        mock_domain_usage.total_tokens = 150
        mock_domain.usage.return_value = mock_domain_usage

        analyzer = PaperAnalyzer(config)

        with patch(
            "information_composer.pubmed.analyzer.analyzer.agents.analyze_domain",
            return_value=mock_domain,
        ):
            result = await analyzer.analyze_paper(paper_input)

        assert result.processing_metadata.tokens_used == 150
