"""Tests for LLM provider streaming functionality."""

import json
from typing import Generator
from unittest.mock import Mock, patch

import pytest

from thetable_poc.config.llm_provider import LLMProvider, OllamaProvider, OpenRouterProvider


class TestLLMProviderAbstract:
    """Test abstract LLMProvider interface."""

    def test_chat_stream_is_abstract_method(self):
        """Test that chat_stream is defined as an abstract method."""
        # LLMProvider should not be instantiable
        with pytest.raises(TypeError):
            LLMProvider()

    def test_chat_stream_returns_generator(self):
        """Test that chat_stream returns a Generator."""
        # Create a mock provider that implements chat_stream
        class MockProvider(LLMProvider):
            def chat(self, system: str, user: str, temperature: float = 0.7) -> str:
                return "test response"

            def chat_stream(
                self, system: str, user: str, temperature: float = 0.7
            ) -> Generator[str, None, None]:
                yield "chunk1"
                yield "chunk2"

        provider = MockProvider()
        result = provider.chat_stream("system", "user")

        # Check that it's a generator
        assert isinstance(result, Generator)

        # Check that it yields strings
        chunks = list(result)
        assert chunks == ["chunk1", "chunk2"]


class TestOllamaProviderStreaming:
    """Test OllamaProvider streaming implementation."""

    def test_ollama_chat_stream_parses_ndjson(self):
        """Test that chat_stream correctly parses NDJSON responses."""
        provider = OllamaProvider(model="llama2")

        # Mock response with NDJSON format
        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'{"response": "Hello", "done": false}',
            b'{"response": " world", "done": false}',
            b'{"response": "", "done": true}',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert chunks == ["Hello", " world"]

    def test_ollama_chat_stream_yields_response_field(self):
        """Test that chat_stream extracts 'response' field from NDJSON."""
        provider = OllamaProvider(model="llama2")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'{"response": "Test", "done": false}',
            b'{"response": " chunk", "done": false}',
            b'{"response": "", "done": true}',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert "Test" in chunks
        assert " chunk" in chunks

    def test_ollama_chat_stream_handles_done_signal(self):
        """Test that chat_stream stops when done: true is received."""
        provider = OllamaProvider(model="llama2")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'{"response": "First", "done": false}',
            b'{"response": "", "done": true}',
            b'{"response": "Should not appear", "done": false}',  # After done
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert "First" in chunks
        assert "Should not appear" not in chunks

    def test_ollama_chat_stream_handles_empty_chunks(self):
        """Test that empty response chunks are filtered out."""
        provider = OllamaProvider(model="llama2")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'{"response": "Hello", "done": false}',
            b'{"response": "", "done": false}',  # Empty
            b'{"response": "World", "done": false}',
            b'{"response": "", "done": true}',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        # Empty chunks should not be yielded
        assert chunks == ["Hello", "World"]


class TestOpenRouterProviderStreaming:
    """Test OpenRouterProvider streaming implementation."""

    def test_openrouter_chat_stream_parses_sse(self):
        """Test that chat_stream correctly parses SSE format."""
        provider = OpenRouterProvider(model="openai/gpt-4o-mini", api_key="test-key")

        # Mock SSE response
        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: {"choices": [{"delta": {"content": "Hello"}}]}',
            b'data: {"choices": [{"delta": {"content": " world"}}]}',
            b'data: [DONE]',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert chunks == ["Hello", " world"]

    def test_openrouter_chat_stream_extracts_delta(self):
        """Test that chat_stream extracts delta.content from SSE."""
        provider = OpenRouterProvider(model="openai/gpt-4o-mini", api_key="test-key")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: {"choices": [{"delta": {"content": "Test"}}]}',
            b'data: {"choices": [{"delta": {"content": " chunk"}}]}',
            b'data: [DONE]',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert "Test" in chunks
        assert " chunk" in chunks

    def test_openrouter_chat_stream_handles_done_signal(self):
        """Test that chat_stream stops at [DONE] signal."""
        provider = OpenRouterProvider(model="openai/gpt-4o-mini", api_key="test-key")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: {"choices": [{"delta": {"content": "First"}}]}',
            b'data: [DONE]',
            b'data: {"choices": [{"delta": {"content": "Should not appear"}}]}',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert "First" in chunks
        assert "Should not appear" not in chunks

    def test_openrouter_chat_stream_handles_empty_delta(self):
        """Test that empty delta content is filtered out."""
        provider = OpenRouterProvider(model="openai/gpt-4o-mini", api_key="test-key")

        mock_response = Mock()
        mock_response.iter_lines.return_value = [
            b'data: {"choices": [{"delta": {"content": "Hello"}}]}',
            b'data: {"choices": [{"delta": {}}]}',  # No content
            b'data: {"choices": [{"delta": {"content": "World"}}]}',
            b'data: [DONE]',
        ]
        mock_response.raise_for_status = Mock()

        with patch("requests.post", return_value=mock_response):
            chunks = list(provider.chat_stream("system", "user"))

        assert chunks == ["Hello", "World"]

    def test_openrouter_chat_stream_requires_api_key(self):
        """Test that chat_stream raises error if API key is missing."""
        # Mock environment to ensure no API key from env
        with patch.dict("os.environ", {}, clear=True):
            provider = OpenRouterProvider(model="openai/gpt-4o-mini", api_key=None)

            # Generator is lazy - need to actually iterate to trigger the error
            with pytest.raises(ValueError, match="API key required"):
                gen = provider.chat_stream("system", "user")
                next(gen)  # Force generator execution
