"""tests/test_per_operation_provider.py — Per-operation LLM provider routing (#1247).

Tests:
1. Default config: all operations use the shared main provider (backward compat)
2. Extraction override only: extraction gets dedicated provider, others share main
3. All operations overridden: each gets its own provider
4. Missing API key falls back gracefully to main provider key
5. Config fields have correct env var mapping
6. build_operation_provider returns None when no override configured
"""
from __future__ import annotations

import os
from dataclasses import fields
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.config import CortexConfig
from cortex.llm.provider import CompletionResponse, EmbedResponse, LLMProvider
from cortex.llm.provider_router import ProviderRouter, build_operation_provider


def _cr(t="ok"):
    return CompletionResponse(
        text=t, model="t", tokens_in=5, tokens_out=3,
        cost_estimate=0, latency_ms=10,
    )


def _mock_llm():
    m = MagicMock(spec=LLMProvider)
    m.PROVIDER_NAME = "mock"
    m.complete = AsyncMock(return_value=_cr())
    m.embed = AsyncMock(return_value=EmbedResponse(
        embeddings=[[0.1]], model="t", tokens_in=4,
        tokens_out=0, cost_estimate=0, latency_ms=5,
    ))
    m.usage_summary = MagicMock(return_value={"calls": 0})
    m._usage_log = []
    return m


# ---------------------------------------------------------------------------
# CortexConfig field tests
# ---------------------------------------------------------------------------

class TestConfigFields:
    """Verify per-operation provider fields exist on CortexConfig."""

    def test_extraction_fields_exist(self):
        c = CortexConfig()
        assert c.provider_extraction_base_url == ""
        assert c.provider_extraction_api_key == ""

    def test_reconciliation_fields_exist(self):
        c = CortexConfig()
        assert c.provider_reconciliation_base_url == ""
        assert c.provider_reconciliation_api_key == ""

    def test_dialectic_fields_exist(self):
        c = CortexConfig()
        assert c.provider_dialectic_base_url == ""
        assert c.provider_dialectic_api_key == ""

    def test_env_var_mapping(self):
        """Verify CORTEX_PROVIDER_EXTRACTION_BASE_URL maps to the field."""
        field_names = {f.name for f in fields(CortexConfig)}
        # These field names should exist so the env override mechanism picks them up
        assert "provider_extraction_base_url" in field_names
        assert "provider_extraction_api_key" in field_names
        assert "provider_reconciliation_base_url" in field_names
        assert "provider_reconciliation_api_key" in field_names
        assert "provider_dialectic_base_url" in field_names
        assert "provider_dialectic_api_key" in field_names

    def test_env_override_picks_up_extraction_field(self):
        """CORTEX_PROVIDER_EXTRACTION_BASE_URL should be picked up by _env_overrides."""
        with patch.dict(os.environ, {
            "CORTEX_PROVIDER_EXTRACTION_BASE_URL": "https://extraction.example.com/v1",
            "CORTEX_PROVIDER_EXTRACTION_API_KEY": "sk-extract-123",
        }):
            overrides = CortexConfig._env_overrides()
            assert overrides.get("provider_extraction_base_url") == "https://extraction.example.com/v1"
            assert overrides.get("provider_extraction_api_key") == "sk-extract-123"


# ---------------------------------------------------------------------------
# build_operation_provider tests
# ---------------------------------------------------------------------------

class TestBuildOperationProvider:
    """Tests for the per-operation provider factory."""

    def test_returns_none_when_no_override(self):
        """No per-op config → returns None (caller uses main provider)."""
        c = CortexConfig()
        result = build_operation_provider(c, "extraction")
        assert result is None

    def test_returns_none_for_empty_base_url(self):
        c = CortexConfig(provider_extraction_base_url="")
        result = build_operation_provider(c, "extraction")
        assert result is None

    def test_creates_provider_with_override(self):
        """Per-op base_url set → returns a new provider."""
        c = CortexConfig(
            provider_extraction_base_url="https://extraction.example.com/v1",
            provider_extraction_api_key="sk-extract-123",
            extraction_model="minimax-m2.5",
        )
        provider = build_operation_provider(c, "extraction")
        assert provider is not None
        assert isinstance(provider, LLMProvider)

    def test_falls_back_to_main_api_key(self):
        """Per-op key empty → uses main provider_llm_api_key."""
        c = CortexConfig(
            provider_extraction_base_url="https://extraction.example.com/v1",
            provider_extraction_api_key="",
            provider_llm_api_key="sk-main-key",
        )
        provider = build_operation_provider(c, "extraction")
        assert provider is not None

    def test_each_operation_gets_own_provider(self):
        """All three operations get independent providers."""
        c = CortexConfig(
            provider_extraction_base_url="https://extract.example.com/v1",
            provider_extraction_api_key="sk-e",
            provider_reconciliation_base_url="https://recon.example.com/v1",
            provider_reconciliation_api_key="sk-r",
            provider_dialectic_base_url="https://dial.example.com/v1",
            provider_dialectic_api_key="sk-d",
            extraction_model="model-e",
            reconciliation_model="model-r",
            dialectic_model="model-d",
        )
        pe = build_operation_provider(c, "extraction")
        pr = build_operation_provider(c, "reconciliation")
        pd = build_operation_provider(c, "dialectic")
        assert pe is not None and pr is not None and pd is not None
        # They should be different objects (independent rate limits)
        assert pe is not pr
        assert pr is not pd

    def test_reconciliation_provider(self):
        c = CortexConfig(
            provider_reconciliation_base_url="https://recon.example.com/v1",
            provider_reconciliation_api_key="sk-r",
        )
        provider = build_operation_provider(c, "reconciliation")
        assert provider is not None

    def test_dialectic_provider(self):
        c = CortexConfig(
            provider_dialectic_base_url="https://dial.example.com/v1",
            provider_dialectic_api_key="sk-d",
        )
        provider = build_operation_provider(c, "dialectic")
        assert provider is not None


