"""
tests/test_venice_provider.py — Unit tests for VeniceProvider (#501).

Covers:
    - Provider init with / without VENICE_API_KEY
    - venice_parameters injection into every completion request
    - graceful cascade skip when API key is missing
    - JSON extraction flow
    - embed() raises NotImplementedError
    - Error handling: 429 retry behaviour, 503, non-retryable 400
    - Usage record logging
    - PROVIDER_NAME / model defaults

All tests are mock-based — no real API calls, no Venice API key required.
"""
from __future__ import annotations

import json
import os
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch, call

import pytest
import pytest_asyncio

from cortex.llm.provider import (
    CompletionResponse,
    ExtractionResponse,
    EmbedResponse,
    UsageRecord,
)


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

def _make_openai_response(
    text: str = "hello from venice",
    tokens_in: int = 10,
    tokens_out: int = 20,
    headers: dict | None = None,
) -> MagicMock:
    """Build a mock openai.ChatCompletion response object."""
    message = MagicMock()
    message.content = text

    choice = MagicMock()
    choice.message = message

    usage = MagicMock()
    usage.prompt_tokens = tokens_in
    usage.completion_tokens = tokens_out

    resp = MagicMock()
    resp.choices = [choice]
    resp.usage = usage
    resp.headers = headers or {}
    return resp


def _make_rate_limit_error(status_code: int = 429) -> Exception:
    """Simulate an openai RateLimitError or similar HTTP error."""
    exc = Exception(f"HTTP {status_code}: rate limit exceeded")
    exc.status_code = status_code  # type: ignore[attr-defined]
    return exc


def _make_http_error(status_code: int) -> Exception:
    exc = Exception(f"HTTP {status_code}")
    exc.status_code = status_code  # type: ignore[attr-defined]
    return exc


# ===========================================================================
# Fixtures
# ===========================================================================

@pytest.fixture()
def venice_env(monkeypatch):
    """Set VENICE_API_KEY so VeniceProvider can initialise."""
    monkeypatch.setenv("VENICE_API_KEY", "test-venice-key-abc123")
    # Clear any real base-URL override so tests use the default
    monkeypatch.delenv("VENICE_BASE_URL", raising=False)


@pytest.fixture()
def no_venice_env(monkeypatch):
    """Ensure VENICE_API_KEY is absent so VeniceProvider fails to init."""
    monkeypatch.delenv("VENICE_API_KEY", raising=False)


@pytest.fixture()
def provider(venice_env):
    """Return a VeniceProvider with a stubbed openai client."""
    from cortex.llm.venice_provider import VeniceProvider
    p = VeniceProvider()
    return p


# ===========================================================================
# 1. Initialisation
# ===========================================================================

class TestVeniceProviderInit:
    def test_init_success_with_api_key(self, venice_env):
        """Provider initialises cleanly when VENICE_API_KEY is set."""
        from cortex.llm.venice_provider import VeniceProvider
        p = VeniceProvider()
        assert p.PROVIDER_NAME == "venice"
        assert p.DEFAULT_COMPLETION_MODEL == "grok-41-fast"
        assert p._completion_model == "grok-41-fast"
        assert "venice.ai" in p._base_url

    def test_init_raises_without_api_key(self, no_venice_env):
        """Provider raises RuntimeError when VENICE_API_KEY is absent."""
        from cortex.llm.venice_provider import VeniceProvider
        with pytest.raises(RuntimeError, match="VENICE_API_KEY"):
            VeniceProvider()

    def test_init_custom_model(self, venice_env):
        """completion_model kwarg overrides the default."""
        from cortex.llm.venice_provider import VeniceProvider
        p = VeniceProvider(completion_model="gemini-3-flash-preview")
        assert p._completion_model == "gemini-3-flash-preview"

    def test_init_custom_base_url_kwarg(self, venice_env):
        """base_url kwarg overrides the default."""
        from cortex.llm.venice_provider import VeniceProvider
        p = VeniceProvider(base_url="https://custom.venice.test/v1")
        assert p._base_url == "https://custom.venice.test/v1"

    def test_init_custom_base_url_env(self, monkeypatch, venice_env):
        """VENICE_BASE_URL env var overrides the default."""
        from cortex.llm.venice_provider import VeniceProvider
        monkeypatch.setenv("VENICE_BASE_URL", "https://env-override.venice.test/v1")
        p = VeniceProvider()
        assert p._base_url == "https://env-override.venice.test/v1"

    def test_provider_name(self, venice_env):
        """PROVIDER_NAME is 'venice' — used by FallbackLLMProvider."""
        from cortex.llm.venice_provider import VeniceProvider
        assert VeniceProvider.PROVIDER_NAME == "venice"


