"""
tests/test_llm_robustness.py — Tests for LLM robustness fixes (#fix/llm-robustness).

Covers:
    1. Retry/backoff: transient errors retried, non-retryable errors propagated
    2. Retry exhaustion triggers cascade fallover (not premature cascade)
    3. Ollama embed_model passthrough: cloud models rejected
    4. Embedding model change warning detection
    5. Retrieval weights wired from config

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

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

from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    LLMProvider,
    retry_with_backoff,
    _is_retryable_exception,
)
from cortex.llm.fallback import FallbackLLMProvider, build_provider_cascade
from cortex.llm.ollama_provider import OllamaProvider
from cortex.core.retrieval import RetrievalPipeline


# ===========================================================================
# Helpers
# ===========================================================================

def _make_completion(text: str = "ok", model: str = "test") -> CompletionResponse:
    return CompletionResponse(
        text=text, 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=5, tokens_out=0, cost_estimate=0.0, latency_ms=5.0,
    )


class _FakeRateLimitError(Exception):
    """Simulates an HTTP 429 from an LLM SDK."""
    def __init__(self):
        super().__init__("Rate limit exceeded")
        self.status_code = 429


class _FakeServerError(Exception):
    """Simulates an HTTP 500 from an LLM SDK."""
    def __init__(self):
        super().__init__("Internal server error")
        self.status_code = 500


class _FakeBadRequestError(Exception):
    """Simulates an HTTP 400 — should NOT be retried."""
    def __init__(self):
        super().__init__("Bad request")
        self.status_code = 400


class _FakeAuthError(Exception):
    """Simulates an HTTP 401 — should NOT be retried."""
    def __init__(self):
        super().__init__("Unauthorized")
        self.status_code = 401


class _StubProvider(LLMProvider):
    """Minimal LLMProvider for testing."""
    PROVIDER_NAME = "stub"
    DEFAULT_COMPLETION_MODEL = "stub-model"
    DEFAULT_EMBED_MODEL = "stub-embed"

    def __init__(self, complete_fn=None, embed_fn=None):
        super().__init__()
        self._complete_fn = complete_fn or AsyncMock(return_value=_make_completion())
        self._embed_fn = embed_fn or AsyncMock(return_value=_make_embed())

    async def complete(self, system, user, model=None, **kw):
        return await self._complete_fn()

    async def extract(self, system, user, schema, model=None, **kw):
        raise NotImplementedError

    async def embed(self, texts, model=None, input_type="query"):
        return await self._embed_fn()


# ===========================================================================
# Test 1: retry_with_backoff succeeds after transient 429
# ===========================================================================

@pytest.mark.asyncio
async def test_retry_succeeds_after_transient_429():
    """A 429 on first attempt should be retried; second attempt succeeds."""
    call_count = 0

    async def flaky_call():
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise _FakeRateLimitError()
        return "success"

    result = await retry_with_backoff(
        flaky_call,
        max_retries=3,
        backoff_base=0.01,  # fast for tests
        operation_name="test_429",
    )
    assert result == "success"
    assert call_count == 2  # one failure + one success


@pytest.mark.asyncio
async def test_retry_succeeds_after_server_error():
    """HTTP 500 on first attempt retried; second attempt succeeds."""
    call_count = 0

    async def flaky_call():
        nonlocal call_count
        call_count += 1
        if call_count <= 2:
            raise _FakeServerError()
        return "recovered"

    result = await retry_with_backoff(
        flaky_call,
        max_retries=3,
        backoff_base=0.01,
        operation_name="test_500",
    )
    assert result == "recovered"
    assert call_count == 3


# ===========================================================================
# Test 2: retry exhaustion triggers cascade fallover
# ===========================================================================

@pytest.mark.asyncio
async def test_retry_exhaustion_raises_for_cascade():
    """When all retries exhausted, the error propagates so FallbackLLMProvider can cascade."""
    async def always_429():
        raise _FakeRateLimitError()

    with pytest.raises(_FakeRateLimitError):
        await retry_with_backoff(
            always_429,
            max_retries=2,
            backoff_base=0.01,
            operation_name="test_exhaust",
        )


@pytest.mark.asyncio
async def test_cascade_triggers_after_retry_exhaustion():
    """FallbackLLMProvider should cascade to next provider only after retries fail."""
    call_counts = {"primary": 0, "secondary": 0}

    async def primary_complete():
        call_counts["primary"] += 1
        raise _FakeRateLimitError()

    async def secondary_complete():
        call_counts["secondary"] += 1
        return _make_completion("from-secondary")

    primary = _StubProvider(complete_fn=primary_complete)
    secondary = _StubProvider(complete_fn=secondary_complete)
    fallback = FallbackLLMProvider([primary, secondary])

    result = await fallback.complete(system="s", user="u")
    assert result.text == "from-secondary"
    # Primary was called once (then its internal retry_with_backoff would have
    # retried — but since _StubProvider.complete doesn't use retry, it raises immediately).
    # In the real providers, retry_with_backoff handles the retries internally.
    assert call_counts["primary"] == 1
    assert call_counts["secondary"] == 1


# ===========================================================================
# Test 3: Non-retryable errors propagate immediately (no retry)
# ===========================================================================

@pytest.mark.asyncio
async def test_non_retryable_400_not_retried():
    """HTTP 400 should NOT be retried."""
    call_count = 0

    async def bad_request():
        nonlocal call_count
        call_count += 1
        raise _FakeBadRequestError()

    with pytest.raises(_FakeBadRequestError):
        await retry_with_backoff(
            bad_request,
            max_retries=3,
            backoff_base=0.01,
            operation_name="test_400",
        )
    assert call_count == 1  # only one attempt — not retried


@pytest.mark.asyncio
async def test_non_retryable_401_not_retried():
    """HTTP 401 should NOT be retried."""
    call_count = 0

    async def auth_error():
        nonlocal call_count
        call_count += 1
        raise _FakeAuthError()

    with pytest.raises(_FakeAuthError):
        await retry_with_backoff(
            auth_error,
            max_retries=3,
            backoff_base=0.01,
            operation_name="test_401",
        )
    assert call_count == 1


# ===========================================================================
# Test 4: _is_retryable_exception classification
# ===========================================================================

def test_retryable_classification():
    """Verify exception classification for retry decisions."""
    assert _is_retryable_exception(_FakeRateLimitError()) is True
    assert _is_retryable_exception(_FakeServerError()) is True
    assert _is_retryable_exception(ConnectionError("reset")) is True
    assert _is_retryable_exception(TimeoutError("timeout")) is True
    assert _is_retryable_exception(_FakeBadRequestError()) is False
    assert _is_retryable_exception(_FakeAuthError()) is False
    assert _is_retryable_exception(ValueError("bad value")) is False


# ===========================================================================
# Test 5: Ollama doesn't receive OpenAI/Voyage embed model names
# ===========================================================================

def test_ollama_rejects_openai_embed_model():
    """OllamaProvider should fall back to default when given an OpenAI embed model."""
    provider = OllamaProvider(embed_model="text-embedding-3-small")
    assert provider._embed_model == OllamaProvider.DEFAULT_EMBED_MODEL


def test_ollama_rejects_voyage_embed_model():
    """OllamaProvider should fall back to default when given a Voyage embed model."""
    provider = OllamaProvider(embed_model="voyage-code-3")
    assert provider._embed_model == OllamaProvider.DEFAULT_EMBED_MODEL


def test_ollama_accepts_valid_embed_model():
    """OllamaProvider should accept valid Ollama embed models."""
    provider = OllamaProvider(embed_model="nomic-embed-text")
    assert provider._embed_model == "nomic-embed-text"

    provider2 = OllamaProvider(embed_model="mxbai-embed-large")
    assert provider2._embed_model == "mxbai-embed-large"


def test_build_cascade_skips_embed_model_for_ollama():
    """build_provider_cascade should not pass cloud embed_model to Ollama."""
    with patch("cortex.llm.fallback._import_provider_class") as mock_import:
        # Create a mock OllamaProvider class
        mock_ollama_cls = MagicMock()
        mock_ollama_instance = MagicMock(spec=LLMProvider)
        mock_ollama_cls.return_value = mock_ollama_instance

        mock_import.return_value = mock_ollama_cls

        build_provider_cascade(
            ["ollama"],
            embed_model="text-embedding-3-small",
        )

        # Ollama should be called WITHOUT embed_model
        mock_ollama_cls.assert_called_once()
        call_kwargs = mock_ollama_cls.call_args[1]
        assert "embed_model" not in call_kwargs


# ===========================================================================
# Test 6: Embedding model change warning
# ===========================================================================

def test_embedding_model_change_detection(tmp_path, caplog):
    """Plugin should warn when embedding model changes between runs."""
    import sqlite3

    db_path = str(tmp_path / "test_cortex.db")

    # Create a minimal database with the old model recorded
    conn = sqlite3.connect(db_path)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS cortex_config (
            key TEXT PRIMARY KEY,
            value TEXT
        )
    """)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS embedding_index (
            id TEXT PRIMARY KEY,
            item_type TEXT,
            item_id TEXT,
            embedding BLOB
        )
    """)
    # Record old model and some embeddings
    conn.execute(
        "INSERT INTO cortex_config (key, value) VALUES ('embedding_model', 'text-embedding-3-small')"
    )
    conn.execute(
        "INSERT INTO embedding_index (id, item_type, item_id, embedding) "
        "VALUES ('e1', 'claim', 'c1', X'00')"
    )
    conn.commit()
    conn.close()

    # Mock config with a DIFFERENT model
    mock_config = MagicMock()
    mock_config.embedding_model = "voyage-code-3"
    mock_config.storage_db_path = db_path
    mock_config.storage_enable_wal = False
    mock_config.storage_busy_timeout_ms = 5000
    mock_config.owner_id = "test"
    mock_config.retrieval_token_budget = 3000
    mock_config.retrieval_vector_weight = 0.7
    mock_config.retrieval_bm25_weight = 0.3
    mock_config.llm_provider_priority = []
    mock_config.extraction_model = "test"
    mock_config.max_llm_calls_per_session = 0

    # Create a mock storage with _get_connection
    mock_storage = MagicMock()
    mock_conn = sqlite3.connect(db_path)
    mock_storage._get_connection.return_value = mock_conn

    # Manually test the check method
    from cortex.api.plugin import CortexPlugin

    # We can't easily init the full plugin, so test the method directly
    plugin = object.__new__(CortexPlugin)
    plugin.config = mock_config
    plugin.storage = mock_storage
    plugin._embedding_model_changed = False

    with caplog.at_level(logging.WARNING):
        plugin._check_embedding_model_consistency()

    assert plugin._embedding_model_changed is True
    assert "EMBEDDING MODEL CHANGED" in caplog.text
    assert "text-embedding-3-small" in caplog.text
    assert "voyage-code-3" in caplog.text

    mock_conn.close()


# ===========================================================================
# Test 7: Retrieval weights from config are used
# ===========================================================================

def test_retrieval_pipeline_accepts_weights():
    """RetrievalPipeline should use weights from constructor kwargs."""
    mock_storage = MagicMock()
    mock_config = MagicMock()
    mock_config.retrieval_token_budget = 3000

    pipeline = RetrievalPipeline(
        storage=mock_storage,
        config=mock_config,
        vector_weight=0.8,
        bm25_weight=0.2,
    )
    assert pipeline.vector_weight == 0.8
    assert pipeline.bm25_weight == 0.2


def test_retrieval_weights_default_values():
    """RetrievalPipeline should use 0.7/0.3 defaults when not specified."""
    mock_storage = MagicMock()
    mock_config = MagicMock()

    pipeline = RetrievalPipeline(
        storage=mock_storage,
        config=mock_config,
    )
    assert pipeline.vector_weight == 0.7
    assert pipeline.bm25_weight == 0.3


def test_rrf_fusion_basic():
    """rrf_fusion() should combine ranked lists correctly (replaces hybrid_score test)."""
    from cortex.core.retrieval import rrf_fusion
    result = rrf_fusion([
        [("a", 1.0), ("b", 0.5)],
        [("b", 1.0), ("a", 0.5)],
    ], k=60)
    # Both items appear in both lists, so they should have equal RRF scores
    scores = dict(result)
    assert abs(scores["a"] - scores["b"]) < 1e-9


# ===========================================================================
# Test 8: ConnectionError / TimeoutError are retryable
# ===========================================================================

@pytest.mark.asyncio
async def test_retry_on_connection_error():
    """ConnectionError should be retried."""
    call_count = 0

    async def flaky():
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise ConnectionError("Connection refused")
        return "connected"

    result = await retry_with_backoff(
        flaky, max_retries=2, backoff_base=0.01,
        operation_name="test_conn",
    )
    assert result == "connected"
    assert call_count == 2


@pytest.mark.asyncio
async def test_retry_on_timeout_error():
    """TimeoutError should be retried."""
    call_count = 0

    async def flaky():
        nonlocal call_count
        call_count += 1
        if call_count <= 2:
            raise TimeoutError("Request timed out")
        return "ok"

    result = await retry_with_backoff(
        flaky, max_retries=3, backoff_base=0.01,
        operation_name="test_timeout",
    )
    assert result == "ok"
    assert call_count == 3


# ===========================================================================
# Test 9: Embedding model change — no warning when model unchanged
# ===========================================================================

def test_no_warning_when_model_unchanged(tmp_path, caplog):
    """No warning when embedding model hasn't changed."""
    import sqlite3

    db_path = str(tmp_path / "test_cortex.db")
    conn = sqlite3.connect(db_path)
    conn.execute("CREATE TABLE cortex_config (key TEXT PRIMARY KEY, value TEXT)")
    conn.execute("CREATE TABLE embedding_index (id TEXT PRIMARY KEY, item_type TEXT, item_id TEXT, embedding BLOB)")
    conn.execute("INSERT INTO cortex_config (key, value) VALUES ('embedding_model', 'voyage-code-3')")
    conn.execute("INSERT INTO embedding_index (id, item_type, item_id, embedding) VALUES ('e1', 'claim', 'c1', X'00')")
    conn.commit()

    mock_config = MagicMock()
    mock_config.embedding_model = "voyage-code-3"  # Same model

    mock_storage = MagicMock()
    mock_storage._get_connection.return_value = conn

    from cortex.api.plugin import CortexPlugin
    plugin = object.__new__(CortexPlugin)
    plugin.config = mock_config
    plugin.storage = mock_storage
    plugin._embedding_model_changed = False

    with caplog.at_level(logging.WARNING):
        plugin._check_embedding_model_consistency()

    assert plugin._embedding_model_changed is False
    assert "EMBEDDING MODEL CHANGED" not in caplog.text
    conn.close()


