import logging
import subprocess
from typing import Any, Dict, Optional

from mcp.server.fastmcp import FastMCP

from utils.traci import ensure_traci_start_stdout_suppressed
from mcp_tools.simulation import run_simple_simulation
from mcp_tools.network import netconvert, netgenerate, osm_get
from mcp_tools.route import random_trips, duarouter, od2trips
from mcp_tools.signal import tls_cycle_adaptation, tls_coordinator
from mcp_tools.analysis import analyze_fcd
from mcp_tools.vehicle import (
    get_vehicles, get_vehicle_speed, get_vehicle_position, 
    get_vehicle_acceleration, get_vehicle_lane, get_vehicle_route,
    get_simulation_info
)
from mcp_tools.rl import find_sumo_rl_scenario_files, list_rl_scenarios, run_rl_training
from utils.connection import connection_manager
from utils.sumo import find_sumo_binary, find_sumo_home, find_sumo_tools_dir
from workflows.sim_gen import sim_gen_workflow
from workflows.signal_opt import signal_opt_workflow
from workflows.rl_train import rl_train_workflow

# Configure logging to stderr to not interfere with MCP stdio transport
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Ensure TraCI never writes to stdout by default (MCP stdio safety).
ensure_traci_start_stdout_suppressed()

# Initialize MCP Server (official SDK)
server = FastMCP("SUMO-MCP-Server")