# ===========================================================================
# 2. venice_parameters injection
# ===========================================================================

class TestVeniceParametersInjection:
    @pytest.mark.asyncio
    async def test_complete_injects_venice_parameters(self, provider):
        """venice_parameters must appear in extra_body on every complete() call."""
        mock_response = _make_openai_response("test reply")
        mock_client = MagicMock()
        mock_client.chat = MagicMock()
        mock_client.chat.completions = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.complete(system="sys", user="user msg")

        mock_client.chat.completions.create.assert_called_once()
        call_kwargs = mock_client.chat.completions.create.call_args.kwargs

        assert "extra_body" in call_kwargs
        vp = call_kwargs["extra_body"]["venice_parameters"]
        assert vp["include_venice_system_prompt"] is False
        assert vp["disable_thinking"] is True

    @pytest.mark.asyncio
    async def test_complete_venice_params_present_on_every_call(self, provider):
        """venice_parameters are injected on all calls, not just the first."""
        mock_response = _make_openai_response("reply")
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        for _ in range(3):
            await provider.complete(system="s", user="u")

        assert mock_client.chat.completions.create.call_count == 3
        for c in mock_client.chat.completions.create.call_args_list:
            vp = c.kwargs["extra_body"]["venice_parameters"]
            assert vp["include_venice_system_prompt"] is False
            assert vp["disable_thinking"] is True

    @pytest.mark.asyncio
    async def test_complete_passes_model_and_temperature(self, provider):
        """Model and temperature are forwarded to the API correctly."""
        mock_response = _make_openai_response()
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.complete(
            system="s", user="u",
            model="gemini-3-flash-preview",
            temperature=0.7,
        )

        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
        assert call_kwargs["model"] == "gemini-3-flash-preview"
        assert call_kwargs["temperature"] == 0.7

    @pytest.mark.asyncio
    async def test_complete_passes_max_tokens_when_set(self, provider):
        """max_tokens kwarg is forwarded when provided."""
        mock_response = _make_openai_response()
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.complete(system="s", user="u", max_tokens=512)

        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
        assert call_kwargs.get("max_tokens") == 512

    @pytest.mark.asyncio
    async def test_complete_no_response_format_param(self, provider):
        """response_format is NOT passed to Venice API (custom base URL pattern)."""
        mock_response = _make_openai_response()
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.complete(system="s", user="u", response_format="json")

        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
        assert "response_format" not in call_kwargs


# ===========================================================================
# 3. complete() response handling
# ===========================================================================

class TestVeniceCompleteResponse:
    @pytest.mark.asyncio
    async def test_complete_returns_completion_response(self, provider):
        """complete() returns a CompletionResponse with correct fields."""
        mock_response = _make_openai_response(
            text="venice says hi", tokens_in=15, tokens_out=25
        )
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        result = await provider.complete(system="sys", user="usr")

        assert isinstance(result, CompletionResponse)
        assert result.text == "venice says hi"
        assert result.tokens_in == 15
        assert result.tokens_out == 25
        assert result.model == "grok-41-fast"  # default
        assert result.latency_ms > 0

    @pytest.mark.asyncio
    async def test_complete_logs_usage_record(self, provider):
        """complete() appends a UsageRecord to the provider usage log."""
        mock_response = _make_openai_response(tokens_in=8, tokens_out=12)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        assert len(provider._usage_log) == 0
        await provider.complete(system="s", user="u")
        assert len(provider._usage_log) == 1

        rec = provider._usage_log[0]
        assert rec.provider == "venice"
        assert rec.call_type == "complete"
        assert rec.tokens_in == 8
        assert rec.tokens_out == 12

    @pytest.mark.asyncio
    async def test_complete_cost_is_zero_for_free_models(self, provider):
        """grok-41-fast is a free model — cost_estimate should be 0.0."""
        mock_response = _make_openai_response(tokens_in=100, tokens_out=200)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        result = await provider.complete(system="s", user="u")
        assert result.cost_estimate == 0.0

    @pytest.mark.asyncio
    async def test_complete_logs_venice_headers(self, provider, caplog):
        """Venice balance/ratelimit headers are logged at DEBUG level."""
        mock_response = _make_openai_response(
            headers={
                "x-venice-balance-usd": "4.20",
                "x-ratelimit-remaining-requests": "99",
            }
        )
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        import logging
        with caplog.at_level(logging.DEBUG, logger="cortex.llm.venice_provider"):
            await provider.complete(system="s", user="u")

        assert any("4.20" in r.message for r in caplog.records)
        assert any("99" in r.message for r in caplog.records)


