"""tests/test_llm_usage_tracker.py — Tests for LLM usage tracking (Issue #419)."""
from __future__ import annotations
import time
from datetime import datetime, timezone, timedelta
from unittest.mock import AsyncMock, MagicMock
import pytest
from cortex.llm.usage_tracker import LLMUsageTracker, TierStats, _percentile
from cortex.llm.provider import CompletionResponse, ExtractionResponse, LLMProvider
from cortex.llm.router import ModelRouter

@pytest.fixture(autouse=True)
def reset_tracker():
    LLMUsageTracker.reset_instances()
    yield
    LLMUsageTracker.reset_instances()

def _t(owner="test"): return LLMUsageTracker.get(owner)
def _cr(ti=10, to=5): return CompletionResponse(text="ok", model="m", tokens_in=ti, tokens_out=to, cost_estimate=0.0, latency_ms=50.0)
def _er(ti=8, to=4): return ExtractionResponse(data={}, model="m", tokens_in=ti, tokens_out=to, cost_estimate=0.0, latency_ms=30.0)

class MockProvider(LLMProvider):
    DEFAULT_COMPLETION_MODEL = "mock"; DEFAULT_EMBED_MODEL = "me"; PROVIDER_NAME = "mock"
    def __init__(self):
        super().__init__(); self._cr = _cr(); self._er = _er()
    async def complete(self, system, user, model=None, **kw): return self._cr
    async def extract(self, system, user, schema, model=None, **kw): return self._er
    async def embed(self, texts, model=None, input_type="query"):
        from cortex.llm.provider import EmbedResponse
        return EmbedResponse(embeddings=[[0.1]], model="m", tokens_in=1, tokens_out=0, cost_estimate=0.0, latency_ms=1.0)

class TestTierStats:
    def test_empty(self):
        ts = TierStats()
        assert ts.call_count == 0 and ts.avg_latency_ms() == 0.0
    def test_dict_keys(self):
        ts = TierStats(); ts._latencies = [10.0]
        d = ts.to_dict()
        for k in ("call_count","total_input_tokens","total_output_tokens","avg_latency_ms","p50_latency_ms","p95_latency_ms","error_count"):
            assert k in d
    def test_avg(self):
        ts = TierStats(); ts._latencies = [10.0, 20.0, 30.0]
        assert ts.avg_latency_ms() == 20.0
    def test_percentiles(self):
        ts = TierStats(); ts._latencies = list(range(1, 101))
        assert ts.p50_latency_ms() == 50.0 and ts.p95_latency_ms() == 95.0

