import asyncio import os from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.staticfiles import StaticFiles import logging import threading import rclpy from .models import GeneratePlanRequest, ExecuteMissionRequest from .websocket_manager import websocket_manager from .py_tree_generator import py_tree_generator from .ros2_client import MissionActionClient # --- Application Setup --- app = FastAPI( title="Drone Backend Service", description="Handles mission planning, generation, and execution for the drone.", version="1.0.0", ) # --- Mount Static Files for Visualizations --- static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'generated_visualizations')) app.mount("/static", StaticFiles(directory=static_dir), name="static") # --- ROS2 Node and Client Initialization --- rclpy.init() ros2_client = MissionActionClient() def run_ros2_node(): """Spins the ROS2 node in a dedicated thread.""" logging.info("Starting to spin ROS2 node...") rclpy.spin(ros2_client) logging.info("ROS2 node has stopped spinning.") # --- API Endpoints --- @app.post("/generate_plan", response_model=dict) async def generate_plan_endpoint(request: GeneratePlanRequest): """ Receives a user prompt and returns a generated `py_tree.json` with a visualization URL. """ try: pytree_dict = await py_tree_generator.generate(request.user_prompt) return pytree_dict except RuntimeError as e: return {"error": str(e)} @app.post("/execute_mission", response_model=dict) async def execute_mission_endpoint(request: ExecuteMissionRequest): """ Receives a `py_tree.json` and sends it to the drone for execution. """ ros2_client.send_goal(request.py_tree) return {"status": "execution_started"} @app.websocket("/ws/status") async def websocket_endpoint(websocket: WebSocket): """ Handles the WebSocket connection for real-time status updates. """ await websocket_manager.connect(websocket) try: while True: await websocket.receive_text() except WebSocketDisconnect: websocket_manager.disconnect(websocket) logging.info("Client disconnected from WebSocket.") # --- Server Lifecycle --- @app.on_event("startup") async def startup_event(): """ On startup, get the current asyncio event loop and pass it to the websocket manager. Also, start the ROS2 node in a background thread. """ # Configure WebSocket Manager loop = asyncio.get_running_loop() websocket_manager.set_loop(loop) logging.info("WebSocket event loop configured.") # Start ROS2 node in a background thread ros2_thread = threading.Thread(target=run_ros2_node, daemon=True) ros2_thread.start() logging.info("ROS2 node thread started.") @app.on_event("shutdown") async def shutdown_event(): logging.info("Backend service shutting down.") ros2_client.destroy_node() rclpy.shutdown() logging.info("ROS2 node shut down successfully.")