server.py

"""REST API Server for Heliosail-RX Simulation.

Provides HTTP endpoints for:
  • Mission creation, execution, monitoring
  • Simulation result retrieval
  • Multi-objective optimization runs
  • Configuration management

Architecture:
  • Mission Manager: Tracks concurrent missions with UUIDs
  • REST Layer: HTTP endpoints using FastAPI
  • Kernel Integration: Translates API requests to SimulationEngine
  • JSON Config: Schema for SimulationConfig and optimization parameters
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from pathlib import Path
import os
import time
import asyncio
from concurrent.futures import ProcessPoolExecutor
import sys
import uuid
import random
from typing import Dict, Any, List, Optional
from datetime import datetime
from pathlib import Path
from enum import Enum
import math
import sqlite3
import pickle
import json
import logging

# Add propulsion-hybrid to path for IonThrusterModule
_prop_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "propulsion-hybrid"))
if _prop_path not in sys.path:
    sys.path.insert(0, _prop_path)
    
try:
    from propulsion_hybrid.ion_thrust import IonThrusterModule
except ImportError:
    IonThrusterModule = None

try:
    from models.wrappers.lorentz_wrapper import LorentzForceModule
except ImportError:
    LorentzForceModule = None

# Add project root to sys.path to allow imports from models
sys.path.insert(0, str(Path(__file__).parent.parent))

from models.core import (
    SimulationConfig, SpacecraftConfig, EnvironmentConfig, 
    PhysicsConfig, SolverConfig, SpacecraftState, SolverType
)
from api.mission import Mission, MissionBuilder
from models.wrappers.orbital_mechanics_wrapper import CombinedOrbitalDynamicsModule
from models.wrappers.sail_physics_wrapper import CombinedSailDynamicsModule


logger = logging.getLogger(__name__)


# ============================================================================
# REST API Data Models (Pydantic)
# ============================================================================

class OrbitType(str, Enum):
    """Predefined orbital configurations."""
    LEO = "LEO"  # Low Earth Orbit (~400 km)
    GEO = "GEO"  # Geostationary (~36000 km)
    SOLAR = "SOLAR"  # Solar sail mission (~7M km)
    MARS = "MARS"  # Mars transfer orbit
    ASTEROID = "ASTEROID"  # Generic Asteroid rendezvous
    MERCURY = "MERCURY" # Mercury transfer
    VENUS = "VENUS" # Venus transfer
    JUPITER = "JUPITER" # Jupiter transfer
    SATURN = "SATURN" # Saturn transfer
    URANUS = "URANUS" # Uranus transfer
    NEPTUNE = "NEPTUNE" # Neptune transfer
    PLUTO = "PLUTO" # Pluto transfer
    ASTEROID_BELT = "ASTEROID_BELT" # Main Asteroid Belt (Ceres)
    JUPITER_MOONS = "JUPITER_MOONS" # Jovian System (Europa focus)
    EARTH = "EARTH" # Earth (1 AU Heliocentric)


class SailType(str, Enum):
    """Solar sail architecture options."""
    FLAT_PLATE = "flat_plate"
    GOSSAMER = "gossamer"
    HYBRID = "hybrid"  # Solar + ion thrusters


class OptimizerType(str, Enum):
    """Optimization algorithms."""
    SINGLE_OBJECTIVE = "single_objective"
    MULTI_OBJECTIVE = "multi_objective"  # Pareto front
    GENETIC_ALGORITHM = "genetic_algorithm"
    DIFFERENTIAL_EVOLUTION = "differential_evolution"


class SimulationConfigJSON(BaseModel):
    """Mission configuration (JSON schema for REST API).
    
    This maps user-friendly mission parameters to SimulationConfig.
    """
    # Mission Identity
    mission_name: str = Field(default="untitled_mission")
    mission_target: str = Field(default="Mars", description="Target body or destination")
    
    # Orbit
    initial_orbit: OrbitType = Field(default=OrbitType.SOLAR)
    departure_date: str = Field(default="2026-02-23")  # ISO 8601
    
    # Spacecraft Properties
    spacecraft_mass: float = Field(default=260.0, ge=10.0, le=10000.0)
    spacecraft_name: str = Field(default="Heliosail-RX")
    
    # Sail Configuration
    sail_area: float = Field(default=1200.0, ge=10.0, le=10000.0, description="[m²]")
    reflectivity: float = Field(default=0.92, ge=0.5, le=0.99)
    sail_type: SailType = Field(default=SailType.FLAT_PLATE)
    hybrid_mode: bool = Field(default=False, description="Enable ion thrusters")
    
    # Propulsion (if hybrid_mode=True)
    ion_thrust_newtons: float = Field(default=0.0, ge=0.0, le=100.0)
    ion_isp_seconds: float = Field(default=3000.0, ge=300.0, le=10000.0)
    fuel_mass_kg: float = Field(default=0.0, ge=0.0, le=1000.0)
    
    # Simulation Parameters
    simulation_duration_days: float = Field(default=365.0, ge=0.1, le=10000.0)
    time_step_seconds: float = Field(default=10.0, ge=0.1, le=1000.0)
    
    # Advanced Environment Models
    use_cme: bool = Field(default=False, description="Enable Coronal Mass Ejections")
    use_solar_cycle: bool = Field(default=True, description="Enable 11-Year Solar Cycle Irradiance Variance")
    use_atmospheric_drag: bool = Field(default=False, description="Enable Atmospheric Drag")
    
    # Advanced Environment Physics
    enable_gravity_assists: bool = Field(default=False, description="Enable Third-Body perturbations for gravity slingshots")
    spacecraft_charge: float = Field(default=0.0, description="Electrostatic charge (Coulombs) for Lorentz forces")
    
    # Advanced Solver Config
    solver_type: str = Field(default="RK4_SYMPLECTIC", description="ODE Integration Method")
    
    # Optimization (if optimizer is specified)
    optimizer: Optional[OptimizerType] = Field(default=None)
    optimizer_objectives: Optional[List[str]] = Field(
        default=None,
        description="e.g., ['max_final_distance', 'min_time', 'min_fuel']"
    )
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "mission_name": "Heliosail Mars Transfer",
                "mission_target": "Mars",
                "initial_orbit": "SOLAR",
                "spacecraft_mass": 260.0,
                "sail_area": 1200.0,
                "reflectivity": 0.92,
                "sail_type": "flat_plate",
                "hybrid_mode": True,
                "ion_thrust_newtons": 0.08,
                "ion_isp_seconds": 3000.0,
                "fuel_mass_kg": 50.0,
                "simulation_duration_days": 365.0,
                "time_step_seconds": 10.0,
                "use_cme": False,
                "use_atmospheric_drag": False,
                "solver_type": "RK4_SYMPLECTIC",
                "optimizer": "multi_objective",
                "optimizer_objectives": ["max_final_distance", "min_time"],
            }
        }
    }


class MissionResponse(BaseModel):
    """Response from mission creation/status endpoints."""
    mission_id: str = Field(description="UUID for this mission (for queries)")
    status: str = Field(description="queued, running, completed, failed")
    config: Dict[str, Any] = Field(description="Configuration used")
    created_at: str = Field(description="ISO 8601 timestamp")
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    error: Optional[str] = None


class TrajectoryPoint(BaseModel):
    """Single point in spacecraft trajectory."""
    t_sec: float = Field(description="Mission time [seconds]")
    position: tuple = Field(description="[x, y, z] in meters")
    velocity: tuple = Field(description="[vx, vy, vz] in m/s")
    diagnostics: Dict[str, Any] = Field(default_factory=dict)


class SimulationResultsJSON(BaseModel):
    """Simulation results (JSON schema for REST API)."""
    mission_id: str
    status: str  # "success" or "failed"
    config: Dict[str, Any]
    summary: Dict[str, Any] = Field(default_factory=dict)
    trajectory: List[TrajectoryPoint] = Field(default_factory=list)
    telemetry: Dict[str, List[tuple]] = Field(default_factory=dict)
    error_log: Optional[str] = None


class OptimizationRequest(BaseModel):
    """Request for multi-mission optimization sweep."""
    sweep_name: str = Field(default="optimization_run")
    base_config: SimulationConfigJSON
    parameter_ranges: Dict[str, tuple] = Field(
        description="Parameter name -> (min, max) tuples"
    )
    objectives: List[str] = Field(description="Objectives to optimize")
    num_samples: int = Field(default=100, ge=10, le=10000, description="Population size")


class OptimizationResult(BaseModel):
    """Result of optimization sweep."""
    sweep_id: str
    status: str  # "queued", "running" or "completed", "failed"
    message: Optional[str] = None
    base_config: Dict[str, Any]
    parameter_ranges: Optional[Dict[str, tuple]] = None
    objectives: Optional[List[str]] = None
    pareto_front: List[Dict[str, Any]] = Field(default_factory=list)
    num_evaluated: int = 0
    total_samples: int = 0
    best_objectives: Dict[str, float] = Field(default_factory=dict)
    error: Optional[str] = None


# ============================================================================
# Mission Manager (State Storage)
# ============================================================================

DB_FILE = "heliosail.sqlite"

class MissionManager:
    """Manages concurrent missions and optimizations with UUIDs and state tracking."""
    
    def __init__(self):
        self.optimizations: Dict[str, Dict[str, Any]] = {} # sweep_id -> optimization_data
        self._init_db()
        
    def _init_db(self):
        with sqlite3.connect(DB_FILE) as conn:
            conn.execute('''
                CREATE TABLE IF NOT EXISTS missions (
                    id TEXT PRIMARY KEY,
                    config_json TEXT,
                    status TEXT,
                    created_at TEXT,
                    started_at TEXT,
                    completed_at TEXT,
                    error TEXT,
                    result_blob BLOB
                )
            ''')
    
    def create_mission(self, config: SimulationConfigJSON) -> str:
        """Create a new mission, return UUID."""
        mission_id = str(uuid.uuid4())
        config_json = config.model_dump_json() if hasattr(config, "model_dump_json") else json.dumps(config)
        
        with sqlite3.connect(DB_FILE) as conn:
            conn.execute('''
                INSERT INTO missions (id, config_json, status, created_at)
                VALUES (?, ?, ?, ?)
            ''', (mission_id, config_json, "queued", datetime.utcnow().isoformat()))
            
        logger.info(f"Created mission {mission_id}: {config.mission_name}")
        return mission_id
    
    def get_status(self, mission_id: str) -> Dict[str, Any]:
        """Get mission status by UUID."""
        with sqlite3.connect(DB_FILE) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute("SELECT * FROM missions WHERE id=?", (mission_id,)).fetchone()
            if not row:
                raise ValueError(f"Mission {mission_id} not found")
                
            config_dict = json.loads(row["config_json"])
            config = SimulationConfigJSON(**config_dict)
            
            return {
                "status": row["status"],
                "created_at": row["created_at"],
                "started_at": row["started_at"],
                "completed_at": row["completed_at"],
                "error": row["error"],
                "config": config,
                "mission_obj": None
            }
    
    def get_results(self, mission_id: str) -> Optional[Any]:
        """Get simulation results by mission UUID."""
        with sqlite3.connect(DB_FILE) as conn:
            row = conn.execute("SELECT result_blob FROM missions WHERE id=?", (mission_id,)).fetchone()
            if not row or not row[0]:
                raise ValueError(f"Results for {mission_id} not found (may still be running)")
            return pickle.loads(row[0])
    
    def list_missions(self) -> List[str]:
        """Get all mission UUIDs."""
        with sqlite3.connect(DB_FILE) as conn:
            rows = conn.execute("SELECT id FROM missions ORDER BY created_at DESC").fetchall()
            return [r[0] for r in rows]

    def update_mission_status(self, mission_id: str, status: str, error: str = None, result_obj: Any = None):
        """Safely updates a mission's state and persists to disk."""
        with sqlite3.connect(DB_FILE) as conn:
            now = datetime.utcnow().isoformat()
            if status == "running":
                conn.execute("UPDATE missions SET status=?, started_at=? WHERE id=?", (status, now, mission_id))
            elif status == "completed":
                res_blob = pickle.dumps(result_obj)
                conn.execute("UPDATE missions SET status=?, completed_at=?, result_blob=? WHERE id=?", (status, now, res_blob, mission_id))
            elif status == "failed":
                conn.execute("UPDATE missions SET status=?, completed_at=?, error=? WHERE id=?", (status, now, error, mission_id))
    
    def _build_mission_obj(self, config: SimulationConfigJSON) -> Mission:
        """Convert REST config JSON to Mission object."""
        
        # Determine initial orbit
        orbit_params = self._get_orbit_params(config.initial_orbit)
        
        # Create initial spacecraft state
        initial_state = SpacecraftState(
            position=orbit_params["position"],
            velocity=orbit_params["velocity"],
            attitude_quaternion=(1.0, 0.0, 0.0, 0.0),
            angular_velocity=(0.0, 0.0, 0.0),
            mass=config.spacecraft_mass,
            power_available=500.0,  # W (nominal)
            thermal_state={},
            structural_modes=(),
            epoch_sec=0.0,
            mission_elapsed_sec=0.0,
        )
        
        # Build component configurations
        spacecraft_config = SpacecraftConfig(
            mass_dry_kg=max(1.0, config.spacecraft_mass - config.fuel_mass_kg),
            mass_fuel_kg=config.fuel_mass_kg if config.hybrid_mode else 0.0,
            sail_area_m2=config.sail_area,
            sail_reflectivity=config.reflectivity
        )
        
        env_config = EnvironmentConfig(
            use_cme=config.use_cme,
            use_atmospheric_drag=config.use_atmospheric_drag,
            use_solar_cycle=config.use_solar_cycle,
            use_radiation_belts=True,
            use_third_body=True
        )
        
        solver_enum = SolverType.RK4_SYMPLECTIC
        try:
            solver_enum = SolverType[config.solver_type.upper()]
        except KeyError:
            logger.warning(f"Invalid solver {config.solver_type}. Using RK4_SYMPLECTIC.")
            
        solver_config = SolverConfig(
            solver=solver_enum,
            dt_nominal=config.time_step_seconds,
            dt_min=max(0.1, config.time_step_seconds / 10),
            dt_max=config.time_step_seconds * 10
        )
        
        # Convert mission duration to seconds
        t_end = config.simulation_duration_days * 86400.0
        
        # Set telemetry step to target approx 10000 points max to prevent huge IPC transfers
        telemetry_step = max(0.1, config.time_step_seconds, t_end / 10000.0)
        
        # Create SimulationConfig
        sim_config = SimulationConfig(
            name=config.mission_name,
            t_start=0.0,
            t_end=t_end,
            telemetry_step_s=telemetry_step,
            initial_state=initial_state,
            spacecraft=spacecraft_config,
            environment=env_config,
            solver=solver_config
        )
        
        # Create Mission object
        mission = Mission(sim_config)
        
        # Determine 3rd body (Jupiter approximation) for gravity assist
        # Average distance ~5.2 AU, mu ~1.266e17
        jupiter_pos = (5.2 * 1.496e11, 0, 0)
        jupiter_mu = 1.266e17
        
        # Add physics modules
        gravity_module = CombinedOrbitalDynamicsModule(
            mu=orbit_params.get("mu", 1.327e20),
            include_perturbations=True,
            perturbation_config={
                "enable_third_body": config.enable_gravity_assists
            }
        )
        
        if config.enable_gravity_assists:
            gravity_module.initialize({
                "perturbation_config": {
                    "enable_third_body": True,
                    "third_body_position": jupiter_pos,
                    "third_body_name": "Jupiter",
                    "third_body_mu": jupiter_mu
                }
            })
            
        mission.add_physics_module("orbital_mechanics", gravity_module)
        
        sail_module = CombinedSailDynamicsModule(
            sail_area_m2=config.sail_area,
            reflectivity=config.reflectivity,
            mass_kg=config.spacecraft_mass,
            include_thermal=True
        )
        mission.add_physics_module("sail_physics", sail_module)
        
        # Inject Hybrid Ion Propulsion if enabled
        if config.hybrid_mode and IonThrusterModule is not None:
            ion_module = IonThrusterModule()
            ion_module.initialize({
                "isp": config.ion_isp_seconds,
                "max_thrust": config.ion_thrust_newtons,
                "efficiency": 0.6,
                "max_power": 500.0
            })
            mission.add_physics_module("hybrid_propulsion", ion_module)

        # Inject Lorentz Force interactions if charged
        if config.spacecraft_charge != 0.0 and LorentzForceModule is not None:
            lorentz_module = LorentzForceModule(charge_coulombs=config.spacecraft_charge)
            mission.add_physics_module("lorentz", lorentz_module)
        
        logger.info(f"Mission {config.mission_name} configured for {t_end} seconds")
        
        return mission
    
    def _get_orbit_params(self, orbit_type: OrbitType) -> Dict[str, Any]:
        """Get position/velocity/mu for standard orbits."""
        mu_sun = 1.327e20
        mu_earth = 3.986e14
        mu_mars = 4.282e13
        mu_jupiter = 1.266e17
        mu_saturn = 3.793e16
        
        orbits = {
            OrbitType.LEO: {
                "position": (6.7e6, 0.0, 0.0),  # 400 km altitude from Earth Center
                "velocity": (0.0, 7660.0, 0.0),  # m/s (circular around Earth)
                "mu": mu_earth
            },
            OrbitType.EARTH: {
                "position": (1.496e11, 0.0, 0.0), # 1 AU
                "velocity": (0.0, 29.78e3, 0.0), # Earth average orbital velocity
                "mu": mu_sun
            },
            OrbitType.GEO: {
                "position": (4.22e7, 0.0, 0.0),  # 36000 km altitude from Earth Center
                "velocity": (0.0, 3070.0, 0.0),
                "mu": mu_earth
            },
            OrbitType.SOLAR: {
                "position": (7e9, 0.0, 0.0),  # 7M km (solar sail regime)
                "velocity": (0.0, 30e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.MARS: {
                "position": (2.28e11, 0.0, 0.0),  # ~1.52 AU
                "velocity": (0.0, 24e3, 0.0),  # Average orbital velocity
                "mu": mu_sun
            },
            OrbitType.ASTEROID: {
                "position": (4e11, 0.0, 0.0),  # Beyond Mars
                "velocity": (0.0, 20e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.MERCURY: {
                "position": (5.79e10, 0.0, 0.0), # ~0.39 AU
                "velocity": (0.0, 47.36e3, 0.0), 
                "mu": mu_sun
            },
            OrbitType.VENUS: {
                "position": (1.08e11, 0.0, 0.0), # ~0.72 AU
                "velocity": (0.0, 35.02e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.JUPITER: {
                "position": (7.78e11, 0.0, 0.0), # ~5.20 AU
                "velocity": (0.0, 13.07e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.SATURN: {
                "position": (1.43e12, 0.0, 0.0), # ~9.58 AU
                "velocity": (0.0, 9.68e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.URANUS: {
                "position": (2.87e12, 0.0, 0.0), # ~19.2 AU
                "velocity": (0.0, 6.80e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.NEPTUNE: {
                "position": (4.50e12, 0.0, 0.0), # ~30.0 AU
                "velocity": (0.0, 5.43e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.PLUTO: {
                "position": (5.90e12, 0.0, 0.0), # ~39.4 AU
                "velocity": (0.0, 4.74e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.ASTEROID_BELT: {
                "position": (4.13e11, 0.0, 0.0), # ~2.76 AU (Ceres)
                "velocity": (0.0, 17.9e3, 0.0),
                "mu": mu_sun
            },
            OrbitType.JUPITER_MOONS: {
                "position": (7.78e11, 0.0, 0.0), # Same distance as Jupiter system
                "velocity": (0.0, 15e3, 0.0), # Altered entry velocity for Jovian capture
                "mu": mu_sun
            },
        }
        return orbits.get(orbit_type, orbits[OrbitType.SOLAR])


# ============================================================================
# REST API (FastAPI)
# ============================================================================

app = FastAPI(
    title="Heliosail-RX Simulation API",
    version="1.0.0",
    description="REST API for solar sail trajectory simulation and optimization",
)

# CORS for GUI/external clients
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    """Initialize resources like ProcessPoolExecutor on startup."""
    global process_pool
    # Only spawn executor in the main process
    process_pool = ProcessPoolExecutor()

@app.on_event("shutdown")
async def shutdown_event():
    """Clean up resources on shutdown."""
    if process_pool:
        process_pool.shutdown(wait=True)

# Shared mission manager
mission_manager = MissionManager()

# Mount the Web GUI
gui_path = Path(__file__).parent.parent / "web-gui"
if gui_path.exists():
    app.mount("/static", StaticFiles(directory=str(gui_path)), name="static")

@app.get("/", tags=["gui"])
async def serve_gui():
    """Serve the Web GUI dashboard."""
    index_file = gui_path / "index.html"
    if index_file.exists():
        return FileResponse(str(index_file))
    return {"message": "Heliosail-RX API is running, but GUI is not found."}

@app.get("/styles.css", tags=["gui"])
async def serve_css():
    css_file = gui_path / "styles.css"
    if css_file.exists():
        return FileResponse(str(css_file))
    raise HTTPException(status_code=404)

@app.get("/app.js", tags=["gui"])
async def serve_js():
    js_file = gui_path / "app.js"
    if js_file.exists():
        return FileResponse(str(js_file))
    raise HTTPException(status_code=404)


# ════════════════════════════════════════════════════════════════════════════
# MISSION ENDPOINTS
# ════════════════════════════════════════════════════════════════════════════

@app.post("/mission/create", response_model=MissionResponse)
async def create_mission(config: SimulationConfigJSON):
    """Create a new simulation mission.
    
    POST /mission/create
    {
      "mission_name": "Heliosail Mars Transfer",
      "mission_target": "Mars",
      "initial_orbit": "SOLAR",
      "sail_area": 1200.0,
      ...
    }
    
    Returns:
      mission_id: UUID for tracking this mission
      status: "queued" (ready to run)
    """
    try:
        mission_id = mission_manager.create_mission(config)
        status = mission_manager.get_status(mission_id)
        
        return MissionResponse(
            mission_id=mission_id,
            status=status["status"],
            config=config.model_dump(),
            created_at=status["created_at"],
        )
    except Exception as e:
        logger.error(f"Error creating mission: {e}")
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/mission/run/{mission_id}")
async def run_mission(mission_id: str, background_tasks: BackgroundTasks):
    """Execute a queued mission (asynchronous).
    
    POST /mission/run/{mission_id}
    
    Returns:
      status: "running"
    """
    try:
        status = mission_manager.get_status(mission_id)
        if status["status"] != "queued":
            raise ValueError(f"Mission {mission_id} is not in queued state: {status['status']}")
        
        # Mark as running
        mission_manager.update_mission_status(mission_id, "running")
        
        # Extract config to a dictionary for process pool execution
        config_obj = status["config"]
        config_dict = config_obj.model_dump() if hasattr(config_obj, "model_dump") else config_obj
        
        # Add simulation task to background using async wrapper
        background_tasks.add_task(_run_simulation_async, mission_id, config_dict)
        
        return {
            "mission_id": mission_id,
            "status": "running",
            "message": f"Mission {mission_id} started in background",
        }
    except Exception as e:
        logger.error(f"Error running mission {mission_id}: {e}")
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/mission/{mission_id}/status", response_model=MissionResponse)
async def get_mission_status(mission_id: str):
    """Get current status of a mission.
    
    GET /mission/{mission_id}/status
    
    Returns:
      status: "queued" | "running" | "completed" | "failed"
    """
    try:
        status = mission_manager.get_status(mission_id)
        
        return MissionResponse(
            mission_id=mission_id,
            status=status["status"],
            config=status["config"].model_dump(),
            created_at=status["created_at"],
            started_at=status["started_at"],
            completed_at=status["completed_at"],
            error=status.get("error"),
        )
    except Exception as e:
        logger.error(f"Error getting status for {mission_id}: {e}")
        raise HTTPException(status_code=404, detail=str(e))


@app.get("/mission/{mission_id}/results", response_model=SimulationResultsJSON)
def get_mission_results(mission_id: str):
    """Retrieve simulation results (available after completion).
    
    GET /mission/{mission_id}/results
    
    Returns:
      trajectory: List of state points
      telemetry: Time series of diagnostics
      summary: Key statistics
    """
    try:
        status = mission_manager.get_status(mission_id)
        if status["status"] not in ["completed", "failed"]:
            raise ValueError(f"Mission {mission_id} not completed: status={status['status']}")
        
        results = mission_manager.get_results(mission_id)
        
        # Convert trajectory to JSON-serializable format
        traj = []
        if results.trajectory:
            for step in results.trajectory:
                diags = step.diagnostics or {}
                diags["degradation_factor"] = getattr(step.state, "degradation_factor", 1.0)
                
                traj.append(TrajectoryPoint(
                    t_sec=step.time_info.t_epoch,
                    position=step.state.position,
                    velocity=step.state.velocity,
                    diagnostics=diags,
                ))
        
        return SimulationResultsJSON(
            mission_id=mission_id,
            status="success" if status["status"] == "completed" else "failed",
            config=status["config"].model_dump(),
            summary=results.summary() if results else {},
            trajectory=traj,
            telemetry=results.telemetry if results else {},
            error_log=status.get("error"),
        )
    except Exception as e:
        logger.error(f"Error getting results for {mission_id}: {e}")
        raise HTTPException(status_code=404, detail=str(e))


@app.get("/mission/list")
async def list_all_missions():
    """List all mission IDs."""
    mission_ids = mission_manager.list_missions()
    return {
        "missions": mission_ids,
        "count": len(mission_ids),
    }


# ════════════════════════════════════════════════════════════════════════════
# OPTIMIZATION ENDPOINTS
# ════════════════════════════════════════════════════════════════════════════

@app.post("/optimization/run")
async def run_optimization(req: OptimizationRequest, background_tasks: BackgroundTasks):
    """Launch a multi-objective optimization sweep.
    
    POST /optimization/run
    {
      "sweep_name": "Sail Area Optimization",
      "base_config": {...},
      "parameter_ranges": {
        "sail_area": [500, 2000],
        "reflectivity": [0.85, 0.95]
      },
      "objectives": ["max_final_distance", "min_time"],
      "num_samples": 200
    }
    
    Returns:
      sweep_id: UUID for tracking
      status: "running"
    """
    try:
        sweep_id = str(uuid.uuid4())
        
        # Store optimization task info
        mission_manager.optimizations[sweep_id] = {
            "status": "queued",
            "base_config": req.base_config.model_dump(),
            "parameter_ranges": req.parameter_ranges,
            "objectives": req.objectives,
            "num_samples": req.num_samples,
            "created_at": datetime.utcnow().isoformat(),
            "started_at": None,
            "completed_at": None,
            "pareto_front": [],
            "num_evaluated": 0,
            "error": None
        }
        
        logger.info(f"Started optimization {sweep_id}: {req.sweep_name}")
        logger.info(f"  Parameters: {list(req.parameter_ranges.keys())}")
        logger.info(f"  Objectives: {req.objectives}")
        logger.info(f"  Population size: {req.num_samples}")
        
        # Add optimization task to background
        background_tasks.add_task(
            _run_optimization_async,
            sweep_id,
            req.base_config.model_dump(),
            req.parameter_ranges,
            req.objectives,
            req.num_samples,
        )
        
        return OptimizationResult(
            sweep_id=sweep_id,
            status="queued",
            message="Optimization sweep queued for background processing",
            base_config=req.base_config.model_dump(),
            total_samples=req.num_samples
        )
    except Exception as e:
        logger.error(f"Error starting optimization: {e}")
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/optimization/{sweep_id}/status", response_model=OptimizationResult)
async def get_optimization_status(sweep_id: str):
    """Get current status of optimization sweep."""
    logger.info(f"Status query for sweep {sweep_id}")
    
    if sweep_id not in mission_manager.optimizations:
        raise HTTPException(status_code=404, detail="Optimization sweep not found")
        
    opt = mission_manager.optimizations[sweep_id]
    
    return OptimizationResult(
        sweep_id=sweep_id,
        status=opt["status"],
        base_config=opt["base_config"],
        parameter_ranges=opt["parameter_ranges"],
        objectives=opt["objectives"],
        pareto_front=opt["pareto_front"],
        num_evaluated=opt["num_evaluated"],
        total_samples=opt["num_samples"],
        error=opt.get("error")
    )


# ════════════════════════════════════════════════════════════════════════════
# HEALTH CHECK
# ════════════════════════════════════════════════════════════════════════════

@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "service": "Heliosail-RX Simulation API",
        "timestamp": datetime.utcnow().isoformat(),
    }


# ════════════════════════════════════════════════════════════════════════════
# Background Tasks
# ════════════════════════════════════════════════════════════════════════════


def _run_mission_in_process(mission_id: str, config_dict: Dict[str, Any]) -> Any:
    """Runs the simulation in a separate process to avoid GIL blocking."""
    # Move logging to module level or retrieve it
    proc_logger = logging.getLogger("worker")
    try:
        # Reconstruct mission config
        config_obj = SimulationConfigJSON(**config_dict)
        manager = MissionManager()
        mission = manager._build_mission_obj(config_obj)
        
        proc_logger.info(f"[Worker] Executing mission {mission_id}")
        result = mission.run()
        proc_logger.info(f"[Worker] Mission {mission_id} completed successfully")
        return {"status": "success", "result": result}
    except Exception as e:
        proc_logger.error(f"[Worker] Mission {mission_id} failed: {e}")
        return {"status": "error", "error": str(e)}


async def _run_simulation_async(mission_id: str, config_dict: Dict[str, Any]):
    """Background task wrapper that awaits the process pool executor."""
    try:
        loop = asyncio.get_running_loop()
        res = await loop.run_in_executor(
            process_pool, 
            _run_mission_in_process, 
            mission_id, 
            config_dict
        )
        
        if res["status"] == "success":
            mission_manager.update_mission_status(mission_id, "completed", result_obj=res["result"])
        else:
            mission_manager.update_mission_status(mission_id, "failed", error=res["error"])
            
    except Exception as e:
        logger.error(f"[BG] Error in simulation wrapper for {mission_id}: {e}")
        mission_manager.update_mission_status(mission_id, "failed", error=str(e))


def _evaluate_objectives(result, objectives: List[str]) -> Dict[str, float]:
    """Extract requested optimization objectives from a SimulationResult."""
    metrics = {}
    import math
    
    for obj in objectives:
        if obj == "max_final_distance":
            if result.trajectory:
                pos = result.trajectory[-1].state.position
                dist = math.sqrt(sum(x*x for x in pos))
                # For optimizer, we typically minimize objectives, so negative distance
                metrics[obj] = -dist
            else:
                metrics[obj] = 0.0
        elif obj == "min_time":
            metrics[obj] = result.elapsed_wall_time_sec
        elif obj == "min_fuel":
            initial_mass = result.config.initial_state.mass
            final_mass = result.trajectory[-1].state.mass if result.trajectory else initial_mass
            fuel_used = initial_mass - final_mass
            metrics[obj] = fuel_used
        else:
            # Default to 0 for unknown
            metrics[obj] = 0.0
            
    return metrics


async def _run_optimization_async(
    sweep_id: str,
    base_config: Dict[str, Any],
    parameter_ranges: Dict[str, tuple],
    objectives: List[str],
    num_samples: int,
):
    """Background task: Execute optimization sweep asynchronously across process pool."""
    from mission_optimizer.pareto_front_generator import pareto_from_candidates
    opt = mission_manager.optimizations[sweep_id]
    
    try:
        opt["status"] = "running"
        opt["started_at"] = datetime.utcnow().isoformat()
        logger.info(f"[BG] Optimization {sweep_id} evaluating {num_samples} configurations")
        
        loop = asyncio.get_running_loop()
        tasks = []
        candidates_raw = []
        
        # 1. Generate Parameters and Dispatch Sim Tasks
        for i in range(num_samples):
            # Clone base config
            sample_config = base_config.copy()
            
            # Apply random uniform perturbations to specific parameters
            for param, (p_min, p_max) in parameter_ranges.items():
                if param in sample_config:
                    sample_config[param] = random.uniform(p_min, p_max)
                    
            sample_mission_id = f"{sweep_id}_sample_{i}"
            candidates_raw.append({"config": sample_config, "id": sample_mission_id})
            
            # Queue to process pool
            tasks.append(loop.run_in_executor(
                process_pool, 
                _run_mission_in_process, 
                sample_mission_id, 
                sample_config
            ))
            
        # 2. Parallel Evaluation
        evaluations = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 3. Assemble Candidates and compute Objectives
        candidates = []
        for i, res in enumerate(evaluations):
            if isinstance(res, Exception):
                logger.error(f"[BG] Optimization sample failed: {res}")
                continue
                
            if res.get("status") == "success":
                sim_res = res["result"]
                metrics = _evaluate_objectives(sim_res, objectives)
                
                # Maintain order of objectives for pareto comparison
                obj_vector = [metrics[obj] for obj in objectives]
                
                candidates.append({
                    "id": candidates_raw[i]["id"],
                    "parameters": {k: candidates_raw[i]["config"][k] for k in parameter_ranges.keys()},
                    "objectives": obj_vector,
                    "metrics": metrics
                })
                opt["num_evaluated"] += 1

        # 4. Generate Pareto Front
        logger.info(f"[BG] Computing Pareto front from {opt['num_evaluated']} successful candidates")
        pareto = pareto_from_candidates(candidates)
        
        # Clean up negative values used for 'max_' optimization objectives prior to saving
        for p in pareto:
            for obj in objectives:
                if obj.startswith("max_"):
                    p["metrics"][obj] = abs(p["metrics"][obj])

        opt["pareto_front"] = pareto
        opt["status"] = "completed"
        opt["completed_at"] = datetime.utcnow().isoformat()
        logger.info(f"[BG] Optimization {sweep_id} completed. Found {len(pareto)} Pareto optimal solutions.")
        
    except Exception as e:
        logger.error(f"[BG] Optimization {sweep_id} failed: {e}")
        opt["status"] = "failed"
        opt["error"] = str(e)
        opt["completed_at"] = datetime.utcnow().isoformat()


# ════════════════════════════════════════════════════════════════════════════
# Entry Point
# ════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import uvicorn
    
    logging.basicConfig(level=logging.INFO)
    
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        log_level="info",
    )