class TestBasic:
    def test_singleton(self):
        assert LLMUsageTracker.get("a") is LLMUsageTracker.get("a")
        assert LLMUsageTracker.get("a") is not LLMUsageTracker.get("b")
    def test_empty_all_tiers(self):
        s = _t().get_summary()
        for t in ("fast","main","reasoning"):
            assert t in s["tiers"] and s["tiers"][t]["call_count"] == 0
    def test_empty_total_zeros(self):
        s = _t().get_summary()
        assert s["total"]["call_count"] == 0
    def test_record_single(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=42.0)
        s = t.get_summary()
        assert s["tiers"]["fast"]["call_count"] == 1
        assert s["tiers"]["fast"]["total_input_tokens"] == 10
        assert s["tiers"]["fast"]["avg_latency_ms"] == 42.0
    def test_record_multiple(self):
        t = _t()
        t.record_call(tier="main", operation="consolidation", input_tokens=100, output_tokens=50, latency_ms=100.0)
        t.record_call(tier="main", operation="session_summary", input_tokens=200, output_tokens=80, latency_ms=200.0)
        s = t.get_summary()
        assert s["tiers"]["main"]["call_count"] == 2 and s["tiers"]["main"]["total_input_tokens"] == 300
    def test_error_count(self):
        t = _t()
        t.record_call(tier="reasoning", operation="reconciliation", latency_ms=50.0, error=True)
        s = t.get_summary()
        assert s["tiers"]["reasoning"]["error_count"] == 1
    def test_totals(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        t.record_call(tier="main", operation="consolidation", input_tokens=20, output_tokens=10, latency_ms=20.0)
        t.record_call(tier="reasoning", operation="reconciliation", input_tokens=30, output_tokens=15, latency_ms=30.0, error=True)
        s = t.get_summary()
        assert s["total"]["call_count"] == 3 and s["total"]["total_input_tokens"] == 60 and s["total"]["error_count"] == 1
    def test_by_operation(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=5, output_tokens=2, latency_ms=10.0)
        t.record_call(tier="fast", operation="extraction", input_tokens=5, output_tokens=2, latency_ms=20.0)
        s = t.get_summary()
        assert s["by_operation"]["extraction"]["fast"]["call_count"] == 2
    def test_reset(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        t.reset()
        assert t.get_summary()["total"]["call_count"] == 0

class TestFiltering:
    def test_tier_filter(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        t.record_call(tier="main", operation="consolidation", input_tokens=20, output_tokens=10, latency_ms=20.0)
        s = t.get_summary(tier_filter="fast")
        for op_name, tier_map in s["by_operation"].items():
            for tn in tier_map: assert tn == "fast"
    def test_since_excludes_old(self):
        t = _t()
        from cortex.llm.usage_tracker import _CallRecord
        old = _CallRecord(tier="fast", operation="extraction", input_tokens=999, output_tokens=999, latency_ms=999.0, error=False, timestamp=time.time()-3600)
        with t._lock: t._records.append(old)
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        s = t.get_summary(since=datetime.now(timezone.utc) - timedelta(minutes=30))
        assert sum(td.get("call_count", 0) for td in s["tiers"].values()) == 1
    def test_since_future_empty(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        s = t.get_summary(since=datetime.now(timezone.utc) + timedelta(hours=1))
        assert s["total"]["call_count"] == 0
    def test_since_and_tier(self):
        t = _t()
        t.record_call(tier="fast", operation="extraction", input_tokens=10, output_tokens=5, latency_ms=10.0)
        t.record_call(tier="main", operation="consolidation", input_tokens=20, output_tokens=10, latency_ms=20.0)
        s = t.get_summary(since=datetime.now(timezone.utc) - timedelta(minutes=5), tier_filter="main")
        for op_name, tier_map in s["by_operation"].items():
            for tn in tier_map: assert tn == "main"

class TestRouterIntegration:
    def _router(self, enabled=True):
        provider = MockProvider()
        config = MagicMock()
        config.model_routing_enabled = enabled
        config.model_routing_fast = "fast-model"; config.model_routing_main = "main-model"
        config.model_routing_reasoning = "reasoning-model"; config.model_routing_operation_map = {}
        config.owner_id = "router-test"
        return ModelRouter(provider=provider, config=config)

    @pytest.mark.asyncio
    async def test_route_complete_records_call(self):
        r = self._router()
        assert r._usage_tracker is not None
        await r.route_complete(operation="extraction", system="s", user="u")
        s = r._usage_tracker.get_summary()
        assert s["total"]["call_count"] == 1 and s["tiers"]["fast"]["call_count"] == 1
        assert "extraction" in s["by_operation"]

    @pytest.mark.asyncio
    async def test_route_complete_records_tokens(self):
        r = self._router(); r._provider._cr = _cr(ti=42, to=17)
        await r.route_complete(operation="extraction", system="s", user="u")
        s = r._usage_tracker.get_summary()
        assert s["tiers"]["fast"]["total_input_tokens"] == 42

    @pytest.mark.asyncio
    async def test_route_extract_records_call(self):
        r = self._router()
        await r.route_extract(operation="reconciliation", system="s", user="u", schema={})
        s = r._usage_tracker.get_summary()
        assert s["tiers"]["reasoning"]["call_count"] == 1

    @pytest.mark.asyncio
    async def test_route_complete_records_error(self):
        r = self._router(); r._provider.complete = AsyncMock(side_effect=RuntimeError("boom"))
        with pytest.raises(RuntimeError):
            await r.route_complete(operation="extraction", system="s", user="u")
        s = r._usage_tracker.get_summary()
        assert s["tiers"]["fast"]["error_count"] == 1

    @pytest.mark.asyncio
    async def test_routing_disabled_records_fast(self):
        r = self._router(enabled=False)
        await r.route_complete(operation="extraction", system="s", user="u")
        s = r._usage_tracker.get_summary()
        assert s["total"]["call_count"] == 1 and s["tiers"]["fast"]["call_count"] == 1

    @pytest.mark.asyncio
    async def test_multiple_ops_tracked(self):
        r = self._router()
        await r.route_complete(operation="extraction", system="s", user="u")
        await r.route_complete(operation="consolidation", system="s", user="u")
        await r.route_complete(operation="reconciliation", system="s", user="u")
        s = r._usage_tracker.get_summary()
        assert s["total"]["call_count"] == 3
        for op in ("extraction","consolidation","reconciliation"):
            assert op in s["by_operation"]

class TestPercentile:
    def test_empty(self): assert _percentile([], 0.95) == 0.0
    def test_single(self): assert _percentile([42.0], 0.5) == 42.0
    def test_p50(self): assert _percentile(sorted([10.,20.,30.,40.,50.]), 0.50) == 30.0
    def test_p95(self): assert _percentile(sorted(float(i) for i in range(1, 101)), 0.95) == 95.0
