from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from optimization_manager import OptimizationManager

app = FastAPI()
PORT = 7860


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    print("starting")
    await websocket.accept()
    manager = OptimizationManager()
    
    try:
        while True:
            # Receive message from client
            message = await websocket.receive_json()
            message_type = message.get("type")
            
            # Handle different message types
            if message_type == "init":
                await handle_init(websocket, manager, message.get("settings", {}))
            
            elif message_type == "reset":
                await handle_reset(websocket, manager)
            
            elif message_type == "next-step":
                await handle_next_step(websocket, manager)
            
            elif message_type == "prev-step":
                await handle_prev_step(websocket, manager)
            
            elif message_type == "play":
                await handle_play(websocket, manager)
            
            elif message_type == "pause":
                await handle_pause(websocket, manager)
            
            else:
                await send_error(websocket, f"Unknown message type: {message_type}")
    
    except WebSocketDisconnect:
        print("Client disconnected")
    except Exception as e:
        print(f"Error: {e}")
        await send_error(websocket, str(e))


async def handle_init(websocket: WebSocket, manager: OptimizationManager, settings: dict):
    """Handle settings update."""
    data = manager.handle_update_settings(settings)
    if data:
        await websocket.send_json(data)


async def handle_reset(websocket: WebSocket, manager: OptimizationManager):
    """Handle trajectory reset."""
    data = manager.handle_reset()
    if data:
        await websocket.send_json(data)


async def handle_next_step(websocket: WebSocket, manager: OptimizationManager):
    """Handle next step request."""
    data = manager.handle_next_step()
    if data:
        await websocket.send_json(data)


async def handle_prev_step(websocket: WebSocket, manager: OptimizationManager):
    """Handle previous step request."""
    data = manager.handle_prev_step()
    if data:
        await websocket.send_json(data)


async def handle_play(websocket: WebSocket, manager: OptimizationManager):
    """Handle play request (to be implemented)."""
    # TODO: Implement play functionality with automatic step advancement
    pass


async def handle_pause(websocket: WebSocket, manager: OptimizationManager):
    """Handle pause request (to be implemented)."""
    # TODO: Implement pause functionality
    pass


async def send_error(websocket: WebSocket, error_message: str):
    """Send error message to client."""
    await websocket.send_json({
        "type": "error",
        "message": error_message
    })


app.mount("/", StaticFiles(directory="../frontend/dist", html=True), name="frontend")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)