"""
Tests for structured JSON retrieval trace (#1347).
"""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from typing import Any, List, Optional
from unittest.mock import AsyncMock, patch

import pytest

from cortex.core.retrieval_pipeline import RetrievalPipeline
from cortex.core.retrieval_types import RetrievalConstraints, RetrievedItem

@dataclass
class MockConfig:
    retrieval_token_budget: int = 2000
    retrieval_rerank_enabled: bool = False
    retrieval_recency_weight: float = 0.0
    embedding_cache_enabled: bool = False
    retrieval_trace_enabled: bool = True
    retrieval_trace_log_level: str = "INFO"


class MockStorageAdapter:
    async def get_claim(self, claim_id: str) -> None:
        return None
    async def get_claims_sensitivity_batch(self, claim_ids: List[str]) -> dict:
        return {}


class TraceCaptureHandler(logging.Handler):
    def __init__(self):
        super().__init__()
        self.records = []

    def emit(self, record):
        self.records.append(record)


@pytest.fixture
def capture_trace():
    logger = logging.getLogger("cortex.retrieval.trace")
    logger.setLevel(logging.DEBUG)
    handler = TraceCaptureHandler()
    logger.addHandler(handler)
    yield handler
    logger.removeHandler(handler)


@pytest.mark.asyncio
async def test_retrieval_trace_emitted_when_enabled(capture_trace):
    config = MockConfig()
    storage = MockStorageAdapter()
    
    pipeline = RetrievalPipeline(storage, config)
    
    # Mock the underlying search to return a single item
    async def _mock_lightweight_retrieve(*args, **kwargs):
        return [
            RetrievedItem(
                item_type="claim",
                item_id="clm_123",
                content="test content",
                score=0.85,
                metadata={"category": "personal"},
            ),
            RetrievedItem(
                item_type="claim",
                item_id="clm_456",
                content="test content 2",
                score=0.45,
                metadata={"category": "work"},
            )
        ]
        
    pipeline._lightweight_retrieve = AsyncMock(side_effect=_mock_lightweight_retrieve)
    pipeline._apply_budget = AsyncMock(return_value=(50, await _mock_lightweight_retrieve()))
    pipeline._log_retrieval = AsyncMock()

    result = await pipeline.retrieve("test query")
    
    assert len(capture_trace.records) == 1
    record = capture_trace.records[0]
    assert record.levelname == "INFO"
    
    trace_data = json.loads(record.getMessage())
    assert trace_data["event"] == "retrieval_trace"
    assert trace_data["query"] == "test query"
    assert trace_data["mode"] == "lightweight"
    assert trace_data["final"]["count"] == 2
    assert trace_data["final"]["top_score"] == 0.85
    assert "clm_123" in trace_data["final"]["item_ids"]
    assert "personal" in trace_data["final"]["category_distribution"]
    assert "work" in trace_data["final"]["category_distribution"]
    assert "total_latency_ms" in trace_data


@pytest.mark.asyncio
async def test_retrieval_trace_not_emitted_when_disabled(capture_trace):
    config = MockConfig(retrieval_trace_enabled=False)
    storage = MockStorageAdapter()
    
    pipeline = RetrievalPipeline(storage, config)
    
    # Mock the underlying search
    pipeline._lightweight_retrieve = AsyncMock(return_value=[])
    pipeline._apply_budget = AsyncMock(return_value=(0, []))
    pipeline._log_retrieval = AsyncMock()

    await pipeline.retrieve("test query")
    
    assert len(capture_trace.records) == 0
