use std::collections::{BTreeMap, BTreeSet};

use crate::model::component::v1::CapabilityRef;
use crate::model::component::v1::capability::Capability;
use anyhow::{Result, bail};
use serde::{Deserialize, Serialize};

use super::{Robot, Role};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoleAssignment {
    pub capability: CapabilityRef,
    pub roles: BTreeSet<Role>,
    pub source: RoleAssignmentSource,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoleAssignmentSource {
    Explicit,
    Inferred,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoleResolution {
    pub assignments: Vec<RoleAssignment>,
}

impl RoleResolution {
    #[must_use]
    pub fn capabilities_for(&self, role: Role) -> Vec<&CapabilityRef> {
        self.assignments
            .iter()
            .filter(|assignment| assignment.roles.contains(&role))
            .map(|assignment| &assignment.capability)
            .collect()
    }
}

pub fn resolve_roles(
    model: &Robot,
    components_by_type: &BTreeMap<String, crate::model::component::v1::Component>,
) -> Result<RoleResolution> {
    let mut errors = Vec::new();
    let mut capability_index = BTreeMap::new();

    for (component_id, component_instance) in &model.components {
        match components_by_type.get(&component_instance.component) {
            Some(component) => {
                for (capability_id, capability) in &component.capabilities {
                    capability_index.insert(
                        CapabilityRef::new(component_id, capability_id),
                        capability.clone(),
                    );
                }
            }
            None => errors.push(format!(
                "component '{}' uses unstaged component type '{}'",
                component_id, component_instance.component
            )),
        }
    }

    let mut role_to_explicit_capabilities: BTreeMap<Role, Vec<CapabilityRef>> = BTreeMap::new();
    let mut assignments: BTreeMap<CapabilityRef, RoleAssignment> = BTreeMap::new();

    for (component_id, component_instance) in &model.components {
        for (capability_id, roles) in &component_instance.roles {
            let capability_ref = CapabilityRef::new(component_id, capability_id);
            match capability_index.get(&capability_ref) {
                Some(capability) => {
                    let mut seen = BTreeSet::new();
                    for role in roles {
                        if !seen.insert(*role) {
                            errors.push(format!(
                                "components.{component_id}.roles.{capability_id} repeats role '{role}'"
                            ));
                        }
                        if !role_matches_capability(*role, capability) {
                            errors.push(format!(
                                "components.{component_id}.roles.{capability_id} assigns role '{role}' to {} capability",
                                capability.kind_name()
                            ));
                        }
                        role_to_explicit_capabilities
                            .entry(*role)
                            .or_default()
                            .push(capability_ref.clone());
                    }
                    assignments.insert(
                        capability_ref.clone(),
                        RoleAssignment {
                            capability: capability_ref,
                            roles: seen,
                            source: RoleAssignmentSource::Explicit,
                        },
                    );
                }
                None => errors.push(format!(
                    "components.{component_id}.roles.{capability_id} references missing capability"
                )),
            }
        }
    }

    for (role, capabilities) in &role_to_explicit_capabilities {
        if !role.allows_multiple_capabilities() && capabilities.len() > 1 {
            errors.push(format!(
                "role '{role}' is explicitly assigned to multiple capabilities: {}",
                format_capabilities(capabilities)
            ));
        }
    }

    for role in [
        Role::Localization,
        Role::Mapping,
        Role::Traversability,
        Role::Safety,
        Role::Odometry,
    ] {
        if role_to_explicit_capabilities.contains_key(&role) {
            continue;
        }

        let candidates = capability_index
            .iter()
            .filter(|(_, capability)| role_matches_capability(role, capability))
            .map(|(capability_ref, _)| capability_ref.clone())
            .collect::<Vec<_>>();

        match candidates.as_slice() {
            [capability] => {
                assignments
                    .entry(capability.clone())
                    .and_modify(|assignment| {
                        assignment.roles.insert(role);
                    })
                    .or_insert_with(|| RoleAssignment {
                        capability: capability.clone(),
                        roles: [role].into_iter().collect(),
                        source: RoleAssignmentSource::Inferred,
                    });
            }
            [] => {}
            _ if role.allows_multiple_capabilities() => {}
            _ => errors.push(format!(
                "role '{role}' is ambiguous without an explicit model hint; candidates: {}",
                format_capabilities(&candidates)
            )),
        }
    }

    if errors.is_empty() {
        Ok(RoleResolution {
            assignments: assignments.into_values().collect(),
        })
    } else {
        bail!("Role resolution errors:\n{}", errors.join("\n"))
    }
}

fn role_matches_capability(role: Role, capability: &Capability) -> bool {
    match role {
        Role::Localization => matches!(
            capability,
            Capability::Depth(_)
                | Capability::Lidar(_)
                | Capability::Camera(_)
                | Capability::Gnss(_)
                | Capability::Imu(_)
        ),
        Role::Mapping | Role::Traversability | Role::Safety => matches!(
            capability,
            Capability::Range(_) | Capability::Depth(_) | Capability::Lidar(_)
        ),
        Role::Odometry => matches!(capability, Capability::Imu(_)),
        Role::Perception => matches!(capability, Capability::Camera(_) | Capability::Depth(_)),
    }
}

fn format_capabilities(capabilities: &[CapabilityRef]) -> String {
    capabilities
        .iter()
        .map(ToString::to_string)
        .collect::<Vec<_>>()
        .join(", ")
}

#[cfg(test)]
mod tests {
    use crate::model::component::v1::capability::{
        Accelerometer, Camera, CameraMode, Capability, Depth, Gnss, Imu, Lidar, LidarOutput, Range,
        StructuralTarget,
    };

    use super::{Role, role_matches_capability};

    #[test]
    fn localization_role_matches_visual_inertial_backend_inputs() {
        assert!(role_matches_capability(
            Role::Localization,
            &Capability::Depth(Depth {
                target: link_target(),
                publish_rate_hz: 30.0,
                width_px: 640,
                height_px: 480,
                field_of_view_rad: None,
                min_range_m: None,
                max_range_m: None,
            })
        ));
        assert!(role_matches_capability(
            Role::Localization,
            &Capability::Lidar(Lidar {
                target: link_target(),
                publish_rate_hz: 10.0,
                output: LidarOutput::Points,
                min_range_m: None,
                max_range_m: None,
                horizontal_fov_rad: None,
                horizontal_resolution_rad: None,
                vertical_fov_rad: None,
                vertical_resolution_rad: None,
            })
        ));
        assert!(role_matches_capability(
            Role::Localization,
            &Capability::Camera(Camera {
                target: link_target(),
                mode: CameraMode::Rgb,
                publish_rate_hz: 30.0,
                width_px: 640,
                height_px: 480,
                field_of_view_rad: None,
            })
        ));
        assert!(role_matches_capability(
            Role::Localization,
            &Capability::Imu(Imu {
                target: link_target(),
                publish_rate_hz: 100.0,
                axes: None,
            })
        ));
        assert!(!role_matches_capability(
            Role::Localization,
            &Capability::Range(Range {
                target: link_target(),
                publish_rate_hz: 20.0,
                min_range_m: 0.03,
                max_range_m: 4.0,
                field_of_view_rad: 0.47,
            })
        ));
        assert!(!role_matches_capability(
            Role::Localization,
            &Capability::Accelerometer(Accelerometer {
                target: link_target(),
                publish_rate_hz: 100.0,
                axes: None,
            })
        ));
    }

    #[test]
    fn perception_role_matches_camera_and_depth_only() {
        assert!(role_matches_capability(
            Role::Perception,
            &Capability::Camera(Camera {
                target: link_target(),
                mode: crate::model::component::v1::capability::CameraMode::Rgb,
                publish_rate_hz: 30.0,
                width_px: 640,
                height_px: 480,
                field_of_view_rad: None,
            })
        ));
        assert!(role_matches_capability(
            Role::Perception,
            &Capability::Depth(Depth {
                target: link_target(),
                publish_rate_hz: 30.0,
                width_px: 640,
                height_px: 480,
                field_of_view_rad: None,
                min_range_m: None,
                max_range_m: None,
            })
        ));
        assert!(!role_matches_capability(
            Role::Perception,
            &Capability::Range(Range {
                target: link_target(),
                publish_rate_hz: 10.0,
                min_range_m: 0.05,
                max_range_m: 4.0,
                field_of_view_rad: 0.1,
            })
        ));
    }

    #[test]
    fn localization_role_matches_gnss() {
        assert!(role_matches_capability(
            Role::Localization,
            &Capability::Gnss(Gnss {
                target: link_target(),
                publish_rate_hz: 10.0,
                coordinate_system: Default::default(),
            })
        ));
    }

    fn link_target() -> StructuralTarget {
        StructuralTarget::Link {
            id: "sensor_link".to_string(),
        }
    }
}
