import asyncio
import os
import sys
import subprocess
import time
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
import logging
import threading

# --- Sys Path Injection ---
# Add ai_agent root to sys.path
current_dir = os.path.dirname(os.path.abspath(__file__))
# Go up 2 levels: src -> groundcontrol -> ai_agent
ai_agent_root = os.path.abspath(os.path.join(current_dir, "../.."))
# Also add models directory specifically for implicit imports if needed
models_dir = os.path.join(ai_agent_root, "models")

if ai_agent_root not in sys.path:
    sys.path.append(ai_agent_root)
if models_dir not in sys.path:
    sys.path.append(models_dir)

# Now we can import from ai_agent modules
from models.models_client import Models_Client
from tools.core.common_functions import read_json_file

from .models import GeneratePlanRequest, ExecuteMissionRequest
from .websocket_manager import websocket_manager
from .py_tree_generator import py_tree_generator

# --- Global Variables ---
model_server_process = None
MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL", "http://localhost:8000")

# --- 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'))
if not os.path.exists(static_dir):
    os.makedirs(static_dir)
app.mount("/static", StaticFiles(directory=static_dir), name="static")

def start_model_server():
    """Starts the model server if it's not already running."""
    global model_server_process
    client = Models_Client(MODEL_SERVER_URL)
    
    # Check if already running
    try:
        logging.info(f"Checking model server health at {MODEL_SERVER_URL}...")
        health = client.check_health()
        if health.get("status") in ["healthy", "loading"]:
            logging.info("Model server is already running.")
            return
    except Exception:
        pass # Not running

    logging.info("Starting model server...")
    
    # 读取配置文件
    try:
        config_json_file = os.path.join(ai_agent_root, "config", "model_config.json")
        config_data = read_json_file(config_json_file)
        if config_data is None or "models_server" not in config_data:
            logging.error(f"无法读取配置文件或配置文件中缺少 'models_server' 字段: {config_json_file}")
            return
        
        # 如果 models_server 是列表，取第一个；如果是字典，直接使用
        models_server_data = config_data["models_server"]
        model_server_config = models_server_data[0] if isinstance(models_server_data, list) else models_server_data
        
        host = model_server_config.get("model_server_host", "localhost")
        port = model_server_config.get("model_server_port", 8000)
        workers = model_server_config.get("model_controller_workers", 1)
    except Exception as e:
        logging.error(f"读取配置文件失败: {e}，使用默认配置")
        host = "localhost"
        port = 8000
        workers = 1
    
    # Start as subprocess using uvicorn
    try:
        def log_output(pipe, prefix):
            """实时输出子进程的日志"""
            try:
                for line in iter(pipe.readline, b''):
                    if line:
                        line_str = line.decode('utf-8', errors='replace').strip()
                        if line_str:
                            logging.info(f"[Model Server {prefix}] {line_str}")
            except Exception as e:
                logging.error(f"Error reading {prefix}: {e}")
            finally:
                pipe.close()
        
        # 使用 uvicorn 启动模型服务器
        # 使用 models_server:create_app 作为应用入口
        uvicorn_cmd = [
            sys.executable, "-m", "uvicorn",
            "models.models_server:create_app",
            "--host", host,
            "--port", str(port),
            "--workers", str(workers),
            "--factory"
        ]
        
        model_server_process = subprocess.Popen(
            uvicorn_cmd,
            cwd=ai_agent_root,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        
        # 启动线程实时输出日志
        stdout_thread = threading.Thread(
            target=log_output, 
            args=(model_server_process.stdout, "STDOUT"),
            daemon=True
        )
        stderr_thread = threading.Thread(
            target=log_output, 
            args=(model_server_process.stderr, "STDERR"),
            daemon=True
        )
        stdout_thread.start()
        stderr_thread.start()
        
        # 先等待一小段时间，检查进程是否立即退出
        time.sleep(2)
        if model_server_process.poll() is not None:
            # 进程已退出，读取错误信息
            out, err = model_server_process.communicate()
            error_msg = err.decode('utf-8', errors='replace') if err else out.decode('utf-8', errors='replace') if out else "Unknown error"
            logging.error(f"Model server process exited immediately: {error_msg}")
            return
        
        # Wait for it to be ready
        logging.info("Waiting for model server to be ready...")
        if client.wait_for_server(timeout=300, check_interval=5):
            logging.info("Model server started and ready.")
        else:
            logging.error("Model server failed to start within timeout.")
            # Check for errors
            if model_server_process.poll() is not None:
                out, err = model_server_process.communicate()
                error_msg = err.decode('utf-8', errors='replace') if err else out.decode('utf-8', errors='replace') if out else 'Unknown error'
                logging.error(f"Model server process exited with error: {error_msg}")
    except Exception as e:
        logging.error(f"Failed to launch model server: {e}")

def stop_model_server():
    """Stops the model server subprocess if it was started by this service."""
    global model_server_process
    if model_server_process:
        logging.info("Stopping model server subprocess...")
        model_server_process.terminate()
        try:
            model_server_process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            model_server_process.kill()
        logging.info("Model server stopped.")

# --- 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 execution removed
    return {"status": "execution_started (simulation mode)"}

@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.
    Start the ROS2 node in a background thread.
    Ensure the Model Server is running.
    """
    # Configure WebSocket Manager
    loop = asyncio.get_running_loop()
    websocket_manager.set_loop(loop)
    logging.info("WebSocket event loop configured.")

    # Start Model Server
    threading.Thread(target=start_model_server, daemon=True).start()

@app.on_event("shutdown")
async def shutdown_event():
    logging.info("Backend service shutting down.")
    stop_model_server()
    logging.info("Backend service shut down successfully.")
