# tests/test_minima_llm_resilience.py
"""
Tests for MinimaLLM 502 error resilience and infinite retry behavior.

These tests verify:
- 502 is treated as an overload status (triggers cooldown)
- max_attempts=0 allows infinite retries
- Overload warning printed once per request
- Recovery message printed after overload recovery
"""
import os
import asyncio
from dataclasses import replace
from unittest.mock import patch
import pytest

from minima_llm import (
    OpenAIMinimaLlm,
    MinimaLlmConfig,
    MinimaLlmRequest,
    MinimaLlmResponse,
    MinimaLlmFailure,
)
from minima_llm.backend import (
    _is_overload_status,
    _is_retriable_status,
)

def skip_llm_tests() -> bool:
    return os.getenv("SKIP_LLM_ENDPOINT_TESTS", "true").lower() in {"1", "true", "yes"}

class TestOverloadStatusCodes:
    """Test that 502 is properly classified as overload status."""

    def test_502_is_overload_status(self):
        """502 Bad Gateway should be considered an overload status."""
        assert _is_overload_status(502) is True

    def test_502_is_retriable_status(self):
        """502 Bad Gateway should be retriable."""
        assert _is_retriable_status(502) is True

    def test_overload_status_codes(self):
        """All expected overload status codes should return True."""
        overload_codes = [429, 502, 503, 504]
        for code in overload_codes:
            assert _is_overload_status(code) is True, f"{code} should be overload status"

    def test_non_overload_status_codes(self):
        """Non-overload codes should return False."""
        non_overload = [200, 400, 401, 403, 404, 500]
        for code in non_overload:
            assert _is_overload_status(code) is False, f"{code} should not be overload status"


@pytest.mark.skipif(
    skip_llm_tests(),
    reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
)
@pytest.fixture
def base_config():
    """Load config from environment, override timeouts for fast tests."""
    cfg = MinimaLlmConfig.from_env()
    return replace(
        cfg,
        max_outstanding=1,
        timeout_s=1.0,
        base_backoff_s=0.01,
        max_backoff_s=0.02,
        cooldown_floor_s=0.01,
        cooldown_cap_s=0.05,
        cooldown_halflife_s=1.0,
        cache_dir=None,  # Disable cache for tests
    )


class TestInfiniteRetries:
    """Test max_attempts=0 for infinite retries."""



    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_infinite_retries_eventually_succeeds(self, base_config):
        """With max_attempts=0, retries should continue until success."""
        config = replace(base_config, max_attempts=0)
        backend = OpenAIMinimaLlm(config)

        call_count = 0

        async def mock_post_json(url, payload):
            nonlocal call_count
            call_count += 1
            if call_count <= 2:
                return (502, {}, b"Bad Gateway")
            else:
                return (200, {}, b'{"choices":[{"message":{"content":"test response"}}]}')

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                result = await backend.generate(req)
                return result, call_count

        result, calls = asyncio.run(run_test())

        assert isinstance(result, MinimaLlmResponse)
        assert result.text == "test response"
        assert calls == 3  # Two failures, one success

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_limited_retries_fails_after_max(self, base_config):
        """With max_attempts=3, should fail after 3 attempts."""
        config = replace(base_config, max_attempts=3)
        backend = OpenAIMinimaLlm(config)

        async def mock_post_json(url, payload):
            return (502, {}, b"Bad Gateway")

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                result = await backend.generate(req)
                return result

        result = asyncio.run(run_test())

        assert isinstance(result, MinimaLlmFailure)
        assert result.attempts == 3
        assert "502" in result.message


class TestOverloadWarnings:
    """Test overload warning and recovery messages."""

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_overload_warning_printed_once(self, base_config, capsys):
        """Overload warning should be printed only once per request."""
        config = replace(base_config, max_attempts=5)
        backend = OpenAIMinimaLlm(config)

        call_count = 0

        async def mock_post_json(url, payload):
            nonlocal call_count
            call_count += 1
            if call_count <= 3:
                return (502, {}, b"Bad Gateway")
            else:
                return (200, {}, b'{"choices":[{"message":{"content":"ok"}}]}')

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                return await backend.generate(req)

        asyncio.run(run_test())

        captured = capsys.readouterr()
        assert captured.out.count("Server overload") == 1
        assert "HTTP 502" in captured.out
        assert "Ctrl-C" in captured.out

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_recovery_message_printed(self, base_config, capsys):
        """Recovery message should be printed after overload recovery."""
        config = replace(base_config, max_attempts=5)
        backend = OpenAIMinimaLlm(config)

        call_count = 0

        async def mock_post_json(url, payload):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return (502, {}, b"Bad Gateway")
            else:
                return (200, {}, b'{"choices":[{"message":{"content":"ok"}}]}')

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                return await backend.generate(req)

        asyncio.run(run_test())

        captured = capsys.readouterr()
        assert "Server recovered" in captured.out

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_no_warning_on_immediate_success(self, base_config, capsys):
        """No warning should be printed if request succeeds immediately."""
        config = replace(base_config, max_attempts=5)
        backend = OpenAIMinimaLlm(config)

        async def mock_post_json(url, payload):
            return (200, {}, b'{"choices":[{"message":{"content":"ok"}}]}')

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                return await backend.generate(req)

        asyncio.run(run_test())

        captured = capsys.readouterr()
        assert "Server overload" not in captured.out
        assert "Server recovered" not in captured.out


