"""
tests/test_llm_cascade.py — Tests for configurable LLM provider cascade (#105).

Covers:
    - FallbackLLMProvider: cascade ordering, fallback on failure, all-fail raises
    - build_provider_cascade: config-driven instantiation, unknown names skipped
    - CortexConfig: llm_provider_priority field + TOML parsing
    - CortexPlugin: uses config cascade, not hardcoded provider order

All tests use mocks — no real API calls, no API keys required.
"""
from __future__ import annotations

import pytest
import pytest_asyncio
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    ExtractionResponse,
    LLMProvider,
    UsageRecord,
)
from cortex.llm.fallback import FallbackLLMProvider, build_provider_cascade
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 that returns canned responses."""

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

    def __init__(
        self,
        name: str = "mock",
        fail_complete: bool = False,
        fail_extract: bool = False,
        fail_embed: bool = False,
    ) -> None:
        super().__init__()
        self.PROVIDER_NAME = name
        self._fail_complete = fail_complete
        self._fail_extract = fail_extract
        self._fail_embed = fail_embed
        self.complete_call_count = 0
        self.extract_call_count = 0
        self.embed_call_count = 0

    async def complete(self, system, user, model=None, response_format="text",
                       max_tokens=None, temperature=0.0) -> CompletionResponse:
        self.complete_call_count += 1
        if self._fail_complete:
            raise RuntimeError(f"{self.PROVIDER_NAME}: complete() failed")
        rec = UsageRecord(
            provider=self.PROVIDER_NAME, model=self.DEFAULT_COMPLETION_MODEL,
            call_type="complete", tokens_in=5, tokens_out=3,
            cost_estimate=0.001, latency_ms=10.0,
        )
        self._log_usage(rec)
        return _make_completion(model=self.DEFAULT_COMPLETION_MODEL)

    async def extract(self, system, user, schema, model=None,
                      max_tokens=None, temperature=0.0) -> ExtractionResponse:
        self.extract_call_count += 1
        if self._fail_extract:
            raise RuntimeError(f"{self.PROVIDER_NAME}: extract() failed")
        rec = UsageRecord(
            provider=self.PROVIDER_NAME, model=self.DEFAULT_COMPLETION_MODEL,
            call_type="extract", tokens_in=5, tokens_out=3,
            cost_estimate=0.001, latency_ms=10.0,
        )
        self._log_usage(rec)
        return _make_extraction(model=self.DEFAULT_COMPLETION_MODEL)

    async def embed(self, texts, model=None, input_type="query") -> EmbedResponse:
        self.embed_call_count += 1
        if self._fail_embed:
            raise RuntimeError(f"{self.PROVIDER_NAME}: embed() failed")
        rec = UsageRecord(
            provider=self.PROVIDER_NAME, model=self.DEFAULT_EMBED_MODEL,
            call_type="embed", tokens_in=4, tokens_out=0,
            cost_estimate=0.0, latency_ms=5.0,
        )
        self._log_usage(rec)
        return _make_embed(model=self.DEFAULT_EMBED_MODEL)


# ===========================================================================
# FallbackLLMProvider — basic cascade
# ===========================================================================

class TestFallbackLLMProvider:

    def _make_cascade(self, *providers) -> FallbackLLMProvider:
        return FallbackLLMProvider(list(providers))

    def test_requires_at_least_one_provider(self):
        with pytest.raises(ValueError, match="at least one provider"):
            FallbackLLMProvider([])

    def test_provider_names_reflects_cascade_order(self):
        p1 = MockProvider("alpha")
        p2 = MockProvider("beta")
        cascade = self._make_cascade(p1, p2)
        assert cascade.provider_names() == ["alpha", "beta"]

    def test_active_provider_is_none_initially(self):
        cascade = self._make_cascade(MockProvider())
        assert cascade.active_provider is None

    # --- complete() cascade -------------------------------------------------

    @pytest.mark.asyncio
    async def test_complete_uses_first_provider_when_healthy(self):
        p1 = MockProvider("primary")
        p2 = MockProvider("secondary")
        cascade = self._make_cascade(p1, p2)

        result = await cascade.complete("sys", "user")

        assert isinstance(result, CompletionResponse)
        assert p1.complete_call_count == 1
        assert p2.complete_call_count == 0
        assert cascade.active_provider is p1

    @pytest.mark.asyncio
    async def test_complete_falls_back_on_first_failure(self):
        p1 = MockProvider("failing", fail_complete=True)
        p2 = MockProvider("fallback")
        cascade = self._make_cascade(p1, p2)

        result = await cascade.complete("sys", "user")

        assert isinstance(result, CompletionResponse)
        assert p1.complete_call_count == 1
        assert p2.complete_call_count == 1
        assert cascade.active_provider is p2

    @pytest.mark.asyncio
    async def test_complete_falls_back_through_multiple_failures(self):
        p1 = MockProvider("fail1", fail_complete=True)
        p2 = MockProvider("fail2", fail_complete=True)
        p3 = MockProvider("success")
        cascade = self._make_cascade(p1, p2, p3)

        result = await cascade.complete("sys", "user")

        assert isinstance(result, CompletionResponse)
        assert p1.complete_call_count == 1
        assert p2.complete_call_count == 1
        assert p3.complete_call_count == 1
        assert cascade.active_provider is p3

    @pytest.mark.asyncio
    async def test_complete_raises_when_all_providers_fail(self):
        p1 = MockProvider("fail1", fail_complete=True)
        p2 = MockProvider("fail2", fail_complete=True)
        cascade = self._make_cascade(p1, p2)

        with pytest.raises(RuntimeError, match="fail2"):
            await cascade.complete("sys", "user")

    # --- extract() cascade --------------------------------------------------

    @pytest.mark.asyncio
    async def test_extract_uses_first_provider_when_healthy(self):
        p1 = MockProvider("primary")
        p2 = MockProvider("secondary")
        cascade = self._make_cascade(p1, p2)

        schema = {"type": "object"}
        result = await cascade.extract("sys", "user", schema)

        assert isinstance(result, ExtractionResponse)
        assert p1.extract_call_count == 1
        assert p2.extract_call_count == 0

    @pytest.mark.asyncio
    async def test_extract_falls_back_on_failure(self):
        p1 = MockProvider("failing", fail_extract=True)
        p2 = MockProvider("fallback")
        cascade = self._make_cascade(p1, p2)

        result = await cascade.extract("sys", "user", {"type": "object"})

        assert isinstance(result, ExtractionResponse)
        assert p2.extract_call_count == 1

    @pytest.mark.asyncio
    async def test_extract_raises_when_all_fail(self):
        p1 = MockProvider("f1", fail_extract=True)
        p2 = MockProvider("f2", fail_extract=True)
        cascade = self._make_cascade(p1, p2)

        with pytest.raises(RuntimeError):
            await cascade.extract("sys", "user", {})

    # --- embed() cascade ----------------------------------------------------

    @pytest.mark.asyncio
    async def test_embed_uses_first_provider_when_healthy(self):
        p1 = MockProvider("primary")
        cascade = self._make_cascade(p1)

        result = await cascade.embed(["hello"])

        assert isinstance(result, EmbedResponse)
        assert p1.embed_call_count == 1

    @pytest.mark.asyncio
    async def test_embed_falls_back_on_failure(self):
        p1 = MockProvider("failing", fail_embed=True)
        p2 = MockProvider("fallback")
        cascade = self._make_cascade(p1, p2)

        result = await cascade.embed(["hello"])

        assert isinstance(result, EmbedResponse)
        assert p1.embed_call_count == 1
        assert p2.embed_call_count == 1
        assert cascade.active_provider is p2

    @pytest.mark.asyncio
    async def test_embed_raises_when_all_fail(self):
        p1 = MockProvider("f1", fail_embed=True)
        p2 = MockProvider("f2", fail_embed=True)
        cascade = self._make_cascade(p1, p2)

        with pytest.raises(RuntimeError):
            await cascade.embed(["hello"])

    # --- usage tracking -----------------------------------------------------

    @pytest.mark.asyncio
    async def test_usage_summary_reflects_calls(self):
        p1 = MockProvider("primary")
        cascade = self._make_cascade(p1)

        await cascade.complete("sys", "user")
        await cascade.embed(["text"])
        summary = cascade.usage_summary()

        assert summary["calls"] == 2
        assert summary["active_provider"] == "primary"
        assert "primary" in summary["cascade"]

    @pytest.mark.asyncio
    async def test_usage_summary_active_provider_none_before_calls(self):
        cascade = self._make_cascade(MockProvider())
        summary = cascade.usage_summary()
        assert summary["active_provider"] is None

    @pytest.mark.asyncio
    async def test_usage_summary_after_fallback_shows_correct_active(self):
        p1 = MockProvider("fail", fail_complete=True)
        p2 = MockProvider("backup")
        cascade = self._make_cascade(p1, p2)

        await cascade.complete("sys", "user")
        summary = cascade.usage_summary()
        assert summary["active_provider"] == "backup"

    # --- cascade order is preserved -----------------------------------------

    def test_provider_order_preserved(self):
        names = ["alpha", "beta", "gamma"]
        providers = [MockProvider(n) for n in names]
        cascade = FallbackLLMProvider(providers)
        assert cascade.provider_names() == names


# ===========================================================================
# build_provider_cascade — config-driven factory
# ===========================================================================

class TestBuildProviderCascade:

    def _mock_all_providers(self):
        """Patch all three provider classes to return MockProvider instances."""
        mock_anthropic = MockProvider("anthropic")
        mock_openai = MockProvider("openai")
        mock_ollama = MockProvider("ollama")

        def make_patch(instance):
            cls = MagicMock(return_value=instance)
            return cls

        patches = {
            "cortex.llm.anthropic_provider.AnthropicProvider": make_patch(mock_anthropic),
            "cortex.llm.openai_provider.OpenAIProvider": make_patch(mock_openai),
            "cortex.llm.ollama_provider.OllamaProvider": make_patch(mock_ollama),
        }
        return patches, mock_anthropic, mock_openai, mock_ollama

    def test_raises_on_empty_priority(self):
        with pytest.raises(RuntimeError, match="empty"):
            build_provider_cascade([])

    def test_unknown_provider_names_skipped(self):
        # Only known providers that can init go in; unknown are skipped.
        # Mock anthropic to succeed so we get at least one.
        mock_instance = MockProvider("anthropic")
        with patch(
            "cortex.llm.fallback._import_provider_class",
            side_effect=lambda path: (
                MagicMock(return_value=mock_instance)
                if "anthropic" in path
                else (_ for _ in ()).throw(ImportError("sdk missing"))
            ),
        ):
            cascade = build_provider_cascade(["anthropic", "totally_unknown"])
            # "totally_unknown" is not in _PROVIDER_FACTORIES so it gets skipped
            # before _import_provider_class is even called. Anthropic succeeds.
            assert len(cascade.providers) >= 1

    def test_unknown_name_skipped_gracefully(self):
        """Unknown provider name is logged and skipped — does not raise."""
        mock_instance = MockProvider("anthropic")
        with patch(
            "cortex.llm.fallback._import_provider_class",
            return_value=MagicMock(return_value=mock_instance),
        ):
            # Mix of known and completely unknown names
            cascade = build_provider_cascade(["anthropic", "nonexistent_provider"])
            # nonexistent_provider is not in _PROVIDER_FACTORIES → skipped silently
            assert cascade is not None

    def test_raises_when_no_provider_initialises(self):
        """All providers fail to init → RuntimeError about empty cascade."""
        with patch(
            "cortex.llm.fallback._import_provider_class",
            side_effect=ImportError("sdk not installed"),
        ):
            with pytest.raises(RuntimeError, match="no LLM provider could be initialised"):
                build_provider_cascade(["anthropic", "openai", "ollama"])

    def test_strict_mode_raises_on_first_failure(self):
        """strict=True → raise immediately when first provider fails."""
        with patch(
            "cortex.llm.fallback._import_provider_class",
            side_effect=RuntimeError("API key missing"),
        ):
            with pytest.raises(RuntimeError, match="API key missing"):
                build_provider_cascade(["anthropic"], strict=True)

    def test_cascade_contains_providers_in_order(self):
        """Priority order must be preserved in the resulting cascade."""
        instances = {
            "anthropic": MockProvider("anthropic"),
            "openai": MockProvider("openai"),
            "ollama": MockProvider("ollama"),
        }

        def fake_import(path: str):
            for name, inst in instances.items():
                if name in path:
                    return MagicMock(return_value=inst)
            raise ImportError(f"unknown: {path}")

        with patch("cortex.llm.fallback._import_provider_class", side_effect=fake_import):
            cascade = build_provider_cascade(["anthropic", "openai", "ollama"])

        assert cascade.provider_names() == ["anthropic", "openai", "ollama"]

    def test_partial_init_failure_still_builds_cascade(self):
        """If some providers fail to init, the rest form a valid cascade."""
        ok_instance = MockProvider("ollama")

        def fake_import(path: str):
            if "ollama" in path:
                return MagicMock(return_value=ok_instance)
            raise RuntimeError("missing key")

        with patch("cortex.llm.fallback._import_provider_class", side_effect=fake_import):
            cascade = build_provider_cascade(["anthropic", "openai", "ollama"])

        assert "ollama" in cascade.provider_names()
        assert len(cascade.providers) == 1


# ===========================================================================
# CortexConfig — llm_provider_priority field
# ===========================================================================

class TestCortexConfigLLMPriority:

    def test_default_priority_is_anthropic_openai_ollama(self):
        config = CortexConfig()
        assert config.llm_provider_priority == ["anthropic", "openai", "ollama"]

    def test_from_toml_default_loads_priority(self):
        """Loading from defaults.toml should populate llm_provider_priority."""
        config = CortexConfig.from_toml()
        assert isinstance(config.llm_provider_priority, list)
        assert len(config.llm_provider_priority) >= 1
        assert "anthropic" in config.llm_provider_priority

    def test_from_raw_parses_llm_section(self):
        raw = {
            "cortex": {
                "llm": {
                    "provider_priority": ["openai", "ollama"],
                }
            }
        }
        config = CortexConfig._from_raw(raw)
        assert config.llm_provider_priority == ["openai", "ollama"]

    def test_from_raw_custom_order(self):
        raw = {
            "cortex": {
                "llm": {
                    "provider_priority": ["ollama", "anthropic"],
                }
            }
        }
        config = CortexConfig._from_raw(raw)
        assert config.llm_provider_priority == ["ollama", "anthropic"]

    def test_from_raw_single_provider(self):
        raw = {
            "cortex": {
                "llm": {
                    "provider_priority": ["anthropic"],
                }
            }
        }
        config = CortexConfig._from_raw(raw)
        assert config.llm_provider_priority == ["anthropic"]

    def test_from_raw_missing_llm_section_uses_default(self):
        raw = {"cortex": {}}
        config = CortexConfig._from_raw(raw)
        # Falls back to dataclass default
        assert config.llm_provider_priority == ["anthropic", "openai", "ollama"]

    def test_from_raw_ignores_non_list_priority(self):
        """If provider_priority is not a list of strings, it should be ignored
        and the dataclass default should be used instead."""
        raw = {
            "cortex": {
                "llm": {
                    "provider_priority": "anthropic",  # wrong type: str, not list
                }
            }
        }
        config = CortexConfig._from_raw(raw)
        # Type check in _from_raw rejects it → keeps default
        assert config.llm_provider_priority == ["anthropic", "openai", "ollama"]

    def test_overlay_overrides_priority(self):
        """In-process overlay dict takes precedence over defaults."""
        config = CortexConfig.from_toml(
            overlay={"cortex": {"llm": {"provider_priority": ["ollama"]}}}
        )
        assert config.llm_provider_priority == ["ollama"]


# ===========================================================================
# CortexPlugin — uses config cascade
# ===========================================================================

class TestCortexPluginCascadeIntegration:
    """
    Verify that CortexPlugin._init uses build_provider_cascade with config priority.
    Does NOT test the full plugin pipeline — just the LLM wiring.
    """

    def _make_mock_llm(self, name: str = "mock") -> FallbackLLMProvider:
        return FallbackLLMProvider([MockProvider(name)])

    def test_plugin_uses_build_provider_cascade(self):
        """CortexPlugin __init__ source must reference build_provider_cascade, not
        the old hardcoded AnthropicProvider / OpenAIProvider import chain."""
        import cortex.api.plugin as plugin_mod
        import inspect
        src = inspect.getsource(plugin_mod.CortexPlugin.__init__)
        assert "build_provider_cascade" in src
        # Old hardcoded fallback chain must be gone
        assert "AnthropicProvider" not in src
        assert "except ImportError" not in src or "AnthropicProvider" not in src

    def test_plugin_init_respects_config_priority_ordering(self):
        """Priority in config must flow through to build_provider_cascade call."""
        from cortex.api.plugin import CortexPlugin

        captured_priority: list = []

        def fake_build(priority, **kw):
            captured_priority.extend(priority)
            return FallbackLLMProvider([MockProvider("mock")])

        # Build a minimal config with custom priority
        custom_config = CortexConfig()
        custom_config.llm_provider_priority = ["ollama", "openai"]

        with patch("cortex.llm.fallback.build_provider_cascade", side_effect=fake_build), \
             patch("cortex.api.plugin.build_provider_cascade", side_effect=fake_build):
            with patch("cortex.config.CortexConfig.from_toml", return_value=custom_config):
                with patch("cortex.storage.sqlite_adapter.SQLiteAdapter.__init__", return_value=None):
                    with patch("cortex.storage.sqlite_adapter.SQLiteAdapter.ping", new_callable=AsyncMock):
                        # Suppress all other module inits to keep test fast
                        with patch("cortex.extraction.extractor.ExtractionPipeline.__init__", return_value=None):
                            with patch("cortex.core.entity_resolution.EntityResolver.__init__", return_value=None):
                                with patch("cortex.core.reconciliation.ReconciliationEngine.__init__", return_value=None):
                                    with patch("cortex.identity.cornerstone.CornerstoneGuardian.__init__", return_value=None):
                                        with patch("cortex.core.token_budget.TokenBudgetManager.__init__", return_value=None):
                                            with patch("cortex.core.retrieval.RetrievalPipeline.__init__", return_value=None):
                                                with patch("cortex.core.feedback.FeedbackSystem.__init__", return_value=None):
                                                    plugin = CortexPlugin.__new__(CortexPlugin)
                                                    plugin.__init__()

        assert captured_priority == ["ollama", "openai"]

    def test_plugin_source_uses_config_priority(self):
        """Smoke test: plugin __init__ source references config priority, not hardcoded names."""
        import cortex.api.plugin as plugin_mod
        import inspect
        src = inspect.getsource(plugin_mod.CortexPlugin.__init__)
        assert "llm_provider_priority" in src
        assert "build_provider_cascade" in src


# ===========================================================================
# Circuit breaker integration tests
# ===========================================================================

class TestCircuitBreakerIntegration:
    """Verify that the circuit breaker is actually wired into FallbackLLMProvider."""

    def _make_breaker(self, threshold=2, cooldown_sec=60):
        from cortex.llm.circuit_breaker import CircuitBreaker
        return CircuitBreaker(enabled=True, threshold=threshold, cooldown_sec=cooldown_sec)

    @pytest.mark.asyncio
    async def test_complete_skips_open_provider(self):
        """When a provider's breaker is OPEN, complete() skips it entirely."""
        from cortex.llm.circuit_breaker import CircuitBreaker, BreakerState

        breaker = self._make_breaker(threshold=1)
        good = MockProvider(name="good")
        bad = MockProvider(name="bad")

        # Trip the breaker for 'bad' provider
        breaker.record_failure("bad", RuntimeError("429 rate limit"))
        assert breaker.get_state("bad") == BreakerState.OPEN

        fallback = FallbackLLMProvider([bad, good], circuit_breaker=breaker)
        result = await fallback.complete(system="sys", user="usr")

        # Should have skipped 'bad' and used 'good'
        assert result.text == "ok"
        assert fallback.active_provider is good

    @pytest.mark.asyncio
    async def test_extract_skips_open_provider(self):
        """When a provider's breaker is OPEN, extract() skips it entirely."""
        from cortex.llm.circuit_breaker import BreakerState

        breaker = self._make_breaker(threshold=1)
        good = MockProvider(name="good")
        bad = MockProvider(name="bad")

        breaker.record_failure("bad", RuntimeError("429 rate limit"))
        assert breaker.get_state("bad") == BreakerState.OPEN

        fallback = FallbackLLMProvider([bad, good], circuit_breaker=breaker)
        result = await fallback.extract(system="sys", user="usr", schema={})
        assert result.data == {"result": "ok"}

    @pytest.mark.asyncio
    async def test_embed_never_blocked_by_breaker(self):
        """embed() is a read path — providers are NEVER skipped even with open breaker."""
        from cortex.llm.circuit_breaker import BreakerState

        breaker = self._make_breaker(threshold=1)
        provider = MockProvider(name="only")

        breaker.record_failure("only", RuntimeError("429 rate limit"))
        assert breaker.get_state("only") == BreakerState.OPEN

        fallback = FallbackLLMProvider([provider], circuit_breaker=breaker)
        result = await fallback.embed(texts=["hello"])
        assert len(result.embeddings) == 1

    @pytest.mark.asyncio
    async def test_complete_records_failure_and_trips_breaker(self):
        """Failures during complete() are recorded and can trip the breaker."""
        from cortex.llm.circuit_breaker import BreakerState

        breaker = self._make_breaker(threshold=2)
        failing = MockProvider(name="fail", fail_complete=True)
        good = MockProvider(name="good")

        fallback = FallbackLLMProvider([failing, good], circuit_breaker=breaker)

        # First call — failing provider fails, good provider succeeds
        await fallback.complete(system="sys", user="usr")
        # 'fail' now has 0 consecutive (transient errors don't trip by default)
        # Let's use a quota error directly
        breaker.record_failure("fail", RuntimeError("429 rate limit exceeded quota"))
        breaker.record_failure("fail", RuntimeError("429 rate limit exceeded quota"))
        assert breaker.get_state("fail") == BreakerState.OPEN

    @pytest.mark.asyncio
    async def test_complete_records_success_and_resets_breaker(self):
        """A successful call resets the breaker state."""
        from cortex.llm.circuit_breaker import BreakerState

        breaker = self._make_breaker(threshold=5)
        provider = MockProvider(name="p1")

        # Manually add some failures
        breaker.record_failure("p1", RuntimeError("429 rate limit"))
        breaker.record_failure("p1", RuntimeError("429 rate limit"))
        assert breaker._get_breaker("p1").consecutive_failures == 2

        fallback = FallbackLLMProvider([provider], circuit_breaker=breaker)
        await fallback.complete(system="sys", user="usr")

        # Success should reset
        assert breaker._get_breaker("p1").consecutive_failures == 0
        assert breaker.get_state("p1") == BreakerState.CLOSED

    @pytest.mark.asyncio
    async def test_all_providers_breaker_open_raises(self):
        """If ALL providers have open breakers, complete() raises CircuitBreakerOpenError."""
        from cortex.llm.circuit_breaker import CircuitBreakerOpenError, BreakerState

        breaker = self._make_breaker(threshold=1)
        p1 = MockProvider(name="p1")
        p2 = MockProvider(name="p2")

        breaker.record_failure("p1", RuntimeError("429 rate limit"))
        breaker.record_failure("p2", RuntimeError("429 rate limit"))

        fallback = FallbackLLMProvider([p1, p2], circuit_breaker=breaker)

        with pytest.raises(CircuitBreakerOpenError):
            await fallback.complete(system="sys", user="usr")
