use std::collections::{HashMap, HashSet};

use anyhow::{bail, Result};
use serde::{Deserialize, Serialize};

use crate::{
    config::{HarnessConfig, Profile},
    transport::StreamTransportBackend,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
#[serde(rename_all = "kebab-case")]
pub enum ModuleState {
    Planned,
    Starting,
    Running,
    Stopped,
    Failed,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
#[serde(rename_all = "kebab-case")]
pub enum StreamDirection {
    Input,
    Output,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
pub struct StreamDescriptor {
    pub name: String,
    pub direction: StreamDirection,
    pub message_type: String,
    pub transport: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
pub struct ModuleHealth {
    pub ok: bool,
    pub state: ModuleState,
    pub message: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
pub struct ModuleInfo {
    pub id: usize,
    pub name: String,
    pub module_type: String,
    pub state: ModuleState,
    pub health: ModuleHealth,
    pub dependencies: Vec<String>,
    pub inputs: Vec<StreamDescriptor>,
    pub outputs: Vec<StreamDescriptor>,
    pub capabilities: Vec<String>,
    pub physical: bool,
    pub required: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[cfg_attr(feature = "mcp", derive(schemars::JsonSchema))]
pub struct ModuleGraph {
    pub modules: Vec<ModuleInfo>,
}

impl ModuleGraph {
    pub fn new(modules: Vec<ModuleInfo>) -> Result<Self> {
        let mut names = HashSet::with_capacity(modules.len());
        let mut ids = HashSet::with_capacity(modules.len());
        for module in &modules {
            if !ids.insert(module.id) {
                bail!("duplicate module id '{}'", module.id);
            }
            if module.name.trim().is_empty() {
                bail!("module name cannot be empty");
            }
            if !names.insert(module.name.clone()) {
                bail!("duplicate module name '{}'", module.name);
            }
        }
        for module in &modules {
            for dependency in &module.dependencies {
                if dependency == &module.name {
                    bail!("module '{}' cannot depend on itself", module.name);
                }
                if !names.contains(dependency) {
                    bail!(
                        "module '{}' depends on unknown module '{}'",
                        module.name,
                        dependency
                    );
                }
            }
        }
        Ok(Self { modules })
    }

    pub fn modules(&self) -> &[ModuleInfo] {
        &self.modules
    }

    pub fn stream_count(&self) -> usize {
        self.modules
            .iter()
            .map(|module| module.inputs.len() + module.outputs.len())
            .sum()
    }

    pub fn to_dot(&self) -> String {
        let mut dot = String::from("digraph leash_module_graph {\n");
        dot.push_str("  graph [rankdir=LR];\n");
        dot.push_str("  node [shape=box];\n");

        for module in &self.modules {
            let risk = if module.physical {
                "physical"
            } else {
                "non-physical"
            };
            dot.push_str(&format!(
                "  {} [label=\"{}\\n{}\\n{}\", color=\"{}\"];\n",
                dot_id(module.id),
                escape_dot(&module.name),
                escape_dot(&module.module_type),
                risk,
                if module.physical { "red" } else { "black" }
            ));
        }

        for from in &self.modules {
            for output in &from.outputs {
                for to in &self.modules {
                    if from.id == to.id {
                        continue;
                    }
                    for input in &to.inputs {
                        if output.name == input.name && output.message_type == input.message_type {
                            dot.push_str(&format!(
                                "  {} -> {} [label=\"{}:{} via {}\"];\n",
                                dot_id(from.id),
                                dot_id(to.id),
                                escape_dot(&output.name),
                                escape_dot(&output.message_type),
                                escape_dot(&output.transport)
                            ));
                        }
                    }
                }
            }
        }

        dot.push_str("}\n");
        dot
    }
}

#[derive(Debug, Clone)]
pub struct ModuleCoordinator {
    modules: Vec<ModuleInfo>,
    start_failures: HashMap<String, String>,
}

impl ModuleCoordinator {
    pub fn new(graph: ModuleGraph) -> Self {
        Self {
            modules: graph.modules,
            start_failures: HashMap::new(),
        }
    }

    pub fn with_start_failure(mut self, module_name: &str, message: impl Into<String>) -> Self {
        self.start_failures
            .insert(module_name.to_string(), message.into());
        self
    }

    pub fn start(&mut self) -> Result<()> {
        let order = self.dependency_order()?;
        let mut started = Vec::new();

        for index in order {
            self.set_health(index, ModuleState::Starting, true, "starting");
            let module_name = self.modules[index].name.clone();
            if let Some(message) = self.start_failures.get(&module_name).cloned() {
                self.set_health(index, ModuleState::Failed, false, message.clone());
                self.stop_started_after_failure(&started);
                bail!("module '{}' failed to start: {}", module_name, message);
            }
            self.set_health(index, ModuleState::Running, true, "running");
            started.push(index);
        }

        Ok(())
    }

    pub fn stop(&mut self) {
        let Ok(order) = self.dependency_order() else {
            for index in (0..self.modules.len()).rev() {
                if matches!(
                    self.modules[index].state,
                    ModuleState::Starting | ModuleState::Running
                ) {
                    self.set_health(index, ModuleState::Stopped, true, "stopped");
                }
            }
            return;
        };
        for index in order.into_iter().rev() {
            if matches!(
                self.modules[index].state,
                ModuleState::Starting | ModuleState::Running
            ) {
                self.set_health(index, ModuleState::Stopped, true, "stopped");
            }
        }
    }

    pub fn graph(&self) -> ModuleGraph {
        ModuleGraph::new(self.modules.clone()).expect("coordinator preserves valid module graph")
    }

    pub fn modules(&self) -> &[ModuleInfo] {
        &self.modules
    }

    pub fn is_healthy(&self) -> bool {
        self.modules
            .iter()
            .all(|module| module.health.ok && module.state == ModuleState::Running)
    }

    fn dependency_order(&self) -> Result<Vec<usize>> {
        #[derive(Clone, Copy, PartialEq, Eq)]
        enum Visit {
            New,
            Active,
            Done,
        }

        let by_name: HashMap<&str, usize> = self
            .modules
            .iter()
            .enumerate()
            .map(|(index, module)| (module.name.as_str(), index))
            .collect();
        let mut visits = vec![Visit::New; self.modules.len()];
        let mut order = Vec::with_capacity(self.modules.len());

        fn visit(
            index: usize,
            modules: &[ModuleInfo],
            by_name: &HashMap<&str, usize>,
            visits: &mut [Visit],
            order: &mut Vec<usize>,
        ) -> Result<()> {
            match visits[index] {
                Visit::Done => return Ok(()),
                Visit::Active => bail!("module dependency cycle at '{}'", modules[index].name),
                Visit::New => {}
            }

            visits[index] = Visit::Active;
            for dependency in &modules[index].dependencies {
                let Some(dependency_index) = by_name.get(dependency.as_str()).copied() else {
                    bail!(
                        "module '{}' depends on unknown module '{}'",
                        modules[index].name,
                        dependency
                    );
                };
                visit(dependency_index, modules, by_name, visits, order)?;
            }
            visits[index] = Visit::Done;
            order.push(index);
            Ok(())
        }

        for index in 0..self.modules.len() {
            visit(index, &self.modules, &by_name, &mut visits, &mut order)?;
        }
        Ok(order)
    }

    fn stop_started_after_failure(&mut self, started: &[usize]) {
        for index in started.iter().copied().rev() {
            self.set_health(
                index,
                ModuleState::Stopped,
                true,
                "stopped after startup failure",
            );
        }
    }

    fn set_health(
        &mut self,
        index: usize,
        state: ModuleState,
        ok: bool,
        message: impl Into<String>,
    ) {
        let module = &mut self.modules[index];
        module.state = state;
        module.health = ModuleHealth {
            ok,
            state,
            message: Some(message.into()),
        };
    }
}

pub fn default_module_graph(config: &HarnessConfig, capabilities: Vec<String>) -> ModuleGraph {
    let driver_name = match config.profile {
        Profile::Sim => "sim-driver",
        Profile::Replay => "replay-driver",
        Profile::WaveshareUgv => "waveshare-ugv-driver",
    };
    let transport = config.stream_transport;
    ModuleGraph::new(vec![
        ModuleInfo {
            id: 0,
            name: "harness-runtime".to_string(),
            module_type: "runtime".to_string(),
            state: ModuleState::Planned,
            health: planned_health(),
            dependencies: Vec::new(),
            inputs: vec![stream(
                "capability_request",
                StreamDirection::Input,
                "json",
                transport,
            )],
            outputs: vec![
                stream("health", StreamDirection::Output, "Health", transport),
                stream(
                    "capabilities",
                    StreamDirection::Output,
                    "Capabilities",
                    transport,
                ),
            ],
            capabilities,
            physical: false,
            required: true,
        },
        ModuleInfo {
            id: 1,
            name: driver_name.to_string(),
            module_type: "driver".to_string(),
            state: ModuleState::Planned,
            health: planned_health(),
            dependencies: vec!["harness-runtime".to_string()],
            inputs: vec![stream(
                "drive_command",
                StreamDirection::Input,
                "DriveReq",
                transport,
            )],
            outputs: vec![stream(
                "odometry",
                StreamDirection::Output,
                "OdometryStatus",
                transport,
            )],
            capabilities: vec!["drive".to_string(), "stop".to_string(), "estop".to_string()],
            physical: config.profile.is_physical(),
            required: true,
        },
        ModuleInfo {
            id: 2,
            name: "telemetry".to_string(),
            module_type: "telemetry".to_string(),
            state: ModuleState::Planned,
            health: planned_health(),
            dependencies: vec![driver_name.to_string()],
            inputs: vec![stream(
                "odometry",
                StreamDirection::Input,
                "OdometryStatus",
                transport,
            )],
            outputs: vec![
                stream(
                    "telemetry",
                    StreamDirection::Output,
                    "TelemetryFrame",
                    transport,
                ),
                stream(
                    "sensors",
                    StreamDirection::Output,
                    "SensorSnapshot",
                    transport,
                ),
            ],
            capabilities: vec!["observe".to_string(), "capture".to_string()],
            physical: false,
            required: true,
        },
    ])
    .expect("default module graph uses unique non-empty names")
}

fn dot_id(id: usize) -> String {
    format!("module_{id}")
}

fn escape_dot(value: &str) -> String {
    value.replace('\\', "\\\\").replace('"', "\\\"")
}

fn stream(
    name: &str,
    direction: StreamDirection,
    message_type: &str,
    transport: StreamTransportBackend,
) -> StreamDescriptor {
    StreamDescriptor {
        name: name.to_string(),
        direction,
        message_type: message_type.to_string(),
        transport: transport.as_str().to_string(),
    }
}

fn planned_health() -> ModuleHealth {
    ModuleHealth {
        ok: true,
        state: ModuleState::Planned,
        message: Some("planned".to_string()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn module(name: &str) -> ModuleInfo {
        ModuleInfo {
            id: 0,
            name: name.to_string(),
            module_type: "test".to_string(),
            state: ModuleState::Planned,
            health: planned_health(),
            dependencies: Vec::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
            capabilities: Vec::new(),
            physical: false,
            required: true,
        }
    }

    fn module_with_id(id: usize, name: &str) -> ModuleInfo {
        ModuleInfo { id, ..module(name) }
    }

    #[test]
    fn rejects_duplicate_module_names() {
        let err = ModuleGraph::new(vec![module_with_id(0, "a"), module_with_id(1, "a")])
            .unwrap_err()
            .to_string();
        assert!(err.contains("duplicate module name"));
    }

    #[test]
    fn rejects_duplicate_module_ids() {
        let err = ModuleGraph::new(vec![module_with_id(0, "a"), module_with_id(0, "b")])
            .unwrap_err()
            .to_string();
        assert!(err.contains("duplicate module id"));
    }

    #[test]
    fn counts_declared_streams() {
        let graph = default_module_graph(&HarnessConfig::default(), vec!["health".to_string()]);
        assert_eq!(graph.modules().len(), 3);
        assert_eq!(graph.stream_count(), 8);
    }

    #[test]
    fn exports_dot_with_edges_and_risk_metadata() {
        let graph = default_module_graph(
            &HarnessConfig {
                profile: Profile::WaveshareUgv,
                ..HarnessConfig::default()
            },
            vec!["drive".to_string()],
        );
        let dot = graph.to_dot();
        assert!(dot.contains("digraph leash_module_graph"));
        assert!(dot.contains("waveshare-ugv-driver"));
        assert!(dot.contains("physical"));
        assert!(dot.contains("module_1 -> module_2"));
        assert!(dot.contains("odometry:OdometryStatus via local-pubsub"));
    }

    #[test]
    fn graph_stream_transport_can_switch_backends() {
        let graph = default_module_graph(
            &HarnessConfig {
                stream_transport: StreamTransportBackend::Memory,
                ..HarnessConfig::default()
            },
            vec!["observe".to_string()],
        );

        assert!(graph
            .modules()
            .iter()
            .flat_map(|module| module.inputs.iter().chain(module.outputs.iter()))
            .all(|stream| stream.transport == "memory"));
    }

    #[test]
    fn coordinator_starts_in_dependency_order() {
        let graph = default_module_graph(&HarnessConfig::default(), vec!["health".to_string()]);
        let mut coordinator = ModuleCoordinator::new(graph);
        coordinator.start().unwrap();

        let modules = coordinator.modules();
        assert_eq!(modules[0].state, ModuleState::Running);
        assert_eq!(modules[1].state, ModuleState::Running);
        assert_eq!(modules[2].state, ModuleState::Running);
        assert!(coordinator.is_healthy());
    }

    #[test]
    fn coordinator_cleans_up_started_modules_after_start_failure() {
        let graph = default_module_graph(&HarnessConfig::default(), vec!["drive".to_string()]);
        let mut coordinator =
            ModuleCoordinator::new(graph).with_start_failure("telemetry", "loop refused");

        let err = coordinator.start().unwrap_err().to_string();
        assert!(err.contains("telemetry"));

        let modules = coordinator.modules();
        assert_eq!(modules[0].state, ModuleState::Stopped);
        assert_eq!(modules[1].state, ModuleState::Stopped);
        assert_eq!(modules[2].state, ModuleState::Failed);
        assert!(!coordinator.is_healthy());
    }

    #[test]
    fn coordinator_stop_is_idempotent_during_partial_startup() {
        let mut runtime = module_with_id(0, "runtime");
        runtime.state = ModuleState::Running;
        runtime.health = ModuleHealth {
            ok: true,
            state: ModuleState::Running,
            message: Some("running".to_string()),
        };
        let mut driver = module_with_id(1, "driver");
        driver.dependencies = vec!["runtime".to_string()];
        driver.state = ModuleState::Starting;
        driver.health = ModuleHealth {
            ok: true,
            state: ModuleState::Starting,
            message: Some("starting".to_string()),
        };
        let graph = ModuleGraph::new(vec![runtime, driver]).unwrap();
        let mut coordinator = ModuleCoordinator::new(graph);

        coordinator.stop();
        coordinator.stop();

        assert!(coordinator
            .modules()
            .iter()
            .all(|module| module.state == ModuleState::Stopped));
    }

    #[test]
    fn failed_optional_module_does_not_leave_runtime_starting() {
        let mut optional = module_with_id(1, "optional-viewer");
        optional.required = false;
        optional.dependencies = vec!["runtime".to_string()];
        let graph = ModuleGraph::new(vec![module_with_id(0, "runtime"), optional]).unwrap();
        let mut coordinator =
            ModuleCoordinator::new(graph).with_start_failure("optional-viewer", "not available");

        let _ = coordinator.start().unwrap_err();

        assert!(coordinator
            .modules()
            .iter()
            .all(|module| { matches!(module.state, ModuleState::Stopped | ModuleState::Failed) }));
    }
}
