"""
tests/test_model_router.py — Tests for Model Routing Policy (Issue #375).

Covers:
    - ModelTier enum values
    - Default operation-to-tier mapping
    - Tier selection per operation
    - Fallback chain (REASONING → MAIN → FAST)
    - Config override (custom operation→tier mapping)
    - Disabled mode (passthrough — zero behavior change)
    - ModelRouter as LLMProvider duck-type proxy
    - resolve_routed_model helper
    - Config parsing from TOML
    - All callsites annotated (no un-routed LLM calls)
    - Logging verification
"""

from __future__ import annotations

import logging
import pytest
from typing import Any
from unittest.mock import AsyncMock, MagicMock

from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    ExtractionResponse,
    LLMProvider,
)
from cortex.llm.router import (
    DEFAULT_OPERATION_MAP,
    ModelRouter,
    ModelRoutingConfig,
    ModelTier,
    get_router,
    resolve_routed_model,
)
from cortex.config import CortexConfig


# ===========================================================================
# Helpers / fixtures
# ===========================================================================

def _make_completion(text: str = "ok", model: str = "test-model") -> CompletionResponse:
    return CompletionResponse(
        text=text, model=model,
        tokens_in=5, tokens_out=3, cost_estimate=0.0, latency_ms=10.0,
    )


def _make_extraction(data: dict | None = None, model: str = "test-model") -> ExtractionResponse:
    return ExtractionResponse(
        data=data or {"result": "ok"}, model=model,
        tokens_in=5, tokens_out=3, cost_estimate=0.0, latency_ms=10.0,
    )


def _make_embed(model: str = "test-embed") -> EmbedResponse:
    return EmbedResponse(
        embeddings=[[0.1, 0.2]], model=model,
        tokens_in=4, tokens_out=0, cost_estimate=0.0, latency_ms=5.0,
    )


class MockProvider(LLMProvider):
    """Test double for LLMProvider."""

    DEFAULT_COMPLETION_MODEL = "mock-model"
    DEFAULT_EMBED_MODEL = "mock-embed"
    PROVIDER_NAME = "mock"

    def __init__(self):
        super().__init__()
        self.complete_calls = []
        self.extract_calls = []
        self.embed_calls = []

    async def complete(self, system, user, model=None, response_format="text",
                       max_tokens=None, temperature=0.0):
        self.complete_calls.append({
            "system": system, "user": user, "model": model,
            "response_format": response_format, "max_tokens": max_tokens,
            "temperature": temperature,
        })
        return _make_completion(model=model or self.DEFAULT_COMPLETION_MODEL)

    async def extract(self, system, user, schema, model=None,
                      max_tokens=None, temperature=0.0):
        self.extract_calls.append({
            "system": system, "user": user, "schema": schema, "model": model,
        })
        return _make_extraction(model=model or self.DEFAULT_COMPLETION_MODEL)

    async def embed(self, texts, model=None, input_type="query"):
        self.embed_calls.append({"texts": texts, "model": model, "input_type": input_type})
        return _make_embed(model=model or self.DEFAULT_EMBED_MODEL)


def _make_config(
    enabled: bool = False,
    fast: str = "",
    main: str = "",
    reasoning: str = "",
    operation_map: dict | None = None,
) -> MagicMock:
    """Build a mock CortexConfig with model routing fields."""
    cfg = MagicMock()
    cfg.model_routing_enabled = enabled
    cfg.model_routing_fast = fast
    cfg.model_routing_main = main
    cfg.model_routing_reasoning = reasoning
    cfg.model_routing_operation_map = operation_map or {}
    return cfg


# ===========================================================================
# ModelTier enum tests
# ===========================================================================

class TestModelTier:
    def test_tier_values(self):
        assert ModelTier.FAST.value == "fast"
        assert ModelTier.MAIN.value == "main"
        assert ModelTier.REASONING.value == "reasoning"

    def test_tier_from_string(self):
        assert ModelTier("fast") == ModelTier.FAST
        assert ModelTier("main") == ModelTier.MAIN
        assert ModelTier("reasoning") == ModelTier.REASONING


# ===========================================================================
# Default operation map tests
# ===========================================================================

