#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Combined coverage test for fallback.py in OneLLM.

This test specifically focuses on the remaining uncovered lines:
- Line 150: Error handling in streaming generator
- Line 178: FallbackExhaustionError in streaming from max_fallbacks
- Line 207: Chat streaming generator creation
- Lines 231-246: Completion streaming setup
- Line 303: Completion streaming fallback
"""

import pytest
from unittest import mock
import asyncio

from onellm.errors import APIError, FallbackExhaustionError
from onellm.providers.fallback import FallbackProviderProxy
from onellm.providers.base import Provider
from onellm.utils.fallback import FallbackConfig


class CombinedCoverageProvider(Provider):
    """Provider for comprehensive coverage testing."""

    def __init__(self, name="test", behavior="normal"):
        """
        Initialize provider with configurable behavior.

        Args:
            name: Provider name for error messages
            behavior: One of:
                - "normal": Returns successful responses
                - "error": Raises API errors for all requests
                - "fail_stream_init": Raises API errors for streaming requests
                - "fail_during_stream": Returns a generator that fails during iteration
        """
        self.name = name
        self.behavior = behavior
        self.calls = []

    async def create_chat_completion(self, messages, model, stream=False, **kwargs):
        """Mock chat completion with configurable behaviors."""
        self.calls.append(("create_chat_completion", model, stream, kwargs))

        # Handle error behavior
        if self.behavior == "error":
            raise APIError(f"{self.name} error")

        # Special behavior for streaming
        if stream:
            if self.behavior == "fail_stream_init":
                raise APIError(f"{self.name} streaming initialization error")

            if self.behavior == "fail_during_stream":
                # Return a generator that raises during iteration
                async def failing_generator():
                    # Yield one chunk first
                    yield {"choices": [{"delta": {"content": "First chunk"}}]}
                    # Then raise on the second iteration
                    raise APIError(f"{self.name} streaming iteration error")
                return failing_generator()

            # Normal streaming
            async def normal_generator():
                for i in range(3):
                    yield {"choices": [{"delta": {"content": f"Chunk {i}"}}]}
                    # Small delay to simulate real streaming
                    await asyncio.sleep(0.01)
            return normal_generator()

        # Normal response
        return {
            "choices": [{"message": {"content": f"{self.name} response"}}]
        }

    async def create_completion(self, prompt, model, stream=False, **kwargs):
        """Mock completion with configurable behaviors."""
        self.calls.append(("create_completion", model, stream, kwargs))

        # Handle error behavior
        if self.behavior == "error":
            raise APIError(f"{self.name} error")

        # Special behavior for streaming
        if stream:
            if self.behavior == "fail_stream_init":
                raise APIError(f"{self.name} streaming initialization error")

            if self.behavior == "fail_during_stream":
                # Return a generator that raises during iteration
                async def failing_generator():
                    # Yield one chunk first
                    yield {"choices": [{"text": "First chunk"}]}
                    # Then raise on the second iteration
                    raise APIError(f"{self.name} streaming iteration error")
                return failing_generator()

            # Normal streaming
            async def normal_generator():
                for i in range(3):
                    yield {"choices": [{"text": f"Chunk {i}"}]}
                    # Small delay to simulate real streaming
                    await asyncio.sleep(0.01)
            return normal_generator()

        # Normal response
        return {
            "choices": [{"text": f"{self.name} response"}]
        }

    # Implement other required abstract methods
    async def create_embedding(self, input, model, **kwargs):
        self.calls.append(("create_embedding", model, kwargs))
        if self.behavior == "error":
            raise APIError(f"{self.name} error")
        return {"data": [{"embedding": [0.1, 0.2, 0.3]}]}

    async def upload_file(self, file, purpose, **kwargs):
        self.calls.append(("upload_file", purpose, kwargs))
        if self.behavior == "error":
            raise APIError(f"{self.name} error")
        return {"id": "file-123"}

    async def download_file(self, file_id, **kwargs):
        self.calls.append(("download_file", file_id, kwargs))
        if self.behavior == "error":
            raise APIError(f"{self.name} error")
        return b"content"

    async def create_speech(self, input, model, **kwargs):
        self.calls.append(("create_speech", model, kwargs))
        if self.behavior == "error":
            raise APIError(f"{self.name} error")
        return b"audio"

    async def create_image(self, prompt, model, **kwargs):
        self.calls.append(("create_image", model, kwargs))
        if self.behavior == "error":
            raise APIError(f"{self.name} error")
        return {"data": [{"url": "https://example.com/image.png"}]}


class TestCombinedCoverage:
    """Tests targeting all remaining uncovered lines in fallback.py."""

    def setup_method(self):
        """Set up test environment."""
        # Patch the get_provider function
        self.mock_get_provider = mock.patch("onellm.providers.fallback.get_provider").start()

    def teardown_method(self):
        """Clean up test environment."""
        mock.patch.stopall()

    @pytest.mark.xfail(reason="Test is designed to fail with APIError")
    @pytest.mark.asyncio
    async def test_fail_during_streaming_line_150(self):
        """Test streaming generator that fails during iteration (line 150)."""
        # Create providers
        provider1 = CombinedCoverageProvider("primary", "fail_during_stream")
        provider2 = CombinedCoverageProvider("fallback", "normal")

        # Configure get_provider
        def get_provider_side_effect(provider_name):
            if provider_name == "provider1":
                return provider1
            elif provider_name == "provider2":
                return provider2
            raise ValueError(f"Unknown provider: {provider_name}")

        self.mock_get_provider.side_effect = get_provider_side_effect

        # Create fallback proxy
        proxy = FallbackProviderProxy(
            ["provider1/model1", "provider2/model2"],
            FallbackConfig(retriable_errors=[APIError])
        )

        # Call the streaming method
        generator = await proxy.create_chat_completion(
            messages=[{"role": "user", "content": "Hello"}],
            stream=True
        )

        # Iterate through the generator - should switch to fallback
        chunks = []
        async for chunk in generator:
            chunks.append(chunk)

        # Verify fallback provider was used
        assert len(chunks) == 3
        assert provider1.calls[0][0] == "create_chat_completion"
        assert provider1.calls[0][2] is True  # stream=True
        assert provider2.calls[0][0] == "create_chat_completion"
        assert provider2.calls[0][2] is True  # stream=True

    @pytest.mark.asyncio
    async def test_streaming_max_fallbacks_line_178(self):
        """Test streaming with max_fallbacks exhaustion (line 178)."""
        # Create failing providers
        provider1 = CombinedCoverageProvider("provider1", "fail_stream_init")
        provider2 = CombinedCoverageProvider("provider2", "fail_stream_init")
        provider3 = CombinedCoverageProvider("provider3", "normal")

        # Configure get_provider
        def get_provider_side_effect(provider_name):
            if provider_name == "provider1":
                return provider1
            elif provider_name == "provider2":
                return provider2
            elif provider_name == "provider3":
                return provider3
            raise ValueError(f"Unknown provider: {provider_name}")

        self.mock_get_provider.side_effect = get_provider_side_effect

        # Create fallback proxy with max_fallbacks=1
        proxy = FallbackProviderProxy(
            ["provider1/model1", "provider2/model2", "provider3/model3"],
            FallbackConfig(
                retriable_errors=[APIError],
                max_fallbacks=1  # Only try one fallback
            )
        )

        # Should raise FallbackExhaustionError after trying 2 providers
        with pytest.raises(FallbackExhaustionError) as excinfo:
            await proxy.create_chat_completion(
                messages=[{"role": "user", "content": "Hello"}],
                stream=True
            )

        # Verify error and call counts
        assert "All models failed" in str(excinfo.value)
        assert len(provider1.calls) == 1
        assert len(provider2.calls) == 1
        assert len(provider3.calls) == 0  # Should not be called due to max_fallbacks=1

    @pytest.mark.asyncio
    async def test_chat_streaming_fallback_line_207(self):
        """Test chat completion streaming generator creation (line 207)."""
        # Create providers
        primary = CombinedCoverageProvider("primary", "error")
        fallback = CombinedCoverageProvider("fallback", "normal")

        # Configure get_provider
        def get_provider_side_effect(provider_name):
            if provider_name == "primary":
                return primary
            elif provider_name == "fallback":
                return fallback
            raise ValueError(f"Unknown provider: {provider_name}")

        self.mock_get_provider.side_effect = get_provider_side_effect

        # Create fallback proxy
        proxy = FallbackProviderProxy(
            ["primary/model1", "fallback/model2"],
            FallbackConfig(retriable_errors=[APIError])
        )

        # Call the streaming method
        generator = await proxy.create_chat_completion(
            messages=[{"role": "user", "content": "Hello"}],
            stream=True
        )

        # Collect chunks to verify fallback worked
        chunks = []
        async for chunk in generator:
            chunks.append(chunk)

        # Verify fallback provider was used
        assert len(chunks) == 3
        assert primary.calls[0][0] == "create_chat_completion"
        assert fallback.calls[0][0] == "create_chat_completion"

    @pytest.mark.asyncio
    async def test_completion_streaming_lines_231_246_303(self):
        """Test completion streaming (lines 231-246, 303)."""
        # Create providers
        primary = CombinedCoverageProvider("primary", "error")
        fallback = CombinedCoverageProvider("fallback", "normal")

        # Configure get_provider
        def get_provider_side_effect(provider_name):
            if provider_name == "primary":
                return primary
            elif provider_name == "fallback":
                return fallback
            raise ValueError(f"Unknown provider: {provider_name}")

        self.mock_get_provider.side_effect = get_provider_side_effect

        # Create fallback proxy
        proxy = FallbackProviderProxy(
            ["primary/model1", "fallback/model2"],
            FallbackConfig(retriable_errors=[APIError])
        )

        # Call the streaming method
        generator = await proxy.create_completion(
            prompt="Hello",
            stream=True
        )

        # Collect chunks to verify fallback worked
        chunks = []
        async for chunk in generator:
            chunks.append(chunk)

        # Verify fallback provider was used
        assert len(chunks) == 3
        assert primary.calls[0][0] == "create_completion"
        assert fallback.calls[0][0] == "create_completion"
