app/api/routes/websocket.py

"""
WebSocket Routes for Real-time Execution Streaming.

Provides live updates during workflow execution.
"""

from typing import Any, Dict, Set
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from uuid import uuid4
import asyncio
import json
import logging

from app.engine.graph import Graph
from app.engine.executor import Executor, ExecutionStep
from app.storage.memory import graph_storage, run_storage


logger = logging.getLogger(__name__)

router = APIRouter(tags=["WebSocket"])


class ConnectionManager:
    """Manages WebSocket connections."""
    
    def __init__(self):
        self.active_connections: Dict[str, Set[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, run_id: str):
        """Accept a new WebSocket connection."""
        await websocket.accept()
        if run_id not in self.active_connections:
            self.active_connections[run_id] = set()
        self.active_connections[run_id].add(websocket)
        logger.info(f"WebSocket connected for run: {run_id}")
    
    def disconnect(self, websocket: WebSocket, run_id: str):
        """Remove a WebSocket connection."""
        if run_id in self.active_connections:
            self.active_connections[run_id].discard(websocket)
            if not self.active_connections[run_id]:
                del self.active_connections[run_id]
        logger.info(f"WebSocket disconnected for run: {run_id}")
    
    async def broadcast(self, run_id: str, message: Dict[str, Any]):
        """Broadcast a message to all connections for a run."""
        if run_id in self.active_connections:
            disconnected = set()
            for websocket in self.active_connections[run_id]:
                try:
                    await websocket.send_json(message)
                except Exception:
                    disconnected.add(websocket)
            
            # Clean up disconnected clients
            for ws in disconnected:
                self.active_connections[run_id].discard(ws)


# Global connection manager
manager = ConnectionManager()


@router.websocket("/ws/run/{graph_id}")
async def websocket_run(websocket: WebSocket, graph_id: str):
    """
    WebSocket endpoint for real-time workflow execution.
    
    Connect to this endpoint and send the initial state as JSON.
    You'll receive step-by-step updates as the workflow executes.
    
    Message format (client -> server):
    ```json
    {"action": "start", "initial_state": {"code": "..."}}
    ```
    
    Message format (server -> client):
    ```json
    {
        "type": "step",
        "step": 1,
        "node": "extract",
        "status": "completed",
        "duration_ms": 15.5,
        "state": {...}
    }
    ```
    """
    # Check if graph exists
    stored = await graph_storage.get(graph_id)
    if not stored:
        await websocket.close(code=4004, reason=f"Graph '{graph_id}' not found")
        return
    
    run_id = str(uuid4())
    await manager.connect(websocket, run_id)
    
    try:
        # Wait for start message
        data = await websocket.receive_json()
        
        if data.get("action") != "start":
            await websocket.send_json({
                "type": "error",
                "error": "Expected 'start' action"
            })
            return
        
        initial_state = data.get("initial_state", {})
        
        # Send acknowledgment
        await websocket.send_json({
            "type": "started",
            "run_id": run_id,
            "graph_id": graph_id,
        })
        
        # Rebuild graph
        graph = await _rebuild_graph(stored.definition)
        
        # Create run record
        await run_storage.create(run_id, graph_id, initial_state)
        
        # Execute with streaming updates
        async def on_step(step: ExecutionStep, state: Dict[str, Any]):
            await manager.broadcast(run_id, {
                "type": "step",
                "step": step.step,
                "node": step.node,
                "status": step.result,
                "duration_ms": step.duration_ms,
                "iteration": step.iteration,
                "route_taken": step.route_taken,
                "error": step.error,
                "state": state,
            })
        
        executor = Executor(graph, run_id=run_id)
        
        # Run with step notifications
        result = await _run_with_streaming(executor, initial_state, on_step)
        
        # Send completion
        await websocket.send_json({
            "type": "completed",
            "run_id": run_id,
            "status": result.status.value,
            "final_state": result.final_state,
            "total_duration_ms": result.total_duration_ms,
            "iterations": result.iterations,
            "error": result.error,
        })
        
        # Update storage
        if result.status.value == "completed":
            await run_storage.complete(
                run_id,
                result.final_state,
                [s.to_dict() for s in result.execution_log],
            )
        else:
            await run_storage.fail(run_id, result.error or "Unknown error")
        
    except WebSocketDisconnect:
        logger.info(f"Client disconnected from run {run_id}")
    except Exception as e:
        logger.exception(f"WebSocket error: {e}")
        try:
            await websocket.send_json({
                "type": "error",
                "error": str(e),
            })
        except Exception:
            pass
    finally:
        manager.disconnect(websocket, run_id)


async def _rebuild_graph(definition: Dict[str, Any]) -> Graph:
    """Rebuild graph from definition (copied from graph.py to avoid circular import)."""
    from app.api.routes.graph import _rebuild_graph_from_definition
    return await _rebuild_graph_from_definition(definition)


async def _run_with_streaming(
    executor: Executor,
    initial_state: Dict[str, Any],
    on_step
):
    """Run executor with async step callbacks."""
    from app.engine.graph import END
    from app.engine.state import StateManager
    import time
    from datetime import datetime
    
    # Execute the workflow
    result = await executor.run(initial_state)
    
    # Stream each step (already executed, but we notify)
    for step in result.execution_log:
        await on_step(step, result.final_state)
        await asyncio.sleep(0.01)  # Small delay for streaming effect
    
    return result


@router.websocket("/ws/subscribe/{run_id}")
async def websocket_subscribe(websocket: WebSocket, run_id: str):
    """
    Subscribe to updates for an existing run.
    
    Use this to watch an async execution started via POST /graph/run.
    """
    # Check if run exists
    stored = await run_storage.get(run_id)
    if not stored:
        await websocket.close(code=4004, reason=f"Run '{run_id}' not found")
        return
    
    await manager.connect(websocket, run_id)
    
    try:
        # Send current state
        await websocket.send_json({
            "type": "current_state",
            "run_id": run_id,
            "status": stored.status,
            "current_node": stored.current_node,
            "iteration": stored.iteration,
            "state": stored.current_state,
        })
        
        # Keep connection open and poll for updates
        last_log_count = len(stored.execution_log)
        
        while True:
            await asyncio.sleep(0.5)  # Poll interval
            
            stored = await run_storage.get(run_id)
            if not stored:
                break
            
            # Send new log entries
            if len(stored.execution_log) > last_log_count:
                for entry in stored.execution_log[last_log_count:]:
                    await websocket.send_json({
                        "type": "step",
                        **entry,
                    })
                last_log_count = len(stored.execution_log)
            
            # Check if completed
            if stored.status in ("completed", "failed", "cancelled"):
                await websocket.send_json({
                    "type": "completed",
                    "run_id": run_id,
                    "status": stored.status,
                    "final_state": stored.final_state,
                    "error": stored.error,
                })
                break
                
    except WebSocketDisconnect:
        logger.info(f"Subscriber disconnected from run {run_id}")
    except Exception as e:
        logger.exception(f"WebSocket error: {e}")
    finally:
        manager.disconnect(websocket, run_id)