# ===========================================================================
# Test 10: Retrieve emits warning when embedding model changed
# ===========================================================================

@pytest.mark.asyncio
async def test_retrieve_warns_on_model_change(caplog):
    """retrieve() should log a warning on first call when embedding model changed."""
    from cortex.api.plugin import CortexPlugin

    plugin = object.__new__(CortexPlugin)
    plugin._embedding_model_changed = True
    plugin._embedding_model_warning_emitted = False
    plugin.config = MagicMock()
    plugin.config.retrieval_token_budget = 3000
    plugin.config.owner_id = "test"

    mock_result = MagicMock()
    mock_retrieval = MagicMock()
    mock_retrieval.retrieve = AsyncMock(return_value=mock_result)
    plugin.retrieval = mock_retrieval

    with caplog.at_level(logging.WARNING):
        # We need to mock validate_retrieve_query
        with patch("cortex.api.plugin.validate_retrieve_query", return_value="test query"):
            result = await plugin.retrieve("test query")

    assert "changed embedding model" in caplog.text.lower() or "degraded" in caplog.text.lower()
    assert plugin._embedding_model_warning_emitted is True

    # Second call should NOT warn again
    caplog.clear()
    with caplog.at_level(logging.WARNING):
        with patch("cortex.api.plugin.validate_retrieve_query", return_value="test query"):
            await plugin.retrieve("test query 2")
    assert "degraded" not in caplog.text.lower()