# ===========================================================================
# 4. extract() flow
# ===========================================================================

class TestVeniceExtract:
    @pytest.mark.asyncio
    async def test_extract_returns_parsed_json(self, provider):
        """extract() parses the model's JSON text into a dict."""
        schema = {"type": "object", "properties": {"name": {"type": "string"}}}
        payload = json.dumps({"name": "Eva"})
        mock_response = _make_openai_response(text=payload)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        result = await provider.extract(
            system="extract name", user="My name is Eva", schema=schema
        )

        assert isinstance(result, ExtractionResponse)
        assert result.data == {"name": "Eva"}

    @pytest.mark.asyncio
    async def test_extract_injects_schema_into_system_prompt(self, provider):
        """extract() augments the system prompt with the JSON schema hint."""
        schema = {"type": "object", "properties": {"x": {"type": "integer"}}}
        payload = json.dumps({"x": 1})
        mock_response = _make_openai_response(text=payload)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.extract(system="base system", user="u", schema=schema)

        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
        system_msg = call_kwargs["messages"][0]["content"]
        assert "base system" in system_msg
        assert '"type": "object"' in system_msg

    @pytest.mark.asyncio
    async def test_extract_venice_params_still_injected(self, provider):
        """venice_parameters are injected even during extract() calls."""
        schema = {"type": "object"}
        payload = json.dumps({"ok": True})
        mock_response = _make_openai_response(text=payload)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.extract(system="s", user="u", schema=schema)

        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
        vp = call_kwargs["extra_body"]["venice_parameters"]
        assert vp["include_venice_system_prompt"] is False
        assert vp["disable_thinking"] is True

    @pytest.mark.asyncio
    async def test_extract_handles_json_parse_error(self, provider):
        """extract() returns a _raw/_parse_error dict on malformed JSON."""
        schema = {"type": "object"}
        mock_response = _make_openai_response(text="not json at all {{{")
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        result = await provider.extract(system="s", user="u", schema=schema)

        assert "_parse_error" in result.data
        assert "_raw" in result.data

    @pytest.mark.asyncio
    async def test_extract_logs_usage_record(self, provider):
        """extract() records usage (two records: one from complete, one from extract)."""
        schema = {"type": "object"}
        mock_response = _make_openai_response(text="{}")
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.extract(system="s", user="u", schema=schema)

        call_types = [r.call_type for r in provider._usage_log]
        # complete() logs "complete", extract() logs "extract"
        assert "complete" in call_types
        assert "extract" in call_types


# ===========================================================================
# 5. embed() — not supported
# ===========================================================================

class TestVeniceEmbed:
    @pytest.mark.asyncio
    async def test_embed_raises_not_implemented(self, provider):
        """embed() always raises NotImplementedError."""
        with pytest.raises(NotImplementedError, match="embeddings"):
            await provider.embed(texts=["hello"])

    @pytest.mark.asyncio
    async def test_embed_raises_regardless_of_model(self, provider):
        """embed() raises even if a model name is supplied."""
        with pytest.raises(NotImplementedError):
            await provider.embed(texts=["test"], model="text-embedding-3-small")


# ===========================================================================
# 6. Error handling — 429 and 503 retries
# ===========================================================================