# --- 1. Network Management ---
@server.tool(description="Manage SUMO network (generate, convert, or download OSM).")
def manage_network(action: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    actions:
    - generate: params={'grid': bool, 'grid_number': int}
    - convert: params={'osm_file': str}
    - download_osm: output_file is treated as output_dir. params={'bbox': str, 'prefix': str}
    """
    params = params or {}
    options = params.get("options")
    
    if action == "generate":
        spider = bool(params.get("spider", False))
        grid = bool(params.get("grid", True))
        grid_number = params.get("grid_number", 3)

        if spider:
            # Spider network takes precedence over grid settings.
            grid = False
            options_list = list(options or [])

            def _strip_flag(flag: str, has_value: bool = False) -> None:
                while flag in options_list:
                    idx = options_list.index(flag)
                    options_list.pop(idx)
                    if has_value and idx < len(options_list):
                        options_list.pop(idx)

            def _set_option(flag: str, value: str) -> None:
                if flag in options_list:
                    idx = options_list.index(flag)
                    if idx + 1 < len(options_list):
                        options_list[idx + 1] = value
                    else:
                        options_list.append(value)
                else:
                    options_list.extend([flag, value])

            # Enforce Spider/Grid mutual exclusion even when the user provided `options`.
            _strip_flag("--grid")
            _strip_flag("--grid.number", has_value=True)

            if "--spider" not in options_list:
                options_list.insert(0, "--spider")

            arms_raw = params.get("arms", params.get("arm_number"))
            if arms_raw is not None:
                try:
                    arms = int(arms_raw)
                except (TypeError, ValueError):
                    return f"Error: arms must be a positive integer, got {arms_raw!r}"
                if arms <= 0:
                    return "Error: arms must be > 0"
                _set_option("--spider.arm-number", str(arms))

            circles_raw = params.get("circles", params.get("circle_number"))
            if circles_raw is not None:
                try:
                    circles = int(circles_raw)
                except (TypeError, ValueError):
                    return f"Error: circles must be a positive integer, got {circles_raw!r}"
                if circles <= 0:
                    return "Error: circles must be > 0"
                _set_option("--spider.circle-number", str(circles))

            space_radius_raw = params.get("ring_radius", params.get("space_radius"))
            if space_radius_raw is not None:
                try:
                    space_radius = float(space_radius_raw)
                except (TypeError, ValueError):
                    return f"Error: ring_radius must be a number, got {space_radius_raw!r}"
                if space_radius <= 0:
                    return "Error: ring_radius must be > 0"
                _set_option("--spider.space-radius", str(space_radius))

            attach_length_raw = params.get("radial_distance", params.get("attach_length"))
            if attach_length_raw is not None:
                try:
                    attach_length = float(attach_length_raw)
                except (TypeError, ValueError):
                    return f"Error: radial_distance must be a number, got {attach_length_raw!r}"
                if attach_length < 0:
                    return "Error: radial_distance must be >= 0"
                _set_option("--spider.attach-length", str(attach_length))

            omit_center_raw = params.get("omit_center")
            if omit_center_raw:
                if "--spider.omit-center" not in options_list:
                    options_list.append("--spider.omit-center")

            options = options_list

        return netgenerate(output_file, grid, grid_number, options)
        
    elif action == "convert" or action == "convert_osm":
        osm_file = params.get("osm_file")
        if not osm_file: return "Error: osm_file required for convert action"
        return netconvert(osm_file, output_file, options)
        
    elif action == "download_osm":
        # output_file here acts as output_dir
        bbox = params.get("bbox")
        prefix = params.get("prefix", "osm")
        if not bbox: return "Error: bbox required for download_osm action"
        return osm_get(bbox, output_file, prefix, options)
        
    return f"Unknown action: {action}"

# --- 2. Demand Management ---
@server.tool(description="Manage traffic demand (random trips, OD matrix, routing).")
def manage_demand(action: str, net_file: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    actions:
    - generate_random: params={'end_time': int, 'period': float}
    - convert_od: params={'od_file': str} (net_file unused but kept for consistency)
    - compute_routes: params={'route_files': str} (input trips)
    """
    params = params or {}
    options = params.get("options")
    
    if action == "generate_random" or action == "random_trips":
        # Backward/compat aliases: some clients use `end` instead of `end_time`.
        end_time_raw = params.get("end_time", params.get("end", 3600))
        period_raw = params.get("period", 1.0)
        try:
            end_time = int(end_time_raw)
        except (TypeError, ValueError):
            return f"Error: end_time must be an integer, got {end_time_raw!r}"
        try:
            period = float(period_raw)
        except (TypeError, ValueError):
            return f"Error: period must be a number, got {period_raw!r}"
        return random_trips(net_file, output_file, end_time, period, options)
        
    elif action == "convert_od" or action == "od_matrix":
        od_file = params.get("od_file")
        if not od_file: return "Error: od_file required for convert_od"
        return od2trips(od_file, output_file, options)
        
    elif action == "compute_routes" or action == "routing":
        route_files = params.get("route_files") # Input trips file
        if not route_files: return "Error: route_files required for compute_routes"
        return duarouter(net_file, route_files, output_file, options)
        
    return f"Unknown action: {action}"

# --- 3. Simulation Control ---
@server.tool(description="Control SUMO simulation (connect, step, disconnect).")
def control_simulation(action: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    actions:
    - connect: params={'config_file': str, 'gui': bool}
    - step: params={'step': float}
    - disconnect: no params
    """
    params = params or {}
    
    try:
        timeout_s_raw = params.get("timeout_s", params.get("timeout"))
        timeout_s: Optional[float] = None
        if timeout_s_raw is not None:
            try:
                timeout_s = float(timeout_s_raw)
            except (TypeError, ValueError):
                return f"Error: timeout_s must be a number, got {timeout_s_raw!r}"

        if action == "connect":
            config_file = params.get("config_file")
            gui = params.get("gui", False)
            port = params.get("port", 8813)
            host = params.get("host", "localhost")
            if timeout_s is None:
                connection_manager.connect(config_file, gui, port, host)
            else:
                connection_manager.connect(config_file, gui, port, host, timeout_s=timeout_s)
            return "Successfully connected to SUMO."
            
        elif action == "step":
            step = params.get("step", 0)
            if timeout_s is None:
                connection_manager.simulation_step(step)
            else:
                connection_manager.simulation_step(step, timeout_s=timeout_s)
            return "Simulation advanced."
            
        elif action == "disconnect":
            if timeout_s is None:
                connection_manager.disconnect()
            else:
                connection_manager.disconnect(timeout_s=timeout_s)
            return "Successfully disconnected from SUMO."
            
    except Exception as e:
        return f"Error in control_simulation ({action}): {type(e).__name__}: {e}"
        
    return f"Unknown action: {action}"

# --- 4. Query State ---
@server.tool(description="Query simulation state (vehicles, speed, position). Requires active connection.")
def query_simulation_state(target: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    targets:
    - vehicle_list: no params
    - vehicle_variable: params={'vehicle_id': str, 'variable': 'speed'|'position'|'lane'|'acceleration'|'route'}
    """
    params = params or {}
    
    try:
        if target == "vehicle_list" or target == "vehicles":
            vehs = get_vehicles()
            return f"Active vehicles: {vehs}"
            
        elif target == "vehicle_variable":
            v_id = params.get("vehicle_id")
            var = params.get("variable")
            if not v_id or not var: return "Error: vehicle_id and variable required"
            
            if var == "speed": return f"Speed: {get_vehicle_speed(v_id)}"
            if var == "position": return f"Position: {get_vehicle_position(v_id)}"
            if var == "acceleration": return f"Acceleration: {get_vehicle_acceleration(v_id)}"
            if var == "lane": return f"Lane: {get_vehicle_lane(v_id)}"
            if var == "route": return f"Route: {get_vehicle_route(v_id)}"
            
            return f"Unknown variable: {var}"
        
        elif target == "simulation":
            info = get_simulation_info()
            return f"Simulation Info: {info}"
            
    except Exception as e:
        return f"Error querying state: {type(e).__name__}: {e}"
        
    return f"Unknown target: {target}"

# --- 5. Optimize Signals ---
@server.tool(description="Optimize traffic signals.")
def optimize_traffic_signals(method: str, net_file: str, route_file: str, output_file: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    methods:
    - cycle_adaptation: adapt TLS cycles
    - coordination: TLS coordination
    """
    params = params or {}
    options = params.get("options")
    
    if method == "cycle_adaptation" or method == "Websters":
        return tls_cycle_adaptation(net_file, route_file, output_file)
    elif method == "coordination":
        return tls_coordinator(net_file, route_file, output_file, options)
        
    return f"Unknown method: {method}"

# --- 6. Workflows ---
@server.tool(
    description="""Run high-level SUMO workflows. Available workflows:

**sim_gen_eval** - Generate grid network, simulate traffic, analyze results.
  params:
  - grid_number (int): Grid size NxN. Default=3. Aliases: grid_size, size
  - sim_seconds (int): Simulation duration in seconds. Default=100. Aliases: steps, duration, end_time
  - output_dir (str): Output directory. Default="output"
  Example: run_workflow("sim_gen_eval", {"grid_number": 3, "sim_seconds": 1000})

**signal_opt** - Optimize traffic signals for existing network.
  params:
  - net_file (str): Path to .net.xml file. REQUIRED
  - route_file (str): Path to .rou.xml file. REQUIRED
  - sim_seconds (int): Simulation duration. Default=3600. Aliases: steps, duration
  - use_coordinator (bool): Use tlsCoordinator instead of tlsCycleAdaptation. Default=false
  - output_dir (str): Output directory. Default="output"

**rl_train** - Train RL agent for traffic signal control.
  params:
  - scenario_name (str): Built-in scenario name (use manage_rl_task("list_scenarios") to see options). Aliases: scenario
  - episodes (int): Number of training episodes. Default=5. Aliases: num_episodes
  - steps (int): Steps per episode. Default=1000. Aliases: steps_per_episode
  - output_dir (str): Output directory. Default="output"
"""
)
def run_workflow(workflow_name: str, params: Dict[str, Any]) -> str:
    """Execute a high-level workflow."""

    # Helper to get param with aliases
    def get_param(keys: list, default=None):
        for k in keys:
            if k in params:
                return params[k]
        return default

    if workflow_name in ("sim_gen_eval", "sim_gen_workflow", "sim_gen"):
        grid_number = get_param(["grid_number", "grid_size", "size"], 3)
        sim_seconds = get_param(["sim_seconds", "steps", "duration", "end_time"], 100)
        output_dir = get_param(["output_dir"], "output")

        return sim_gen_workflow(output_dir, int(grid_number), int(sim_seconds))

    elif workflow_name in ("signal_opt", "signal_opt_workflow"):
        net_file = get_param(["net_file"], "")
        route_file = get_param(["route_file"], "")

        if not net_file or not route_file:
            return "Error: signal_opt requires net_file and route_file parameters."

        sim_seconds = get_param(["sim_seconds", "steps", "duration"], 3600)
        use_coordinator = get_param(["use_coordinator"], False)
        output_dir = get_param(["output_dir"], "output")

        return signal_opt_workflow(net_file, route_file, output_dir, int(sim_seconds), bool(use_coordinator))

    elif workflow_name == "rl_train":
        scenario_name = get_param(["scenario_name", "scenario"], "")
        episodes = get_param(["episodes", "num_episodes"], 5)
        steps = get_param(["steps", "steps_per_episode"], 1000)
        output_dir = get_param(["output_dir"], "output")

        return rl_train_workflow(scenario_name, output_dir, int(episodes), int(steps))

    return f"Unknown workflow: {workflow_name}. Available: sim_gen_eval, signal_opt, rl_train"

# --- 7. RL Task Management ---
@server.tool(description="Manage RL tasks (list scenarios, custom training).")
def manage_rl_task(action: str, params: Optional[Dict[str, Any]] = None) -> str:
    """
    actions:
    - list_scenarios: no params
    - train_custom: params={'net_file', 'route_file', 'out_dir', 'episodes', 'steps', 'algorithm', 'reward_type'}
    """
    params = params or {}
    
    if action == "list_scenarios":
        return str(list_rl_scenarios())
        
    elif action == "train_custom":
        scenario_name = params.get("scenario") or params.get("scenario_name")
        net_file = params.get("net_file")
        route_file = params.get("route_file")

        if scenario_name:
            net_file, route_file, err = find_sumo_rl_scenario_files(str(scenario_name))
            if err:
                return err

        if not net_file or not route_file:
            return (
                "Error: train_custom requires either:\n"
                "  - scenario/scenario_name (built-in sumo-rl scenario), OR\n"
                "  - net_file + route_file (custom files)\n"
                "Hint: Use manage_rl_task(list_scenarios) to see available built-in scenarios."
            )

        out_dir = params.get("out_dir") or params.get("output_dir") or "output"

        episodes_raw = params.get("episodes", params.get("num_episodes", 1))
        steps_raw = params.get("steps", params.get("steps_per_episode", 1000))
        try:
            episodes = int(episodes_raw)
        except (TypeError, ValueError):
            return f"Error: episodes must be an integer, got {episodes_raw!r}"
        try:
            steps_per_episode = int(steps_raw)
        except (TypeError, ValueError):
            return f"Error: steps must be an integer, got {steps_raw!r}"

        if episodes <= 0:
            return "Error: episodes must be > 0"
        if steps_per_episode <= 0:
            return "Error: steps must be > 0"

        algorithm = str(params.get("algorithm", "ql"))
        reward_type = str(params.get("reward_type", "diff-waiting-time"))

        return run_rl_training(
            net_file=str(net_file),
            route_file=str(route_file),
            out_dir=str(out_dir),
            episodes=episodes,
            steps_per_episode=steps_per_episode,
            algorithm=algorithm,
            reward_type=reward_type,
        )
        
    return f"Unknown action: {action}"

# --- Legacy/Misc ---
@server.tool(name="get_sumo_info", description="Get the version and path of the installed SUMO.")
def get_sumo_info() -> str:
    try:
        sumo_binary = find_sumo_binary("sumo")
        if not sumo_binary:
            return (
                "Error: Could not locate SUMO executable. "
                "Please ensure SUMO is installed and either `sumo` is available in PATH or `SUMO_HOME` is set."
            )

        result = subprocess.run(
            [sumo_binary, "--version"],
            capture_output=True,
            text=True,
            check=True,
            timeout=10,
        )
        version_output = (result.stdout.splitlines() or ["Unknown"])[0]

        sumo_home = find_sumo_home()
        tools_dir = find_sumo_tools_dir()
        return "\n".join(
            [
                f"SUMO Binary: {sumo_binary}",
                f"SUMO Version: {version_output}",
                f"SUMO_HOME: {sumo_home or 'Not Set'}",
                f"SUMO Tools Dir: {tools_dir or 'Not Found'}",
            ]
        )
    except Exception as e:
        return f"Error checking SUMO: {str(e)}"

@server.tool(name="run_simple_simulation", description="Run a SUMO simulation using a config file.")
def run_simple_simulation_tool(config_path: str, steps: int = 100) -> str:
    return run_simple_simulation(config_path, steps)

@server.tool(description="Analyze FCD output.")
def run_analysis(fcd_file: str) -> str:
    return analyze_fcd(fcd_file)

if __name__ == "__main__":
    # NOTE:
    # MCP stdio transport relies on AnyIO/asyncio to process thread callbacks.
    # In some environments, a lack of scheduled timers can cause the event loop to
    # block indefinitely while waiting for stdio worker-thread results. A small
    # periodic sleep keeps the loop responsive without emitting any stdout output.
    import anyio

    async def _wakeup_task() -> None:
        while True:
            await anyio.sleep(0.1)

    async def _run_stdio_with_wakeup() -> None:
        async with anyio.create_task_group() as tg:
            tg.start_soon(_wakeup_task)
            await server.run_stdio_async()

    anyio.run(_run_stdio_with_wakeup)
