from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.api.dependencies import ServiceContainer
import json


def create_workspace_router(services: ServiceContainer) -> APIRouter:
    router = APIRouter()

    @router.websocket("/ws/workspace/{session_id}")
    async def websocket_endpoint(websocket: WebSocket, session_id: str):
        await websocket.accept()
        print(f"WebSocket connection accepted for session: {session_id}")
        
        # Send a welcome message to confirm the connection is active
        await websocket.send_text(json.dumps({
            "type": "connection_established",
            "message": f"Connected to AI Hub. Session ID: {session_id}"
        }))

        try:
            await websocket.send_text(json.dumps({
                "type": "connection_established",
                "message": f"Connected to AI Hub. Session ID: {session_id}"
            }))

            while True:
                message = await websocket.receive_text()
                data = json.loads(message)
                
                # The endpoint's only job is to dispatch the message to the service
                await services.workspace_service.dispatch_message(websocket, data)
                
        except WebSocketDisconnect:
            print(f"WebSocket connection disconnected for session: {session_id}")
        except Exception as e:
            print(f"An error occurred: {e}")
            
        finally:
            print(f"Closing WebSocket for session: {session_id}")
    return router