class TestVeniceErrorHandling:
    @pytest.mark.asyncio
    async def test_complete_retries_on_429(self, provider):
        """complete() retries on HTTP 429 (rate limit) and succeeds on retry."""
        mock_response = _make_openai_response("ok after retry")
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(
            side_effect=[
                _make_rate_limit_error(429),
                mock_response,
            ]
        )
        provider._openai_client = mock_client

        with patch("asyncio.sleep", new_callable=AsyncMock):
            result = await provider.complete(system="s", user="u")

        assert result.text == "ok after retry"
        assert mock_client.chat.completions.create.call_count == 2

    @pytest.mark.asyncio
    async def test_complete_retries_on_503(self, provider):
        """complete() retries on HTTP 503 (service unavailable)."""
        mock_response = _make_openai_response("recovered")
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(
            side_effect=[
                _make_http_error(503),
                mock_response,
            ]
        )
        provider._openai_client = mock_client

        with patch("asyncio.sleep", new_callable=AsyncMock):
            result = await provider.complete(system="s", user="u")

        assert result.text == "recovered"

    @pytest.mark.asyncio
    async def test_complete_does_not_retry_on_400(self, provider):
        """complete() does NOT retry on HTTP 400 (bad request — non-retryable)."""
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(
            side_effect=_make_http_error(400)
        )
        provider._openai_client = mock_client

        with pytest.raises(Exception, match="400"):
            await provider.complete(system="s", user="u")

        # Should only be called once — no retry on 400
        assert mock_client.chat.completions.create.call_count == 1

    @pytest.mark.asyncio
    async def test_complete_raises_after_all_retries_exhausted(self, provider):
        """complete() raises after max retries (3) on repeated 429s."""
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(
            side_effect=_make_rate_limit_error(429)
        )
        provider._openai_client = mock_client

        with patch("asyncio.sleep", new_callable=AsyncMock):
            with pytest.raises(Exception, match="429"):
                await provider.complete(system="s", user="u")

        # 1 initial attempt + 3 retries = 4 total calls
        assert mock_client.chat.completions.create.call_count == 4


# ===========================================================================
# 7. Cascade integration — graceful skip when key is missing
# ===========================================================================

class TestVeniceCascadeIntegration:
    def test_build_cascade_skips_venice_without_key(self, no_venice_env):
        """build_provider_cascade skips Venice silently when key is absent."""
        from cortex.llm.fallback import build_provider_cascade

        # Mock anthropic and openai so they also fail (no keys in test env)
        # We only care that Venice is skipped, not that the cascade succeeds.
        with pytest.raises(RuntimeError, match="no LLM provider"):
            # All providers fail in this env — RuntimeError is expected.
            build_provider_cascade(["venice"])

    def test_build_cascade_venice_registered_in_factories(self):
        """'venice' is a known provider name in _PROVIDER_FACTORIES."""
        from cortex.llm.fallback import _PROVIDER_FACTORIES
        assert "venice" in _PROVIDER_FACTORIES
        assert "venice_provider.VeniceProvider" in _PROVIDER_FACTORIES["venice"]

    def test_build_cascade_with_venice_in_priority(self, venice_env):
        """Venice is instantiated when its key is present in cascade build."""
        from cortex.llm.fallback import build_provider_cascade
        # Build a cascade with only venice — should succeed
        cascade = build_provider_cascade(["venice"])
        names = cascade.provider_names()
        assert "venice" in names

    def test_build_cascade_venice_priority_position(self, venice_env):
        """Venice appears after openai, before ollama in the cascade."""
        from cortex.llm.fallback import build_provider_cascade
        # Only include venice and ollama for a minimal test
        cascade = build_provider_cascade(["venice", "ollama"])
        names = cascade.provider_names()
        assert names.index("venice") < names.index("ollama")


# ===========================================================================
# 8. Usage summary
# ===========================================================================

class TestVeniceUsageSummary:
    @pytest.mark.asyncio
    async def test_usage_summary_after_calls(self, provider):
        """usage_summary() aggregates across multiple calls."""
        mock_response = _make_openai_response(tokens_in=10, tokens_out=5)
        mock_client = MagicMock()
        mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
        provider._openai_client = mock_client

        await provider.complete(system="s", user="u")
        await provider.complete(system="s2", user="u2")

        summary = provider.usage_summary()
        # complete() + extract() path: we called complete twice directly
        assert summary["calls"] >= 2
        assert summary["tokens_in"] >= 20
        assert summary["tokens_out"] >= 10
