"""
Streaming Response Integration Tests

Tests end-to-end streaming response functionality, including handling of synchronous and asynchronous generators, Server-Sent Events (SSE) format validation, etc.
"""

import asyncio
import json
import threading
import time
import requests
import pytest
from typing import Generator, AsyncGenerator

from ppio_sandbox.agent_runtime.runtime.app import AgentRuntimeApp
from ppio_sandbox.agent_runtime.runtime.context import RequestContext
from ppio_sandbox.agent_runtime.runtime.models import RuntimeConfig


class TestStreamingBasicFlow:
    """Streaming response basic flow tests"""
    
    def _parse_sse_events(self, response) -> list:
        """Parse Server-Sent Events response"""
        sse_events = []
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                if line_str.startswith('data: '):
                    data_content = line_str[6:]  # Remove "data: " prefix
                    try:
                        # Content is JSON encoded, needs to be decoded
                        event_data = json.loads(data_content)
                        sse_events.append(event_data)
                    except json.JSONDecodeError:
                        sse_events.append(data_content)
        return sse_events
    
    def setup_method(self, method):
        """Setup before each test method"""
        # Assign different ports for each test method
        port_map = {
            "test_sync_generator_streaming_e2e": 8920,
            "test_async_generator_streaming_e2e": 8921,
            "test_large_data_streaming_e2e": 8922,
            "test_streaming_with_context_e2e": 8923
        }
        self.test_port = port_map.get(method.__name__, 8924)
        self.base_url = f"http://localhost:{self.test_port}"
        self.server_thread = None
    
    def teardown_method(self):
        """Cleanup after each test method"""
        # Clean up any running servers
        if self.server_thread and self.server_thread.is_alive():
            pass
    
    def start_server_in_thread(self, app: AgentRuntimeApp):
        """Start server in a separate thread"""
        def run_server():
            try:
                app.run(port=self.test_port, host="127.0.0.1")
            except Exception as e:
                print(f"Server error: {e}")
        
        self.server_thread = threading.Thread(target=run_server, daemon=True)
        self.server_thread.start()
        
        # Wait for server to start
        self._wait_for_server_ready()
    
    def _wait_for_server_ready(self, timeout=5):
        """Wait for server to be ready"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            try:
                response = requests.get(f"{self.base_url}/ping", timeout=1)
                if response.status_code == 200:
                    return
            except requests.exceptions.RequestException:
                pass
            time.sleep(0.1)
        raise TimeoutError("Server did not start within timeout")
    
    @pytest.mark.integration
    def test_sync_generator_streaming_e2e(self):
        """Test end-to-end streaming response with synchronous generator"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        def streaming_agent(request: dict) -> Generator[str, None, None]:
            data = request.get('data', {})
            prompt = data.get('prompt', 'test')
            count = data.get('count', 3)
            
            for i in range(count):
                yield f"Chunk {i+1}/{count}: Processing '{prompt}'"
                time.sleep(0.05)  # Simulate processing time
            yield f"Final: Completed processing '{prompt}'"
        
        self.start_server_in_thread(app)
        
        request_data = {
            "data": {"prompt": "streaming test", "count": 3},
            "sandbox_id": "streaming-test",
            "stream": True  # Enable streaming response
        }
        
        # Send streaming request
        start_time = time.time()
        response = requests.post(
            f"{self.base_url}/invocations",
            json=request_data,
            headers={"Content-Type": "application/json"},
            stream=True  # Enable streaming reception
        )
        end_time = time.time()
        
        assert response.status_code == 200
        assert response.headers.get("content-type") == "text/plain; charset=utf-8"
        
        # Verify streaming data (SSE format)
        sse_events = self._parse_sse_events(response)
        
        assert len(sse_events) == 4  # 3 processing chunks + 1 final chunk
        assert "Chunk 1/3: Processing 'streaming test'" in sse_events[0]
        assert "Chunk 2/3: Processing 'streaming test'" in sse_events[1]
        assert "Chunk 3/3: Processing 'streaming test'" in sse_events[2]
        assert "Final: Completed processing 'streaming test'" in sse_events[3]
        
        # Streaming response time assertion is relatively loose, as network transmission can be fast
        # Main verification is that chunked data was actually received
        assert len(sse_events) > 1  # Ensure it was transmitted in chunks
    
    @pytest.mark.integration
    def test_async_generator_streaming_e2e(self):
        """Test end-to-end streaming response with asynchronous generator"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        async def async_streaming_agent(request: dict) -> AsyncGenerator[str, None]:
            data = request.get('data', {})
            prompt = data.get('prompt', 'test')
            count = data.get('count', 3)
            
            for i in range(count):
                await asyncio.sleep(0.05)  # Async delay
                yield f"AsyncChunk {i+1}/{count}: '{prompt}' processed"
            yield f"AsyncFinal: All '{prompt}' processing complete"
        
        self.start_server_in_thread(app)
        
        request_data = {
            "data": {"prompt": "async streaming", "count": 3},
            "sandbox_id": "async-streaming-test",
            "stream": True
        }
        
        start_time = time.time()
        response = requests.post(
            f"{self.base_url}/invocations",
            json=request_data,
            headers={"Content-Type": "application/json"},
            stream=True
        )
        end_time = time.time()
        
        assert response.status_code == 200
        
        # Verify async streaming data (SSE format)
        sse_events = self._parse_sse_events(response)
        
        assert len(sse_events) == 4
        assert "AsyncChunk 1/3: 'async streaming' processed" in sse_events[0]
        assert "AsyncFinal: All 'async streaming' processing complete" in sse_events[3]
        
        # Ensure it was transmitted in chunks
        assert len(sse_events) > 1
    
    @pytest.mark.integration
    def test_large_data_streaming_e2e(self):
        """Test large data streaming transfer"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        def large_data_agent(request: dict) -> Generator[str, None, None]:
            data = request.get('data', {})
            size = data.get('chunk_size', 1024)
            count = data.get('chunk_count', 5)
            
            for i in range(count):
                # Generate data chunks of specified size
                chunk_data = "x" * size
                yield f"Chunk_{i+1}_Size_{size}: {chunk_data}"
        
        self.start_server_in_thread(app)
        
        request_data = {
            "data": {"chunk_size": 512, "chunk_count": 5},
            "sandbox_id": "large-data-test",
            "stream": True
        }
        
        response = requests.post(
            f"{self.base_url}/invocations",
            json=request_data,
            stream=True
        )
        
        assert response.status_code == 200
        
        # Verify large data stream (SSE format)
        sse_events = self._parse_sse_events(response)
        
        assert len(sse_events) == 5
        # Verify data integrity
        for i, event_data in enumerate(sse_events):
            assert f"Chunk_{i+1}_Size_512" in event_data
            assert "x" * 512 in event_data
    
    @pytest.mark.integration
    def test_streaming_with_context_e2e(self):
        """Test streaming response with context"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        def context_streaming_agent(request: dict, context: RequestContext) -> Generator[str, None, None]:
            data = request.get('data', {})
            message = data.get('message', 'test')
            
            # First chunk contains context information
            yield f"Context: sandbox_id={context.sandbox_id}, request_id={context.request_id}"
            
            # Processing data chunks
            for i in range(3):
                yield f"Processing_{i+1}: {message} in {context.sandbox_id}"
            
            # Final chunk
            yield f"Completed: {message} processing in {context.sandbox_id}"
        
        self.start_server_in_thread(app)
        
        request_data = {
            "data": {"message": "context streaming test"},
            "sandbox_id": "context-streaming-sandbox",
            "stream": True
        }
        
        response = requests.post(
            f"{self.base_url}/invocations",
            json=request_data,
            headers={
                "Content-Type": "application/json",
                "User-Agent": "StreamingTestClient/1.0"
            },
            stream=True
        )
        
        assert response.status_code == 200
        
        sse_events = self._parse_sse_events(response)
        
        assert len(sse_events) == 5  # 1 context + 3 processing + 1 completion
        
        # Verify context information
        context_event = sse_events[0]
        assert "sandbox_id=context-streaming-sandbox" in context_event
        assert "request_id=" in context_event
        
        # Verify processing chunks
        for i in range(1, 4):
            processing_event = sse_events[i]
            assert f"Processing_{i}: context streaming test in context-streaming-sandbox" == processing_event
        
        # Verify completion chunk
        final_event = sse_events[4]
        assert "Completed: context streaming test processing in context-streaming-sandbox" == final_event


class TestStreamingErrorHandling:
    """Streaming response error handling tests"""
    
    def _parse_sse_events(self, response) -> list:
        """Parse Server-Sent Events response"""
        sse_events = []
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                if line_str.startswith('data: '):
                    data_content = line_str[6:]
                    try:
                        event_data = json.loads(data_content)
                        sse_events.append(event_data)
                    except json.JSONDecodeError:
                        sse_events.append(data_content)
        return sse_events
    
    def setup_method(self, method):
        """Setup"""
        port_map = {
            "test_streaming_error_handling_e2e": 8930,
            "test_streaming_interruption_e2e": 8931,
        }
        self.test_port = port_map.get(method.__name__, 8932)
        self.base_url = f"http://localhost:{self.test_port}"
        self.server_thread = None
    
    def teardown_method(self):
        """Cleanup"""
        pass
    
    def start_server_in_thread(self, app: AgentRuntimeApp):
        """Start server"""
        def run_server():
            try:
                app.run(port=self.test_port, host="127.0.0.1")
            except Exception as e:
                print(f"Server error: {e}")
        
        self.server_thread = threading.Thread(target=run_server, daemon=True)
        self.server_thread.start()
        self._wait_for_server_ready()
    
    def _wait_for_server_ready(self, timeout=5):
        """Wait for server to be ready"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            try:
                response = requests.get(f"{self.base_url}/ping", timeout=1)
                if response.status_code == 200:
                    return
            except requests.exceptions.RequestException:
                pass
            time.sleep(0.1)
        raise TimeoutError("Server did not start within timeout")
    
    @pytest.mark.integration
    def test_streaming_error_handling_e2e(self):
        """Test error handling in streaming responses"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        def error_streaming_agent(request: dict) -> Generator[str, None, None]:
            data = request.get('data', {})
            should_error = data.get('should_error', False)
            
            yield "Chunk 1: Starting processing"
            yield "Chunk 2: Processing continues"
            
            if should_error:
                raise ValueError("Error during streaming")
            
            yield "Chunk 3: Processing completed successfully"
        
        self.start_server_in_thread(app)
        
        # Test normal streaming processing
        normal_request = {
            "data": {"should_error": False},
            "sandbox_id": "normal-streaming",
            "stream": True
        }
        
        response = requests.post(f"{self.base_url}/invocations", json=normal_request, stream=True)
        
        assert response.status_code == 200
        
        chunks = []
        for line in response.iter_lines():
            if line:
                chunks.append(line.decode('utf-8'))
        
        assert len(chunks) == 3
        assert "Chunk 1: Starting processing" in chunks[0]
        assert "Chunk 3: Processing completed successfully" in chunks[2]
        
        # Test error during streaming processing
        error_request = {
            "data": {"should_error": True},
            "sandbox_id": "error-streaming",
            "stream": True
        }
        
        response = requests.post(f"{self.base_url}/invocations", json=error_request, stream=True)
        
        # Errors in streaming responses usually occur after transmission has started, so status code may still be 200
        # But should see partial data, then connection will be interrupted
        chunks = []
        try:
            for line in response.iter_lines():
                if line:
                    chunks.append(line.decode('utf-8'))
        except requests.exceptions.ChunkedEncodingError:
            # This is expected, because an error occurred during streaming
            pass
        
        # Should receive at least the first two chunks
        assert len(chunks) >= 2
        assert "Chunk 1: Starting processing" in chunks[0]
        assert "Chunk 2: Processing continues" in chunks[1]
    
    @pytest.mark.integration
    def test_non_streaming_vs_streaming_comparison(self):
        """Test comparison between non-streaming and streaming responses"""
        config = RuntimeConfig()
        app = AgentRuntimeApp(config=config)
        
        @app.entrypoint
        def flexible_agent(request: dict) -> Generator[str, None, None]:
            data = request.get('data', {})
            message = data.get('message', 'test')
            
            for i in range(3):
                yield f"Part {i+1}: {message}"
        
        self.start_server_in_thread(app)
        
        # Test non-streaming request (stream=False)
        non_streaming_request = {
            "data": {"message": "non-streaming"},
            "sandbox_id": "non-streaming-test",
            "stream": False
        }
        
        response = requests.post(f"{self.base_url}/invocations", json=non_streaming_request)
        
        assert response.status_code == 200
        assert response.headers.get("content-type") == "application/json"
        
        # Non-streaming response should be in JSON format
        data = response.json()
        assert data["status"] == "success"
        # Generator results should be collected into an array
        # Note: specific behavior depends on server implementation
        
        # Test streaming request (stream=True)
        streaming_request = {
            "data": {"message": "streaming"},
            "sandbox_id": "streaming-test",
            "stream": True
        }
        
        response = requests.post(f"{self.base_url}/invocations", json=streaming_request, stream=True)
        
        assert response.status_code == 200
        assert response.headers.get("content-type") == "text/plain; charset=utf-8"
        
        # Streaming response should be a text stream
        chunks = []
        for line in response.iter_lines():
            if line:
                chunks.append(line.decode('utf-8'))
        
        assert len(chunks) == 3
        assert "Part 1: streaming" in chunks[0]
        assert "Part 2: streaming" in chunks[1]
        assert "Part 3: streaming" in chunks[2]