"""
Tests for WebSocket functionality and real-time features
"""

import pytest
import asyncio
import json
import websockets
from aiohttp import web, WSMsgType
from plua.runtime import LuaAsyncRuntime


class TestWebSocketFunctionality:
    """Test cases for WebSocket operations"""
    
    async def create_websocket_server(self, port=9005):
        """Create a WebSocket server for testing"""
        
        async def websocket_handler(request):
            ws = web.WebSocketResponse()
            await ws.prepare(request)
            
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        # Echo back with server prefix
                        response = {
                            "type": "echo",
                            "original": data,
                            "server_message": f"Server received: {data.get('message', 'no message')}"
                        }
                        await ws.send_str(json.dumps(response))
                    except json.JSONDecodeError:
                        # Echo plain text
                        await ws.send_str(f"Server echo: {msg.data}")
                elif msg.type == WSMsgType.ERROR:
                    print(f'WebSocket error: {ws.exception()}')
                    break
            
            return ws
        
        app = web.Application()
        app.router.add_get('/ws', websocket_handler)
        
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, 'localhost', port)
        await site.start()
        
        return runner
    
    @pytest.mark.asyncio
    async def test_websocket_basic_connection(self):
        """Test basic WebSocket connection and communication"""
        # Start WebSocket server
        server = await self.create_websocket_server(9005)
        await asyncio.sleep(0.1)  # Give server time to start
        
        try:
            runtime = LuaAsyncRuntime()
            runtime.initialize_lua()
            
            # Test WebSocket connection
            script = """
            local ws_connected = false
            local message_received = false
            local received_data = nil
            local connection_error = false
            
            -- Check if WebSocket support is available
            if net.WebSocket then
                local ws = net.WebSocket()
                ws:connect("ws://localhost:9005/ws", {
                    onopen = function()
                        ws_connected = true
                        -- Send a test message
                        ws:send(json.encode({message = "Hello WebSocket Server"}))
                    end,
                    onmessage = function(data)
                        message_received = true
                        received_data = data
                    end,
                    onerror = function(error)
                        connection_error = true
                    end,
                    onclose = function()
                        -- Connection closed
                    end
                })
            else
                -- WebSocket not supported yet
                connection_error = true
            end
            
            return net.WebSocket ~= nil
            """
            
            websocket_supported = runtime.interpreter.lua.execute(script)
            
            if websocket_supported:
                # Wait for WebSocket communication
                await asyncio.sleep(1.0)
                
                ws_connected = runtime.interpreter.lua.eval("ws_connected")
                message_received = runtime.interpreter.lua.eval("message_received")
                received_data = runtime.interpreter.lua.eval("received_data")
                connection_error = runtime.interpreter.lua.eval("connection_error")
                
                print(f"WebSocket Test - Supported: {websocket_supported}, Connected: {ws_connected}, Message received: {message_received}")
                print(f"Received data: {received_data}")
                
                # Note: Actual WebSocket implementation may not be complete yet
                # This test documents the expected behavior
            else:
                print("WebSocket functionality not yet implemented - test documents expected API")
                
        finally:
            await server.cleanup()
    
    @pytest.mark.asyncio
    async def test_websocket_json_communication(self):
        """Test WebSocket with JSON message exchange"""
        server = await self.create_websocket_server(9006)
        await asyncio.sleep(0.1)
        
        try:
            runtime = LuaAsyncRuntime()
            runtime.initialize_lua()
            
            script = """
            local messages_sent = 0
            local messages_received = 0
            local json_data_received = {}
            
            if net.WebSocket then
                local ws = net.WebSocket()
                ws:connect("ws://localhost:9006/ws", {
                    onopen = function()
                        -- Send multiple JSON messages
                        local messages = {
                            {type = "greeting", message = "Hello"},
                            {type = "data", values = {1, 2, 3, 4, 5}},
                            {type = "status", active = true, count = 42}
                        }
                        
                        for i, msg in ipairs(messages) do
                            ws:send(json.encode(msg))
                            messages_sent = messages_sent + 1
                        end
                    end,
                    onmessage = function(data)
                        messages_received = messages_received + 1
                        local parsed = json.decode(data)
                        if parsed then
                            json_data_received[messages_received] = parsed
                        end
                    end
                })
                
                return true
            else
                return false
            end
            """
            
            websocket_supported = runtime.interpreter.lua.execute(script)
            
            if websocket_supported:
                await asyncio.sleep(1.5)
                
                messages_sent = runtime.interpreter.lua.eval("messages_sent")
                messages_received = runtime.interpreter.lua.eval("messages_received")
                
                print(f"WebSocket JSON Test - Sent: {messages_sent}, Received: {messages_received}")
            else:
                print("WebSocket JSON Test - Not supported yet")
                
        finally:
            await server.cleanup()
    
    @pytest.mark.asyncio
    async def test_realtime_event_simulation(self):
        """Test real-time event simulation with timers and callbacks"""
        runtime = LuaAsyncRuntime()
        runtime.initialize_lua()
        
        script = """
        local events_processed = 0
        local event_log = {}
        local simulation_active = true
        
        -- Simulate real-time events with different intervals
        local event_types = {
            {name = "heartbeat", interval = 100},
            {name = "sensor_data", interval = 200},
            {name = "status_update", interval = 500}
        }
        
        for i, event_type in ipairs(event_types) do
            local function generate_event()
                if simulation_active then
                    events_processed = events_processed + 1
                    event_log[events_processed] = {
                        type = event_type.name,
                        timestamp = os.clock(),
                        data = "Event " .. events_processed
                    }
                    
                    -- Schedule next event
                    setTimeout(generate_event, event_type.interval)
                end
            end
            
            -- Start the event generator
            setTimeout(generate_event, event_type.interval)
        end
        
        -- Stop simulation after 1 second
        setTimeout(function()
            simulation_active = false
        end, 1000)
        
        return #event_types
        """
        
        event_types_count = runtime.interpreter.lua.execute(script)
        
        # Wait for simulation to run
        await asyncio.sleep(1.5)
        
        events_processed = runtime.interpreter.lua.eval("events_processed")
        simulation_active = runtime.interpreter.lua.eval("simulation_active")
        
        print(f"Real-time Event Simulation - Event types: {event_types_count}, Events processed: {events_processed}, Active: {simulation_active}")
        
        assert event_types_count == 3
        # Events processed should be > 0 with fixed timer scoping
        if events_processed:
            assert events_processed > 0
            assert simulation_active == False  # Should have stopped
    
    @pytest.mark.asyncio 
    async def test_websocket_connection_recovery(self):
        """Test WebSocket connection recovery and error handling"""
        runtime = LuaAsyncRuntime()
        runtime.initialize_lua()
        
        script = """
        local connection_attempts = 0
        local connection_errors = 0
        local recovery_successful = false
        
        if net.WebSocket then
            local function attempt_connection()
                connection_attempts = connection_attempts + 1
                local ws = net.WebSocket()
                
                ws:connect("ws://localhost:99999/invalid", {  -- Invalid server
                    onopen = function()
                        recovery_successful = true
                    end,
                    onerror = function(error)
                        connection_errors = connection_errors + 1
                        
                        -- Attempt recovery after delay (max 3 attempts)
                        if connection_attempts < 3 then
                            setTimeout(attempt_connection, 200)
                        end
                    end
                })
            end
            
            attempt_connection()
            return true
        else
            return false
        end
        """
        
        websocket_supported = runtime.interpreter.lua.execute(script)
        
        if websocket_supported:
            # Wait for connection attempts
            await asyncio.sleep(1.0)
            
            connection_attempts = runtime.interpreter.lua.eval("connection_attempts")
            connection_errors = runtime.interpreter.lua.eval("connection_errors")
            recovery_successful = runtime.interpreter.lua.eval("recovery_successful")
            
            print(f"WebSocket Recovery Test - Attempts: {connection_attempts}, Errors: {connection_errors}, Recovery: {recovery_successful}")
            
            # Should have attempted connections and gotten errors for invalid server
            if connection_attempts:
                assert connection_attempts > 0
                assert connection_errors >= 0
                assert recovery_successful == False  # Should fail to connect to invalid server
        else:
            print("WebSocket Recovery Test - Not supported yet")
    
    @pytest.mark.asyncio
    async def test_server_sent_events_simulation(self):
        """Test Server-Sent Events (SSE) style functionality with HTTP streaming"""
        runtime = LuaAsyncRuntime()
        runtime.initialize_lua()
        
        # This test simulates SSE-like functionality using timers and HTTP
        script = """
        local sse_events = {}
        local event_count = 0
        local sse_active = true
        
        -- Simulate SSE event stream with periodic updates
        local function send_sse_event(event_type, data)
            if sse_active then
                event_count = event_count + 1
                sse_events[event_count] = {
                    id = event_count,
                    event = event_type,
                    data = data,
                    timestamp = os.clock()
                }
            end
        end
        
        -- Generate different types of events
        setTimeout(function()
            send_sse_event("data", {temperature = 22.5, humidity = 65})
        end, 100)
        
        setTimeout(function()
            send_sse_event("notification", {message = "System status OK"})
        end, 200)
        
        setTimeout(function()
            send_sse_event("update", {version = "1.2.3", status = "deployed"})
        end, 300)
        
        -- Stop after 500ms
        setTimeout(function()
            sse_active = false
        end, 500)
        
        return true
        """
        
        result = runtime.interpreter.lua.execute(script)
        
        # Wait for events to be generated
        await asyncio.sleep(0.8)
        
        event_count = runtime.interpreter.lua.eval("event_count")
        sse_active = runtime.interpreter.lua.eval("sse_active")
        
        print(f"SSE Simulation Test - Events generated: {event_count}, Active: {sse_active}")
        
        assert result == True
        # With fixed timer scoping, should have generated events
        if event_count:
            assert event_count >= 3
            assert sse_active == False  # Should have stopped
