use bone_kernel::ExtrudeFeature;
use bone_types::{ExtrudeId, FeatureId, SketchId};
use serde::{Deserialize, Serialize};
use slotmap::{Key, KeyData};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum PrincipalPlane {
    Xy,
    Yz,
    Zx,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FeatureNode {
    Origin,
    PrincipalPlane(PrincipalPlane),
    Sketch(SketchId),
    Extrude(ExtrudeId),
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum FeatureEdge {
    SketchToExtrude {
        sketch: FeatureId,
        extrude: FeatureId,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct FeatureEntry {
    id: FeatureId,
    node: FeatureNode,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct FeatureTree {
    entries: Vec<FeatureEntry>,
    #[serde(skip)]
    edges: Vec<FeatureEdge>,
}

impl FeatureTree {
    #[must_use]
    pub fn seeded() -> Self {
        let seeds = [
            FeatureNode::Origin,
            FeatureNode::PrincipalPlane(PrincipalPlane::Xy),
            FeatureNode::PrincipalPlane(PrincipalPlane::Yz),
            FeatureNode::PrincipalPlane(PrincipalPlane::Zx),
        ];
        let entries = seeds
            .into_iter()
            .zip(1u32..)
            .map(|(node, idx)| FeatureEntry {
                id: feature_id_from_idx(idx),
                node,
            })
            .collect();
        Self {
            entries,
            edges: Vec::new(),
        }
    }

    #[must_use]
    pub fn node(&self, id: FeatureId) -> Option<FeatureNode> {
        self.entries.iter().find(|e| e.id == id).map(|e| e.node)
    }

    pub fn iter(&self) -> impl Iterator<Item = (FeatureId, FeatureNode)> + '_ {
        self.entries.iter().map(|e| (e.id, e.node))
    }

    #[must_use]
    pub fn feature_of_sketch(&self, sketch: SketchId) -> Option<FeatureId> {
        self.feature_of(|node| matches!(node, FeatureNode::Sketch(s) if *s == sketch))
    }

    pub fn push_sketch(&mut self, sketch: SketchId) -> FeatureId {
        if let Some(existing) = self.feature_of_sketch(sketch) {
            return existing;
        }
        let id = self.allocate();
        self.entries.push(FeatureEntry {
            id,
            node: FeatureNode::Sketch(sketch),
        });
        id
    }

    pub fn remove_sketch(&mut self, sketch: SketchId) -> Option<FeatureId> {
        let id = self.feature_of_sketch(sketch)?;
        self.entries.retain(|e| e.id != id);
        self.drop_edges_incident(id);
        Some(id)
    }

    #[must_use]
    pub fn feature_of_extrude(&self, extrude: ExtrudeId) -> Option<FeatureId> {
        self.feature_of(|node| matches!(node, FeatureNode::Extrude(x) if *x == extrude))
    }

    pub fn push_extrude(&mut self, extrude: ExtrudeId, feature: &ExtrudeFeature) -> FeatureId {
        let id = self.feature_of_extrude(extrude).unwrap_or_else(|| {
            let id = self.allocate();
            self.entries.push(FeatureEntry {
                id,
                node: FeatureNode::Extrude(extrude),
            });
            id
        });
        let edge = self.sketch_edge(id, feature);
        self.drop_edges_incident(id);
        self.edges.extend(edge);
        self.edges.sort_unstable();
        id
    }

    pub fn remove_extrude(&mut self, extrude: ExtrudeId) -> Option<FeatureId> {
        let id = self.feature_of_extrude(extrude)?;
        self.entries.retain(|e| e.id != id);
        self.drop_edges_incident(id);
        Some(id)
    }

    #[must_use]
    pub fn edges(&self) -> &[FeatureEdge] {
        &self.edges
    }

    pub(crate) fn rebuild_edges<'a>(
        &mut self,
        extrudes: impl Iterator<Item = (ExtrudeId, &'a ExtrudeFeature)>,
    ) {
        let mut rebuilt: Vec<FeatureEdge> = extrudes
            .filter_map(|(extrude, feature)| {
                let extrude = self.feature_of_extrude(extrude)?;
                self.sketch_edge(extrude, feature)
            })
            .collect();
        rebuilt.sort_unstable();
        self.edges = rebuilt;
    }

    fn sketch_edge(&self, extrude: FeatureId, feature: &ExtrudeFeature) -> Option<FeatureEdge> {
        let sketch = self.feature_of_sketch(feature.sketch)?;
        Some(FeatureEdge::SketchToExtrude { sketch, extrude })
    }

    fn drop_edges_incident(&mut self, feature: FeatureId) {
        self.edges.retain(|edge| match edge {
            FeatureEdge::SketchToExtrude { sketch, extrude } => {
                *sketch != feature && *extrude != feature
            }
        });
    }

    fn feature_of(&self, matches_node: impl Fn(&FeatureNode) -> bool) -> Option<FeatureId> {
        self.entries
            .iter()
            .find(|e| matches_node(&e.node))
            .map(|e| e.id)
    }

    fn allocate(&self) -> FeatureId {
        let highest = self.entries.iter().map(|e| idx_of(e.id)).max().unwrap_or(0);
        let Some(next) = highest.checked_add(1) else {
            panic!("FeatureTree exhausted 32-bit feature id space");
        };
        feature_id_from_idx(next)
    }
}

fn feature_id_from_idx(idx: u32) -> FeatureId {
    FeatureId::from(KeyData::from_ffi((1u64 << 32) | u64::from(idx)))
}

fn idx_of(id: FeatureId) -> u32 {
    let Ok(idx) = u32::try_from(id.data().as_ffi() & 0xFFFF_FFFF) else {
        panic!("lower 32 bits of ffi key fit in u32");
    };
    idx
}

#[cfg(test)]
pub(crate) fn sample_blind_extrude(sketch: SketchId) -> ExtrudeFeature {
    use bone_kernel::{ExtrudeDirection, ExtrudeEndCondition, ExtrudeSense, MergeResult};
    use bone_types::{Length, PositiveLength, millimeter};
    let Ok(depth) = PositiveLength::new(Length::new::<millimeter>(10.0)) else {
        panic!("10 mm is a positive length");
    };
    ExtrudeFeature {
        sketch,
        direction: ExtrudeDirection::Normal {
            sense: ExtrudeSense::Forward,
        },
        end_condition: ExtrudeEndCondition::Blind { depth },
        draft: None,
        thin_wall: None,
        merge_result: MergeResult::Merge,
    }
}

#[cfg(test)]
mod tests {
    use super::{FeatureEdge, FeatureNode, FeatureTree, sample_blind_extrude};
    use bone_types::{ExtrudeId, SketchId};

    #[test]
    fn push_extrude_adds_node_and_links_sketch() {
        let mut tree = FeatureTree::seeded();
        let sketch = SketchId::default();
        let sketch_feature = tree.push_sketch(sketch);
        let extrude = ExtrudeId::default();
        let extrude_feature = tree.push_extrude(extrude, &sample_blind_extrude(sketch));

        assert_eq!(
            tree.node(extrude_feature),
            Some(FeatureNode::Extrude(extrude))
        );
        assert_eq!(tree.feature_of_extrude(extrude), Some(extrude_feature));
        assert_eq!(
            tree.edges(),
            &[FeatureEdge::SketchToExtrude {
                sketch: sketch_feature,
                extrude: extrude_feature,
            }]
        );
    }

    #[test]
    fn push_extrude_is_idempotent_on_duplicate_id() {
        let mut tree = FeatureTree::seeded();
        let sketch = SketchId::default();
        tree.push_sketch(sketch);
        let extrude = ExtrudeId::default();
        let feature = sample_blind_extrude(sketch);
        let first = tree.push_extrude(extrude, &feature);
        let count_before = tree.iter().count();
        let second = tree.push_extrude(extrude, &feature);

        assert_eq!(first, second);
        assert_eq!(tree.iter().count(), count_before);
        assert_eq!(tree.edges().len(), 1);
    }

    #[test]
    fn push_extrude_without_sketch_node_records_no_edge() {
        let mut tree = FeatureTree::seeded();
        let extrude = ExtrudeId::default();
        let extrude_feature =
            tree.push_extrude(extrude, &sample_blind_extrude(SketchId::default()));

        assert_eq!(
            tree.node(extrude_feature),
            Some(FeatureNode::Extrude(extrude))
        );
        assert!(tree.edges().is_empty());
    }

    fn sketch_key(n: u64) -> SketchId {
        use slotmap::KeyData;
        SketchId::from(KeyData::from_ffi((1u64 << 32) | n))
    }

    #[test]
    fn push_extrude_retargets_edge_to_new_sketch() {
        let mut tree = FeatureTree::seeded();
        let sketch_a = sketch_key(1);
        let sketch_b = sketch_key(2);
        tree.push_sketch(sketch_a);
        let feat_b = tree.push_sketch(sketch_b);
        let extrude = ExtrudeId::default();
        let ext_feat = tree.push_extrude(extrude, &sample_blind_extrude(sketch_a));
        tree.push_extrude(extrude, &sample_blind_extrude(sketch_b));
        assert_eq!(
            tree.edges(),
            &[FeatureEdge::SketchToExtrude {
                sketch: feat_b,
                extrude: ext_feat,
            }]
        );
    }

    #[test]
    fn remove_extrude_drops_node_and_edge() {
        let mut tree = FeatureTree::seeded();
        let sketch = SketchId::default();
        tree.push_sketch(sketch);
        let extrude = ExtrudeId::default();
        tree.push_extrude(extrude, &sample_blind_extrude(sketch));
        assert_eq!(tree.edges().len(), 1);

        let removed = tree.remove_extrude(extrude);

        assert!(removed.is_some());
        assert_eq!(tree.feature_of_extrude(extrude), None);
        assert!(tree.edges().is_empty());
    }

    #[test]
    fn edges_are_derived_not_persisted() {
        let mut tree = FeatureTree::seeded();
        let sketch = SketchId::default();
        tree.push_sketch(sketch);
        tree.push_extrude(ExtrudeId::default(), &sample_blind_extrude(sketch));
        assert_eq!(tree.edges().len(), 1);
        let Ok(ron) = crate::io::ron_io::to_string(&tree) else {
            panic!("feature tree serializes");
        };
        assert!(!ron.contains("edges"));
        let Ok(back) = crate::io::ron_io::from_str::<FeatureTree>(&ron) else {
            panic!("feature tree deserializes");
        };
        assert!(back.edges().is_empty());
    }

    fn extrude_key(n: u64) -> ExtrudeId {
        use slotmap::KeyData;
        ExtrudeId::from(KeyData::from_ffi((1u64 << 32) | n))
    }

    #[test]
    fn edges_canonical_regardless_of_push_order() {
        let sketch = SketchId::default();
        let high = extrude_key(5);
        let low = extrude_key(2);

        let mut pushed = FeatureTree::seeded();
        pushed.push_sketch(sketch);
        pushed.push_extrude(high, &sample_blind_extrude(sketch));
        pushed.push_extrude(low, &sample_blind_extrude(sketch));

        let mut rebuilt = pushed.clone();
        let extrudes = std::collections::BTreeMap::from([
            (low, sample_blind_extrude(sketch)),
            (high, sample_blind_extrude(sketch)),
        ]);
        rebuilt.rebuild_edges(extrudes.iter().map(|(id, feature)| (*id, feature)));

        assert_eq!(
            pushed.edges(),
            rebuilt.edges(),
            "live edge order must match the load rebuild irrespective of push order"
        );
        assert_eq!(pushed.edges().len(), 2);
    }
}
