"""
WebSocket Client
Handles connection to Spectra backend for receiving jobs
"""
import asyncio
import json
import logging
from typing import Callable, Optional, Any
from datetime import datetime

import websockets
from rich.console import Console

console = Console()
logger = logging.getLogger(__name__)


class WebSocketClient:
    """
    WebSocket client for connecting to Spectra backend.
    
    Handles:
    - Connection with authentication
    - Receiving jobs
    - Sending results
    - Reconnection on disconnect
    """
    
    def __init__(self, backend_url: str, token: str):
        """
        Initialize WebSocket client.
        
        Args:
            backend_url: Backend URL (https://...)
            token: Authentication token
        """
        # Convert HTTP URL to WebSocket URL
        ws_url = backend_url.replace("https://", "wss://").replace("http://", "ws://")
        self.ws_url = f"{ws_url}/ws/runner?token={token}"
        self.token = token
        self.websocket = None
        self.runner_id = None
        self.connected = False
        self.job_handler: Optional[Callable] = None
    
    async def connect(self) -> bool:
        """
        Connect to the WebSocket server.
        
        Returns:
            True if connected successfully
        """
        try:
            self.websocket = await websockets.connect(
                self.ws_url,
                ping_interval=30,
                ping_timeout=10
            )
            
            # Wait for connection confirmation
            response = await self.websocket.recv()
            data = json.loads(response)
            
            if data.get("type") == "connected":
                self.connected = True
                self.runner_id = data.get("runner_id")
                return True
            elif data.get("type") == "error":
                console.print(f"❌ {data.get('message')}", style="red")
                return False
            
            return False
            
        except Exception as e:
            logger.error(f"Connection failed: {e}")
            console.print(f"❌ Connection failed: {e}", style="red")
            return False
    
    async def disconnect(self):
        """Disconnect from server"""
        if self.websocket:
            await self.websocket.close()
            self.connected = False
    
    def set_job_handler(self, handler: Callable):
        """
        Set the handler function for incoming jobs.
        
        Args:
            handler: Async function that takes (job, signature) and returns result
        """
        self.job_handler = handler
    
    async def listen(self):
        """
        Listen for incoming messages.
        
        This is the main loop that receives jobs and sends results.
        """
        if not self.websocket:
            raise RuntimeError("Not connected. Call connect() first.")
        
        try:
            async for message in self.websocket:
                data = json.loads(message)
                msg_type = data.get("type")
                
                if msg_type == "job":
                    # Handle incoming job
                    job = data.get("job", {})
                    signature = data.get("signature", "")
                    
                    console.print()
                    console.print(f"[{datetime.now().strftime('%H:%M:%S')}] 📦 [bold]Job received[/bold]")
                    console.print(f"   File: {job.get('test_file', 'unknown')}")
                    
                    if self.job_handler:
                        # Process job
                        result = await self.job_handler(job, signature)
                        
                        # Send result
                        await self.send_result(job.get("id"), result)
                    else:
                        console.print("   ⚠️ No job handler configured", style="yellow")
                
                elif msg_type == "ack":
                    # Job result acknowledged
                    job_id = data.get("job_id")
                    console.print(f"   ✅ Result acknowledged", style="green")
                
                elif msg_type == "pong":
                    # Heartbeat response
                    pass
                
                elif msg_type == "error":
                    console.print(f"   ❌ Error: {data.get('message')}", style="red")
        
        except websockets.ConnectionClosed:
            console.print("\n🔌 Connection closed", style="yellow")
            self.connected = False
        except Exception as e:
            logger.error(f"Listen error: {e}")
            console.print(f"\n❌ Error: {e}", style="red")
            self.connected = False
    
    async def send_result(self, job_id: str, result: dict):
        """
        Send job result to server.
        
        Args:
            job_id: ID of the completed job
            result: Test result data
        """
        if not self.websocket:
            return
        
        message = {
            "type": "result",
            "job_id": job_id,
            "result": result
        }
        
        await self.websocket.send(json.dumps(message))
        console.print(f"   📤 Result sent", style="dim")
    
    async def send_status(self, status: str):
        """Send status update to server"""
        if not self.websocket:
            return
        
        message = {
            "type": "status",
            "status": status
        }
        
        await self.websocket.send(json.dumps(message))