class TestCooldownOnOverload:
    """Test that 502 triggers cooldown bump."""

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_502_triggers_cooldown_bump(self, base_config):
        """502 should trigger cooldown.bump() call."""
        config = replace(
            base_config,
            max_attempts=2,
            cooldown_floor_s=0.1,
            cooldown_cap_s=1.0,
            cooldown_halflife_s=10.0,
        )
        backend = OpenAIMinimaLlm(config)

        bump_calls = []
        original_bump = backend._cooldown.bump

        async def tracking_bump(suggested_s):
            bump_calls.append(suggested_s)
            await original_bump(suggested_s)

        async def mock_post_json(url, payload):
            return (502, {}, b"Bad Gateway")

        async def run_test():
            backend._cooldown.bump = tracking_bump
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                req = MinimaLlmRequest(
                    request_id="test-1",
                    messages=[{"role": "user", "content": "test"}],
                )
                return await backend.generate(req)

        asyncio.run(run_test())

        # Should have bumped cooldown for each 502 response
        assert len(bump_calls) == 2  # max_attempts=2


class TestDspyAdapterInfiniteRetries:
    """Test DSPy adapter behavior with max_attempts=0.

    When max_attempts=0 (infinite HTTP retries), the DSPy adapter's
    parse_retry_limit should be set to a reasonable default (3),
    not 0 (which would cause range(0) = empty loop).
    """

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_parse_retry_limit_with_max_attempts_zero(self, base_config):
        """With max_attempts=0, parse_retry_limit should be 3, not 0."""
        dspy = pytest.importorskip("dspy")
        from pydantic import BaseModel
        from typing import Optional

        from minima_llm.dspy_adapter import run_dspy_batch, MinimaLlmDSPyLM

        config = replace(base_config, max_attempts=0)
        backend = OpenAIMinimaLlm(config)

        # Simple signature for testing
        class SimpleSignature(dspy.Signature):
            input_text: str = dspy.InputField()
            output_text: str = dspy.OutputField()

        class SimpleAnnotation(BaseModel):
            input_text: str
            output_text: Optional[str] = None

        def convert_output(pred, obj: SimpleAnnotation):
            obj.output_text = pred.output_text

        call_count = 0

        async def mock_acall(self, *args, **kwargs):
            nonlocal call_count
            call_count += 1
            # Return valid DSPy-formatted response
            return ["[[ ## reasoning ## ]]\nReasoning.\n[[ ## output_text ## ]]\nmocked_response"]

        async def run_test():
            with patch.object(MinimaLlmDSPyLM, 'acall', mock_acall):
                annotations = [SimpleAnnotation(input_text="test")]
                return await run_dspy_batch(
                    SimpleSignature,
                    annotations,
                    convert_output,
                    backend=backend
                )

        results = asyncio.run(run_test())

        # With old code (range(0)), call_count would be 0 and we'd get an error
        # With fix, call_count should be >= 1
        assert call_count >= 1, "process_one should run at least once with max_attempts=0"
        assert results[0].output_text == "mocked_response"

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_parse_errors_retry_with_max_attempts_zero(self, base_config):
        """Parse errors should still retry (up to 3 times) with max_attempts=0."""
        dspy = pytest.importorskip("dspy")
        from pydantic import BaseModel
        from typing import Optional

        from minima_llm.dspy_adapter import run_dspy_batch, MinimaLlmDSPyLM

        config = replace(base_config, max_attempts=0)
        backend = OpenAIMinimaLlm(config)

        class SimpleSignature(dspy.Signature):
            input_text: str = dspy.InputField()
            output_text: str = dspy.OutputField()

        class SimpleAnnotation(BaseModel):
            input_text: str
            output_text: Optional[str] = None

        def convert_output(pred, obj: SimpleAnnotation):
            obj.output_text = pred.output_text

        call_count = 0

        async def mock_acall(self, *args, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count < 3:
                # Return malformed response to trigger parse error retry
                return ["malformed response without proper headers"]
            else:
                # Return valid response on 3rd try
                return ["[[ ## reasoning ## ]]\nOK.\n[[ ## output_text ## ]]\nsuccess"]

        async def run_test():
            with patch.object(MinimaLlmDSPyLM, 'acall', mock_acall):
                annotations = [SimpleAnnotation(input_text="test")]
                return await run_dspy_batch(
                    SimpleSignature,
                    annotations,
                    convert_output,
                    backend=backend
                )

        results = asyncio.run(run_test())

        # Should have retried parse errors and eventually succeeded
        assert call_count == 3, "Should retry parse errors up to 3 times"
        assert results[0].output_text == "success"

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_infinite_http_retries_through_dspy_adapter(self, base_config):
        """With max_attempts=0, HTTP 502 errors should retry until success.

        This tests the full path: DSPy adapter -> generate() -> HTTP retries.
        """
        dspy = pytest.importorskip("dspy")
        from pydantic import BaseModel
        from typing import Optional

        from minima_llm.dspy_adapter import run_dspy_batch, MinimaLlmDSPyLM

        config = replace(base_config, max_attempts=0)
        backend = OpenAIMinimaLlm(config)

        class SimpleSignature(dspy.Signature):
            input_text: str = dspy.InputField()
            output_text: str = dspy.OutputField()

        class SimpleAnnotation(BaseModel):
            input_text: str
            output_text: Optional[str] = None

        def convert_output(pred, obj: SimpleAnnotation):
            obj.output_text = pred.output_text

        http_call_count = 0
        num_502_before_success = 10  # Simulate server down for 10 requests

        async def mock_post_json(url, payload):
            nonlocal http_call_count
            http_call_count += 1
            if http_call_count <= num_502_before_success:
                return (502, {}, b"Bad Gateway")
            else:
                # ChainOfThought format with reasoning
                response = {
                    "choices": [{
                        "message": {
                            "content": "[[ ## reasoning ## ]]\nDone.\n[[ ## output_text ## ]]\nrecovered"
                        }
                    }]
                }
                import json
                return (200, {}, json.dumps(response).encode())

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                annotations = [SimpleAnnotation(input_text="test")]
                return await run_dspy_batch(
                    SimpleSignature,
                    annotations,
                    convert_output,
                    backend=backend
                )

        results = asyncio.run(run_test())

        # Should have retried 10 times (502) then succeeded on 11th
        assert http_call_count == num_502_before_success + 1
        assert results[0].output_text == "recovered"

    @pytest.mark.skipif(
        skip_llm_tests(),
        reason="Skipping LLM endpoint tests (SKIP_LLM_ENDPOINT_TESTS set)"
    )
    def test_limited_retries_fails_through_dspy_adapter(self, base_config):
        """With max_attempts=3, HTTP errors should fail after 3 attempts.

        Contrast with infinite retries test above.
        """
        dspy = pytest.importorskip("dspy")
        from pydantic import BaseModel
        from typing import Optional

        from minima_llm.dspy_adapter import run_dspy_batch

        config = replace(base_config, max_attempts=3)
        backend = OpenAIMinimaLlm(config)

        class SimpleSignature(dspy.Signature):
            input_text: str = dspy.InputField()
            output_text: str = dspy.OutputField()

        class SimpleAnnotation(BaseModel):
            input_text: str
            output_text: Optional[str] = None

        def convert_output(pred, obj: SimpleAnnotation):
            obj.output_text = pred.output_text

        http_call_count = 0

        async def mock_post_json(url, payload):
            nonlocal http_call_count
            http_call_count += 1
            return (502, {}, b"Bad Gateway")  # Always fail

        async def run_test():
            with patch.object(backend, '_post_json', side_effect=mock_post_json):
                annotations = [SimpleAnnotation(input_text="test")]
                return await run_dspy_batch(
                    SimpleSignature,
                    annotations,
                    convert_output,
                    backend=backend
                )

        # Should raise after exhausting retries
        with pytest.raises(RuntimeError) as exc_info:
            asyncio.run(run_test())

        # Should have tried exactly max_attempts times per parse retry (3 * 3 = 9)
        # generate() retries 3 times, process_one() retries 3 times
        assert http_call_count >= 9
        assert http_call_count <= 20
        assert "502" in str(exc_info.value) or "failed" in str(exc_info.value).lower()