class TestDefaultOperationMap:
    def test_fast_tier_operations(self):
        fast_ops = ["extraction", "classification", "routing", "summarization_light"]
        for op in fast_ops:
            assert DEFAULT_OPERATION_MAP[op] == ModelTier.FAST, f"{op} should be FAST"

    def test_main_tier_operations(self):
        main_ops = ["consolidation", "handoff_build", "summarization_deep",
                     "session_summary"]
        for op in main_ops:
            assert DEFAULT_OPERATION_MAP[op] == ModelTier.MAIN, f"{op} should be MAIN"

    def test_reasoning_tier_operations(self):
        reasoning_ops = ["contradiction_analysis", "reconciliation",
                         "entity_resolution", "promotion_scoring"]
        for op in reasoning_ops:
            assert DEFAULT_OPERATION_MAP[op] == ModelTier.REASONING, f"{op} should be REASONING"

    def test_all_required_operations_mapped(self):
        """Verify all spec-required operations are in the default map."""
        required = [
            "extraction", "classification", "routing", "summarization_light",
            "consolidation", "handoff_build", "summarization_deep",
            "contradiction_analysis", "reconciliation", "entity_resolution",
            "promotion_scoring",
        ]
        for op in required:
            assert op in DEFAULT_OPERATION_MAP, f"Missing required operation: {op}"


# ===========================================================================
# ModelRouter — disabled mode (passthrough)
# ===========================================================================

