use crate::SdfNode;

pub fn simplify(node: &mut SdfNode) {
    match node {
        SdfNode::Union(children) | SdfNode::Intersection(children)
        | SdfNode::SmoothUnion { children, .. } => {
            for c in children.iter_mut() {
                simplify(c);
            }
        }
        SdfNode::Difference { base, subtract } => {
            simplify(base);
            for s in subtract.iter_mut() {
                simplify(s);
            }
        }
        SdfNode::Translate { child, .. } | SdfNode::Rotate { child, .. }
        | SdfNode::Scale { child, .. } => {
            simplify(child);
        }
        SdfNode::Sphere { .. } | SdfNode::Box { .. } | SdfNode::Cylinder { .. } => {}
    }

    let replacement = match node {
        SdfNode::Union(children) => {
            let mut flattened = Vec::with_capacity(children.len());
            let mut changed = false;
            for c in children.drain(..) {
                if let SdfNode::Union(inner) = c {
                    flattened.extend(inner);
                    changed = true;
                } else {
                    flattened.push(c);
                }
            }
            if flattened.len() == 1 {
                Some(flattened.into_iter().next().unwrap())
            } else if changed {
                Some(SdfNode::Union(flattened))
            } else {
                *children = flattened;
                None
            }
        }

        SdfNode::Intersection(children) => {
            let mut flattened = Vec::with_capacity(children.len());
            let mut changed = false;
            for c in children.drain(..) {
                if let SdfNode::Intersection(inner) = c {
                    flattened.extend(inner);
                    changed = true;
                } else {
                    flattened.push(c);
                }
            }
            if flattened.len() == 1 {
                Some(flattened.into_iter().next().unwrap())
            } else if changed {
                Some(SdfNode::Intersection(flattened))
            } else {
                *children = flattened;
                None
            }
        }

        SdfNode::Translate { offset, child } => {
            if offset[0] == 0.0 && offset[1] == 0.0 && offset[2] == 0.0 {
                let inner = std::mem::replace(child.as_mut(), SdfNode::Sphere { radius: 0.0 });
                Some(inner)
            } else if let SdfNode::Translate { offset: o2, child: c2 } = child.as_mut() {
                let combined = [offset[0] + o2[0], offset[1] + o2[1], offset[2] + o2[2]];
                let inner = std::mem::replace(c2.as_mut(), SdfNode::Sphere { radius: 0.0 });
                Some(SdfNode::Translate { offset: combined, child: Box::new(inner) })
            } else {
                None
            }
        }

        SdfNode::Scale { factor, child } => {
            if factor[0] == 1.0 && factor[1] == 1.0 && factor[2] == 1.0 {
                let inner = std::mem::replace(child.as_mut(), SdfNode::Sphere { radius: 0.0 });
                Some(inner)
            } else if let SdfNode::Scale { factor: f2, child: c2 } = child.as_mut() {
                let combined = [factor[0] * f2[0], factor[1] * f2[1], factor[2] * f2[2]];
                let inner = std::mem::replace(c2.as_mut(), SdfNode::Sphere { radius: 0.0 });
                Some(SdfNode::Scale { factor: combined, child: Box::new(inner) })
            } else {
                None
            }
        }

        SdfNode::Rotate { angle_deg, child, .. } => {
            if *angle_deg == 0.0 {
                let inner = std::mem::replace(child.as_mut(), SdfNode::Sphere { radius: 0.0 });
                Some(inner)
            } else {
                None
            }
        }

        _ => None,
    };

    if let Some(r) = replacement {
        *node = r;
    }
}

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

    #[test]
    fn flatten_nested_union() {
        let mut node = SdfNode::Union(vec![
            SdfNode::Union(vec![
                SdfNode::Sphere { radius: 1.0 },
                SdfNode::Sphere { radius: 2.0 },
            ]),
            SdfNode::Sphere { radius: 3.0 },
        ]);
        simplify(&mut node);
        match &node {
            SdfNode::Union(children) => assert_eq!(children.len(), 3),
            _ => panic!("expected Union"),
        }
    }

    #[test]
    fn singleton_union() {
        let mut node = SdfNode::Union(vec![SdfNode::Sphere { radius: 1.0 }]);
        simplify(&mut node);
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 1.0).abs() < 1e-10));
    }

    #[test]
    fn collapse_nested_translate() {
        let mut node = SdfNode::Translate {
            offset: [1.0, 2.0, 3.0],
            child: Box::new(SdfNode::Translate {
                offset: [4.0, 5.0, 6.0],
                child: Box::new(SdfNode::Sphere { radius: 1.0 }),
            }),
        };
        simplify(&mut node);
        match &node {
            SdfNode::Translate { offset, child } => {
                assert!((offset[0] - 5.0).abs() < 1e-10);
                assert!((offset[1] - 7.0).abs() < 1e-10);
                assert!((offset[2] - 9.0).abs() < 1e-10);
                assert!(matches!(**child, SdfNode::Sphere { .. }));
            }
            _ => panic!("expected Translate"),
        }
    }

    #[test]
    fn identity_translate_removed() {
        let mut node = SdfNode::Translate {
            offset: [0.0, 0.0, 0.0],
            child: Box::new(SdfNode::Sphere { radius: 5.0 }),
        };
        simplify(&mut node);
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 5.0).abs() < 1e-10));
    }

    #[test]
    fn identity_scale_removed() {
        let mut node = SdfNode::Scale {
            factor: [1.0, 1.0, 1.0],
            child: Box::new(SdfNode::Sphere { radius: 3.0 }),
        };
        simplify(&mut node);
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 3.0).abs() < 1e-10));
    }

    #[test]
    fn identity_rotate_removed() {
        let mut node = SdfNode::Rotate {
            axis: [1.0, 0.0, 0.0],
            angle_deg: 0.0,
            child: Box::new(SdfNode::Sphere { radius: 2.0 }),
        };
        simplify(&mut node);
        assert!(matches!(node, SdfNode::Sphere { radius } if (radius - 2.0).abs() < 1e-10));
    }

    #[test]
    fn collapse_nested_scale() {
        let mut node = SdfNode::Scale {
            factor: [2.0, 3.0, 4.0],
            child: Box::new(SdfNode::Scale {
                factor: [0.5, 0.5, 0.5],
                child: Box::new(SdfNode::Sphere { radius: 1.0 }),
            }),
        };
        simplify(&mut node);
        match &node {
            SdfNode::Scale { factor, .. } => {
                assert!((factor[0] - 1.0).abs() < 1e-10);
                assert!((factor[1] - 1.5).abs() < 1e-10);
                assert!((factor[2] - 2.0).abs() < 1e-10);
            }
            _ => panic!("expected Scale"),
        }
    }
}