# ---------------------------------------------------------------------------
# Initialization wiring tests
# ---------------------------------------------------------------------------

class TestInitializationWiring:
    """Verify CortexPlugin creates per-operation LLM instances."""

    @patch("cortex.api.plugin.ProviderRouter")
    @patch("cortex.api.plugin.build_provider_cascade")
    def test_default_all_ops_use_shared_provider(self, mock_cascade, mock_router_cls):
        """When no per-op config, all operations share self.llm."""
        mock_provider = _mock_llm()
        mock_router_cls.return_value = mock_provider

        from cortex.api.plugin import CortexPlugin
        plugin = CortexPlugin()

        # All per-op LLMs should be the same as (or wrapping) the shared provider
        assert plugin._extraction_llm is not None
        assert plugin._reconciliation_llm is not None
        assert plugin._dialectic_llm is not None

    @patch("cortex.api.plugin.ProviderRouter")
    @patch("cortex.api.plugin.build_provider_cascade")
    @patch("cortex.llm.provider_router.build_operation_provider")
    def test_extraction_override_uses_dedicated_provider(
        self, mock_build_op, mock_cascade, mock_router_cls,
    ):
        """When extraction provider configured, extractor gets dedicated LLM."""
        mock_main = _mock_llm()
        mock_router_cls.return_value = mock_main

        mock_extraction_provider = _mock_llm()
        mock_extraction_provider.PROVIDER_NAME = "extraction-dedicated"

        def side_effect(config, operation):
            if operation == "extraction":
                return mock_extraction_provider
            return None

        mock_build_op.side_effect = side_effect

        from cortex.api.plugin import CortexPlugin
        plugin = CortexPlugin()

        # Extraction should use the dedicated provider (wrapped in ModelRouter).
        # Verify the extraction LLM is distinct AND backed by our mock provider.
        assert plugin._extraction_llm is not plugin.llm, (
            "Extraction LLM should be a separate object from the shared LLM"
        )
        # The dedicated provider is wrapped in a ModelRouter — verify inner provider
        assert hasattr(plugin._extraction_llm, '_provider'), (
            "Extraction LLM should be a ModelRouter wrapping the dedicated provider"
        )
        assert plugin._extraction_llm._provider is mock_extraction_provider, (
            "Extraction LLM's inner provider should be the dedicated extraction provider"
        )


# ---------------------------------------------------------------------------
# TOML parsing tests
# ---------------------------------------------------------------------------

class TestTOMLParsing:
    """Test that per-operation provider fields are parsed from TOML-style dicts."""

    def test_from_raw_extraction_override(self):
        """Providers.extraction section is parsed correctly."""
        raw = {
            "cortex": {
                "providers": {
                    "llm": {
                        "provider": "openai",
                        "model": "gpt-4.1-mini",
                        "api_key": "sk-main",
                        "base_url": "https://main.example.com/v1",
                    },
                    "extraction": {
                        "base_url": "https://extraction.example.com/v1",
                        "api_key": "sk-extract",
                    },
                },
            },
        }
        config = CortexConfig._from_raw(raw)
        assert config.provider_extraction_base_url == "https://extraction.example.com/v1"
        assert config.provider_extraction_api_key == "sk-extract"
        # Main provider unchanged
        assert config.provider_llm_base_url == "https://main.example.com/v1"
        assert config.provider_llm_api_key == "sk-main"

    def test_from_raw_all_overrides(self):
        raw = {
            "cortex": {
                "providers": {
                    "extraction": {
                        "base_url": "https://e.example.com/v1",
                        "api_key": "sk-e",
                    },
                    "reconciliation": {
                        "base_url": "https://r.example.com/v1",
                        "api_key": "sk-r",
                    },
                    "dialectic": {
                        "base_url": "https://d.example.com/v1",
                        "api_key": "sk-d",
                    },
                },
            },
        }
        config = CortexConfig._from_raw(raw)
        assert config.provider_extraction_base_url == "https://e.example.com/v1"
        assert config.provider_reconciliation_base_url == "https://r.example.com/v1"
        assert config.provider_dialectic_base_url == "https://d.example.com/v1"

    def test_from_raw_no_overrides(self):
        """When no per-op section, fields remain empty (backward compat)."""
        raw = {
            "cortex": {
                "providers": {
                    "llm": {
                        "model": "gpt-4.1-mini",
                    },
                },
            },
        }
        config = CortexConfig._from_raw(raw)
        assert config.provider_extraction_base_url == ""
        assert config.provider_reconciliation_base_url == ""
        assert config.provider_dialectic_base_url == ""


# ---------------------------------------------------------------------------
# Backward compatibility
# ---------------------------------------------------------------------------

class TestBackwardCompatibility:
    """Verify zero regression when no per-operation config is set."""

    def test_default_config_unchanged(self):
        """CortexConfig() with no per-op fields should have empty strings."""
        c = CortexConfig()
        assert c.provider_extraction_base_url == ""
        assert c.provider_extraction_api_key == ""
        assert c.provider_reconciliation_base_url == ""
        assert c.provider_reconciliation_api_key == ""
        assert c.provider_dialectic_base_url == ""
        assert c.provider_dialectic_api_key == ""

    def test_build_operation_provider_returns_none_for_defaults(self):
        """Default config → no per-op providers created."""
        c = CortexConfig()
        for op in ("extraction", "reconciliation", "dialectic"):
            assert build_operation_provider(c, op) is None