class TestModelRouterDisabled:
    """When routing is disabled, all calls pass through unchanged."""

    def test_disabled_by_default(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert not router.enabled

    def test_model_for_operation_returns_none_when_disabled(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert router.model_for_operation("extraction") is None

    @pytest.mark.asyncio
    async def test_complete_passthrough(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        resp = await router.complete(system="s", user="u", model="my-model")
        assert resp.model == "my-model"
        assert len(provider.complete_calls) == 1
        assert provider.complete_calls[0]["model"] == "my-model"

    @pytest.mark.asyncio
    async def test_extract_passthrough(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        resp = await router.extract(
            system="s", user="u", schema={"type": "object"}, model="my-model",
        )
        assert resp.model == "my-model"
        assert len(provider.extract_calls) == 1

    @pytest.mark.asyncio
    async def test_embed_passthrough(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        resp = await router.embed(texts=["hello"], model="embed-model")
        assert len(provider.embed_calls) == 1

    def test_usage_summary_proxied(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        summary = router.usage_summary()
        assert "calls" in summary

    def test_provider_names_proxied(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        names = router.provider_names()
        assert names == ["MockProvider"]  # no provider_names() on MockProvider


# ===========================================================================
# ModelRouter — enabled mode
# ===========================================================================

class TestModelRouterEnabled:

    def test_tier_for_operation(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        assert router.tier_for_operation("extraction") == ModelTier.FAST
        assert router.tier_for_operation("consolidation") == ModelTier.MAIN
        assert router.tier_for_operation("reconciliation") == ModelTier.REASONING

    def test_unknown_operation_defaults_to_fast(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano")
        router = ModelRouter(provider=provider, config=cfg)
        assert router.tier_for_operation("unknown_op") == ModelTier.FAST

    def test_model_for_operation(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        assert router.model_for_operation("extraction") == "nano"
        assert router.model_for_operation("consolidation") == "mini"
        assert router.model_for_operation("reconciliation") == "big"

    @pytest.mark.asyncio
    async def test_route_complete_selects_tier_model(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        resp = await router.route_complete(
            operation="extraction", system="s", user="u",
        )
        assert provider.complete_calls[0]["model"] == "nano"

    @pytest.mark.asyncio
    async def test_route_complete_explicit_model_overrides(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        resp = await router.route_complete(
            operation="extraction", system="s", user="u", model="override-model",
        )
        assert provider.complete_calls[0]["model"] == "override-model"

    @pytest.mark.asyncio
    async def test_route_extract_selects_tier_model(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        resp = await router.route_extract(
            operation="reconciliation", system="s", user="u",
            schema={"type": "object"},
        )
        assert provider.extract_calls[0]["model"] == "big"


# ===========================================================================
# Fallback chain
# ===========================================================================

class TestFallbackChain:
    """REASONING → MAIN → FAST when a tier's model is empty."""

    def test_reasoning_falls_back_to_main(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="")
        router = ModelRouter(provider=provider, config=cfg)
        model = router.model_for_tier(ModelTier.REASONING)
        assert model == "mini"

    def test_reasoning_falls_back_to_fast(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="", reasoning="")
        router = ModelRouter(provider=provider, config=cfg)
        model = router.model_for_tier(ModelTier.REASONING)
        assert model == "nano"

    def test_main_falls_back_to_fast(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        model = router.model_for_tier(ModelTier.MAIN)
        assert model == "nano"

    def test_all_empty_returns_none(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="", main="", reasoning="")
        router = ModelRouter(provider=provider, config=cfg)
        model = router.model_for_tier(ModelTier.FAST)
        assert model is None

    def test_fast_no_fallback_up(self):
        """FAST should not fall back UP to MAIN or REASONING."""
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        model = router.model_for_tier(ModelTier.FAST)
        assert model is None  # FAST has no fallback — only down


# ===========================================================================
# Config override (custom operation map)
# ===========================================================================

class TestConfigOverride:

    def test_custom_operation_map_override(self):
        provider = MockProvider()
        cfg = _make_config(
            enabled=True, fast="nano", main="mini", reasoning="big",
            operation_map={"extraction": "reasoning"},
        )
        router = ModelRouter(provider=provider, config=cfg)
        # extraction normally FAST, but overridden to REASONING
        assert router.tier_for_operation("extraction") == ModelTier.REASONING
        assert router.model_for_operation("extraction") == "big"

    def test_invalid_tier_in_override_ignored(self):
        provider = MockProvider()
        cfg = _make_config(
            enabled=True, fast="nano",
            operation_map={"extraction": "invalid_tier"},
        )
        router = ModelRouter(provider=provider, config=cfg)
        # Invalid tier ignored, default FAST remains
        assert router.tier_for_operation("extraction") == ModelTier.FAST

    def test_custom_operation_added(self):
        provider = MockProvider()
        cfg = _make_config(
            enabled=True, fast="nano", reasoning="big",
            operation_map={"my_custom_op": "reasoning"},
        )
        router = ModelRouter(provider=provider, config=cfg)
        assert router.tier_for_operation("my_custom_op") == ModelTier.REASONING
        assert router.model_for_operation("my_custom_op") == "big"


# ===========================================================================
# resolve_routed_model helper
# ===========================================================================

class TestResolveRoutedModel:

    def test_returns_fallback_when_no_router(self):
        provider = MockProvider()  # plain provider, not a router
        model = resolve_routed_model(provider, "extraction", "fallback-model")
        assert model == "fallback-model"

    def test_returns_fallback_when_disabled(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)  # disabled by default
        model = resolve_routed_model(router, "extraction", "fallback-model")
        assert model == "fallback-model"

    def test_returns_routed_model_when_enabled(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        model = resolve_routed_model(router, "extraction", "fallback-model")
        assert model == "nano"

    def test_returns_routed_model_for_reasoning(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        model = resolve_routed_model(router, "reconciliation", "fallback-model")
        assert model == "big"

    def test_returns_fallback_when_no_tier_model(self):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="", main="", reasoning="")
        router = ModelRouter(provider=provider, config=cfg)
        model = resolve_routed_model(router, "extraction", "fallback-model")
        assert model == "fallback-model"


# ===========================================================================
# get_router helper
# ===========================================================================

class TestModelRouterIsLLMProvider:
    """Audit fix D-HIGH-1: ModelRouter must be an instance of LLMProvider."""

    def test_isinstance_check(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert isinstance(router, LLMProvider), (
            "ModelRouter must extend LLMProvider — isinstance check failed"
        )

    def test_router_has_usage_log(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert hasattr(router, "_usage_log")

    def test_router_shares_provider_usage_log(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert router._usage_log is provider._usage_log


class TestEmptyStringModelRouting:
    """Audit fix D-HIGH-2: empty string model should still get tier-routed."""

    @pytest.mark.asyncio
    async def test_route_complete_empty_string_model_uses_tier(self):
        """Empty string model= should NOT override the tier-routed model."""
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        # Pass model="" — the tier should still be used, not ""
        await router.route_complete(
            operation="extraction", system="s", user="u", model="",
        )
        # _resolve_model treats "" as falsy → returns None → routing picks tier
        # But explicit_model="" is falsy, so routing kicks in
        # The provider should receive "nano" (FAST tier), not ""
        assert provider.complete_calls[0]["model"] == "nano"

    @pytest.mark.asyncio
    async def test_route_extract_empty_string_model_uses_tier(self):
        """Empty string model= should NOT override the tier-routed model."""
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        await router.route_extract(
            operation="reconciliation", system="s", user="u",
            schema={"type": "object"}, model="",
        )
        assert provider.extract_calls[0]["model"] == "big"


class TestGetRouter:

    def test_returns_router_from_model_router(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        assert get_router(router) is router

    def test_returns_none_from_plain_provider(self):
        provider = MockProvider()
        assert get_router(provider) is None

    def test_returns_router_from_attribute(self):
        provider = MockProvider()
        router = ModelRouter(provider=provider)
        wrapper = MagicMock()
        wrapper._router = router
        assert get_router(wrapper) is router


# ===========================================================================
# Config parsing
# ===========================================================================

class TestConfigParsing:

    def test_default_config_has_routing_disabled(self):
        cfg = CortexConfig()
        assert cfg.model_routing_enabled is False
        assert cfg.model_routing_fast == ""
        assert cfg.model_routing_main == ""
        assert cfg.model_routing_reasoning == ""
        assert cfg.model_routing_operation_map == {}

    def test_config_from_toml_parses_routing(self):
        overlay = {
            "cortex": {
                "model_routing": {
                    "enabled": True,
                    "tiers": {
                        "fast": "gpt-4.1-nano",
                        "main": "gpt-4.1-mini",
                        "reasoning": "gpt-4.1",
                    },
                    "operation_map": {
                        "extraction": "main",
                    },
                },
            },
        }
        cfg = CortexConfig.from_toml(overlay=overlay)
        assert cfg.model_routing_enabled is True
        assert cfg.model_routing_fast == "gpt-4.1-nano"
        assert cfg.model_routing_main == "gpt-4.1-mini"
        assert cfg.model_routing_reasoning == "gpt-4.1"
        assert cfg.model_routing_operation_map == {"extraction": "main"}

    def test_defaults_toml_has_routing_disabled(self):
        cfg = CortexConfig.from_toml()
        assert cfg.model_routing_enabled is False


# ===========================================================================
# Logging verification
# ===========================================================================

class TestLogging:

    def test_routing_logs_tier_selection(self, caplog):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="mini", reasoning="big")
        router = ModelRouter(provider=provider, config=cfg)
        with caplog.at_level(logging.DEBUG, logger="cortex.llm.router"):
            model = resolve_routed_model(router, "extraction", "fallback")
        assert any("extraction" in r.message and "nano" in r.message for r in caplog.records)

    def test_fallback_logs_tier_mismatch(self, caplog):
        provider = MockProvider()
        cfg = _make_config(enabled=True, fast="nano", main="", reasoning="")
        router = ModelRouter(provider=provider, config=cfg)
        with caplog.at_level(logging.INFO, logger="cortex.llm.router"):
            model = router.model_for_tier(ModelTier.REASONING)
        # Should log that REASONING fell back to FAST
        assert any("falling back" in r.message.lower() for r in caplog.records)


# ===========================================================================
# Callsite annotation completeness
# ===========================================================================

class TestCallsiteAnnotations:
    """Verify that all known LLM callsites import and use resolve_routed_model."""

    EXPECTED_CALLSITES = {
        "cortex/extraction/extractor.py": "extraction",
        "cortex/core/reconciliation.py": "reconciliation",
        "cortex/core/entity_resolution.py": "entity_resolution",
        "cortex/core/agentic_retrieval.py": "agentic_refinement",
        "cortex/dialectic/engine.py": "dialectic_synthesis",
        "cortex/dialectic/planner.py": "query_planning",
        "cortex/sleep/s1_consolidator.py": "consolidation",
        "cortex/circadian/sleep_consolidator.py": "session_summary",
        "cortex/circadian/curiosity.py": "curiosity_seeds",
        "cortex/circadian/engine.py": "nap_summary",
        "cortex/identity/drift_calculator.py": "key_phrase_extraction",
        "cortex/identity/cornerstone.py": "cornerstone_evolution",
        "cortex/brain_graph/builder.py": "brain_graph_nodes",
        "cortex/deprecation/pipeline.py": "deprecation_review",
    }

    def test_all_callsites_annotated(self):
        """Check that each expected file imports resolve_routed_model."""
        import pathlib
        repo_root = pathlib.Path(__file__).parent.parent
        for rel_path, operation in self.EXPECTED_CALLSITES.items():
            full_path = repo_root / rel_path
            if not full_path.exists():
                continue  # skip if file doesn't exist (optional modules)
            content = full_path.read_text()
            assert "resolve_routed_model" in content, (
                f"{rel_path} does not use resolve_routed_model — "
                f"expected operation annotation: {operation}"
            )

    def test_all_annotated_operations_in_default_map(self):
        """Every operation name used in callsites should be in DEFAULT_OPERATION_MAP."""
        import pathlib
        import re
        repo_root = pathlib.Path(__file__).parent.parent
        # Grep all operation strings from resolve_routed_model calls
        pattern = re.compile(r'resolve_routed_model\([^,]+,\s*"([^"]+)"')
        found_operations = set()
        for py_file in (repo_root / "cortex").rglob("*.py"):
            if "__pycache__" in str(py_file):
                continue
            content = py_file.read_text()
            found_operations.update(pattern.findall(content))
        # Remove the self-referential example in router.py docstring
        for op in found_operations:
            assert op in DEFAULT_OPERATION_MAP, (
                f"Operation '{op}' used in callsite but missing from DEFAULT_OPERATION_MAP"
            )
