//! Standard library of shapes and transforms
//!
//! These shapes are designed for use in higher-level languages and tools.
//! Each shape is a single `struct` with named member variables.  The "output"
//! of the shape is always a [`Tree`], and can be generated by calling
//! `Tree::from(..)`.  Shape member variables can always be represented by
//! variants of [`Value`](types::Value) (checked by a unit test).
//!
//! Every shape implements `Facet + Clone + Send + Sync + Into<Tree> + 'static`.
//! When generating bindings, users are expected to inspect the shape's type
//! using annotations provided by the [`facet`] crate; to iterate over the
//! entire shape library, use [`ShapeVisitor`] and [`visit_shapes`].
//!
//! For an example of binding shapes into a dynamic language, look at the
//! implementation of `fidget_rhai::shapes` (specifically the internal
//! `register_shape` function).
#![warn(missing_docs)]
use facet::Facet;
use fidget_core::context::Tree;

pub mod types;
use types::{Axis, Plane, Vec2, Vec3};

////////////////////////////////////////////////////////////////////////////////
// 2D shapes

/// 2D circle
#[derive(Clone, Facet)]
pub struct Circle {
    /// Center of the circle (in XY)
    #[facet(default = Vec2::new(0.0, 0.0))]
    pub center: Vec2,
    /// Circle radius
    #[facet(default = 1.0)]
    pub radius: f64,
}

impl From<Circle> for Tree {
    fn from(v: Circle) -> Self {
        let (x, y, _) = Tree::axes();
        ((x - v.center.x).square() + (y - v.center.y).square()).sqrt()
            - v.radius
    }
}

/// Rectangle defined by lower and upper corners
#[derive(Clone, Facet)]
pub struct Rectangle {
    /// Lower corner of the rectangle
    pub lower: Vec2,
    /// Upper corner of the rectangle
    pub upper: Vec2,
}

impl From<Rectangle> for Tree {
    fn from(v: Rectangle) -> Self {
        let (x, y, _) = Tree::axes();
        (v.lower.x - x.clone())
            .max(x - v.upper.x)
            .max((v.lower.y - y.clone()).max(y - v.upper.y))
    }
}

////////////////////////////////////////////////////////////////////////////////
// 3D shapes

/// 3D sphere
#[derive(Clone, Facet)]
pub struct Sphere {
    /// Center of the circle (in XYZ)
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub center: Vec3,
    /// Sphere radius
    #[facet(default = 1.0)]
    pub radius: f64,
}

impl From<Sphere> for Tree {
    fn from(v: Sphere) -> Self {
        let (x, y, z) = Tree::axes();
        ((x - v.center.x).square()
            + (y - v.center.y).square()
            + (z - v.center.z).square())
        .sqrt()
            - v.radius
    }
}

/// Box defined by lower and upper corners
#[derive(Clone, Facet)]
pub struct Box {
    /// Lower corner of the rectangle
    pub lower: Vec3,
    /// Upper corner of the rectangle
    pub upper: Vec3,
}

impl From<Box> for Tree {
    fn from(v: Box) -> Self {
        let (x, y, z) = Tree::axes();
        (v.lower.x - x.clone())
            .max(x - v.upper.x)
            .max((v.lower.y - y.clone()).max(y - v.upper.y))
            .max((v.lower.z - z.clone()).max(z - v.upper.z))
    }
}

////////////////////////////////////////////////////////////////////////////////
// CSG operations

/// Take the union of a set of shapes
///
/// If the input is empty, returns an constant empty tree (at +∞)
#[derive(Clone, Facet)]
pub struct Union {
    /// List of shapes to merge
    pub input: Vec<Tree>,
}

impl From<Union> for Tree {
    fn from(v: Union) -> Self {
        if v.input.is_empty() {
            // XXX should this be an error instead?
            Tree::constant(f64::INFINITY)
        } else {
            fn recurse(s: &[Tree]) -> Tree {
                match s.len() {
                    1 => s[0].clone(),
                    n => recurse(&s[..n / 2]).min(recurse(&s[n / 2..])),
                }
            }
            recurse(&v.input)
        }
    }
}

/// Smooth quadratic blend of two shapes
///
/// This formula is taken from "Lipschitz Pruning: Hierarchical Simplification
/// of Primitive-Based SDFs" (Barbier _et al_ '25), which in turn cites
/// [Quilez '20](https://iquilezles.org/articles/smin/)
#[derive(Clone, Facet)]
pub struct Blend {
    /// First shape input
    pub a: Tree,
    /// Second shape input
    pub b: Tree,
    /// Blending radius
    pub radius: f64,
}

impl From<Blend> for Tree {
    fn from(v: Blend) -> Self {
        if v.radius > 0.0 {
            v.a.clone().min(v.b.clone())
                - 1.0 / (4.0 * v.radius)
                    * (v.radius - (v.a - v.b).abs()).max(0.0).square()
        } else {
            v.a.min(v.b)
        }
    }
}

/// Take the intersection of a set of shapes
///
/// If the input is empty, returns a constant full tree (at -∞)
#[derive(Clone, Facet)]
pub struct Intersection {
    /// List of shapes to intersect
    pub input: Vec<Tree>,
}

impl From<Intersection> for Tree {
    fn from(v: Intersection) -> Self {
        if v.input.is_empty() {
            // XXX should this be an error instead?
            Tree::constant(-f64::INFINITY)
        } else {
            fn recurse(s: &[Tree]) -> Tree {
                match s.len() {
                    1 => s[0].clone(),
                    n => recurse(&s[..n / 2]).max(recurse(&s[n / 2..])),
                }
            }
            recurse(&v.input)
        }
    }
}

/// Computes the inverse of a shape
#[derive(Clone, Facet)]
pub struct Inverse {
    /// Shape to invert
    pub shape: Tree,
}

impl From<Inverse> for Tree {
    fn from(v: Inverse) -> Self {
        -v.shape
    }
}

/// Take the difference of two shapes
#[derive(Clone, Facet)]
pub struct Difference {
    /// Original shape
    pub shape: Tree,
    /// Shape to be subtracted from the original
    pub cutout: Tree,
}

impl From<Difference> for Tree {
    fn from(v: Difference) -> Self {
        v.shape.max(-v.cutout)
    }
}

////////////////////////////////////////////////////////////////////////////////
// Transforms

/// Move a shape
#[derive(Clone, Facet)]
pub struct Move {
    /// Shape to move
    pub shape: Tree,
    /// Position offset
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub offset: Vec3,
}

impl From<Move> for Tree {
    fn from(v: Move) -> Self {
        v.shape.remap_affine(nalgebra::convert(
            nalgebra::Translation3::<f64>::new(
                -v.offset.x,
                -v.offset.y,
                -v.offset.z,
            ),
        ))
    }
}

/// Non-uniform scaling
#[derive(Clone, Facet)]
pub struct Scale {
    /// Shape to scale
    pub shape: Tree,
    /// Scale to apply on each axis
    #[facet(default = Vec3::new(1.0, 1.0, 1.0))]
    pub scale: Vec3,
}

impl From<Scale> for Tree {
    fn from(v: Scale) -> Self {
        v.shape
            .remap_affine(nalgebra::convert(nalgebra::Scale3::<f64>::new(
                1.0 / v.scale.x,
                1.0 / v.scale.y,
                1.0 / v.scale.z,
            )))
    }
}

/// Uniform scaling
#[derive(Clone, Facet)]
pub struct ScaleUniform {
    /// Shape to scale
    pub shape: Tree,
    /// Scale to apply
    #[facet(default = 1.0)]
    pub scale: f64,
}

impl From<ScaleUniform> for Tree {
    fn from(v: ScaleUniform) -> Self {
        let s = 1.0 / v.scale;
        v.shape
            .remap_affine(nalgebra::convert(nalgebra::Scale3::<f64>::new(
                s, s, s,
            )))
    }
}

/// Reflection
#[derive(Clone, Facet)]
pub struct Reflect {
    /// Shape to reflect
    pub shape: Tree,

    /// Plane about which to reflect the shape
    #[facet(default = Plane::YZ)]
    pub plane: Plane,
}

impl From<Reflect> for Tree {
    fn from(v: Reflect) -> Self {
        let a = v.plane.axis.vec();
        let (x, y, z) = Tree::axes();
        let d = a.x * x.clone() + a.y * y.clone() + a.z * z.clone()
            - v.plane.offset;
        let scale: Tree = 2.0 * d;
        // TODO could we use nalgebra::Reflection3 here to make the transform
        // affine?  For some reason, it doesn't implement the right SubSet
        // https://github.com/dimforge/nalgebra/issues/1527
        v.shape.remap_xyz(
            x - scale.clone() * a.x,
            y - scale.clone() * a.y,
            z - scale * a.z,
        )
    }
}

/// Reflection about the X axis
#[derive(Clone, Facet)]
pub struct ReflectX {
    /// Shape to reflect
    pub shape: Tree,

    /// Plane about which to reflect the shape
    #[facet(default = 0.0)]
    pub offset: f64,
}

impl From<ReflectX> for Tree {
    fn from(v: ReflectX) -> Self {
        Reflect {
            shape: v.shape,
            plane: Plane {
                axis: Axis::X,
                offset: v.offset,
            },
        }
        .into()
    }
}

/// Reflection about the Y axis
#[derive(Clone, Facet)]
pub struct ReflectY {
    /// Shape to reflect
    pub shape: Tree,

    /// Plane about which to reflect the shape
    #[facet(default = 0.0)]
    pub offset: f64,
}

impl From<ReflectY> for Tree {
    fn from(v: ReflectY) -> Self {
        Reflect {
            shape: v.shape,
            plane: Plane {
                axis: Axis::Y,
                offset: v.offset,
            },
        }
        .into()
    }
}

/// Reflection about the Z axis
#[derive(Clone, Facet)]
pub struct ReflectZ {
    /// Shape to reflect
    pub shape: Tree,

    /// Plane about which to reflect the shape
    #[facet(default = 0.0)]
    pub offset: f64,
}

impl From<ReflectZ> for Tree {
    fn from(v: ReflectZ) -> Self {
        Reflect {
            shape: v.shape,
            plane: Plane {
                axis: Axis::Z,
                offset: v.offset,
            },
        }
        .into()
    }
}

/// Rotates an object about an arbitrary axis and rotation center
#[derive(Clone, Facet)]
pub struct Rotate {
    /// Shape to rotate
    pub shape: Tree,

    /// Axis about which to rotate
    #[facet(default = Axis::Z)]
    pub axis: Axis,

    /// Angle to rotate (in degrees)
    #[facet(default = 0.0)]
    pub angle: f64,

    /// Center of rotation
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub center: Vec3,
}

impl From<Rotate> for Tree {
    fn from(v: Rotate) -> Self {
        let shape = Tree::from(Move {
            shape: v.shape,
            offset: -v.center,
        });
        let d = -v.angle.to_radians();
        let axis = v.axis.vec();
        let shape = shape.remap_affine(nalgebra::convert(
            nalgebra::Rotation3::<f64>::new(nalgebra::Vector3::from(d * *axis)),
        ));
        Move {
            shape,
            offset: v.center,
        }
        .into()
    }
}

/// Rotates an object about the X axis
#[derive(Clone, Facet)]
pub struct RotateX {
    /// Shape to rotate
    pub shape: Tree,

    /// Angle to rotate (in degrees)
    #[facet(default = 0.0)]
    pub angle: f64,

    /// Center of rotation
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub center: Vec3,
}

impl From<RotateX> for Tree {
    fn from(v: RotateX) -> Self {
        Rotate {
            shape: v.shape,
            angle: v.angle,
            center: v.center,
            axis: Axis::X,
        }
        .into()
    }
}

/// Rotates an object about the Y axis
#[derive(Clone, Facet)]
pub struct RotateY {
    /// Shape to rotate
    pub shape: Tree,

    /// Angle to rotate (in degrees)
    #[facet(default = 0.0)]
    pub angle: f64,

    /// Center of rotation
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub center: Vec3,
}

impl From<RotateY> for Tree {
    fn from(v: RotateY) -> Self {
        Rotate {
            shape: v.shape,
            angle: v.angle,
            center: v.center,
            axis: Axis::Y,
        }
        .into()
    }
}

/// Rotates an object about the Z axis
#[derive(Clone, Facet)]
pub struct RotateZ {
    /// Shape to rotate
    pub shape: Tree,

    /// Angle to rotate (in degrees)
    #[facet(default = 0.0)]
    pub angle: f64,

    /// Center of rotation
    #[facet(default = Vec3::new(0.0, 0.0, 0.0))]
    pub center: Vec3,
}

impl From<RotateZ> for Tree {
    fn from(v: RotateZ) -> Self {
        Rotate {
            shape: v.shape,
            angle: v.angle,
            center: v.center,
            axis: Axis::Z,
        }
        .into()
    }
}

// TODO figure out a generic Revolve?  The matrix math is a bit tricky!

/// Revolve a shape about the Y axis, creating a 3D volume
#[derive(Clone, Facet)]
pub struct RevolveY {
    /// Shape to revolve
    pub shape: Tree,
    /// X offset about which to revolve
    #[facet(default = 0.0)]
    pub offset: f64,
}

impl From<RevolveY> for Tree {
    fn from(v: RevolveY) -> Self {
        let offset = Vec3::new(-v.offset, 0.0, 0.0);
        let shape = Tree::from(Move {
            shape: v.shape.clone(),
            offset: -offset,
        });
        let (x, y, z) = Tree::axes();
        let r = (x.square() + y.square()).sqrt();
        let shape = shape.remap_xyz(r, y, z);
        Move { shape, offset }.into()
    }
}

/// Extrude an XY shape in the Z direction
#[derive(Clone, Facet)]
pub struct ExtrudeZ {
    /// Shape to extrude
    pub shape: Tree,
    /// Lower bounds of the extrusion
    #[facet(default = 0.0)]
    pub lower: f64,
    /// Upper bounds of the extrusion
    #[facet(default = 1.0)]
    pub upper: f64,
}

impl From<ExtrudeZ> for Tree {
    fn from(v: ExtrudeZ) -> Self {
        let (x, y, z) = Tree::axes();
        let t = v.shape.remap_xyz(x, y, Tree::constant(0.0));
        t.max((v.lower - z.clone()).max(z - v.upper))
    }
}

/// Loft between two XY shape in the Z direction
#[derive(Clone, Facet)]
pub struct LoftZ {
    /// Lower shape
    pub a: Tree,
    /// Upper shape
    pub b: Tree,
    /// Lower bounds of the loft
    #[facet(default = 0.0)]
    pub lower: f64,
    /// Upper bounds of the loft
    #[facet(default = 1.0)]
    pub upper: f64,
}

impl From<LoftZ> for Tree {
    fn from(v: LoftZ) -> Self {
        let (x, y, z) = Tree::axes();
        let ta = v.a.remap_xyz(x.clone(), y.clone(), Tree::constant(0.0));
        let tb = v.b.remap_xyz(x, y, Tree::constant(0.0));
        let t = ((z.clone() - v.lower) * tb + (v.upper - z.clone()) * ta)
            / (v.upper - v.lower);
        t.max((v.lower - z.clone()).max(z - v.upper))
    }
}

////////////////////////////////////////////////////////////////////////////////

/// Trait for a type which can visit each of the shapes in our library
pub trait ShapeVisitor {
    /// Process the given type
    fn visit<T: Facet<'static> + Clone + Send + Sync + Into<Tree> + 'static>(
        &mut self,
    );
}

/// Maps a shape visitor across all shape definitions
pub fn visit_shapes<V: ShapeVisitor>(visitor: &mut V) {
    visitor.visit::<Sphere>();
    visitor.visit::<Box>();
    visitor.visit::<Plane>();

    visitor.visit::<Circle>();
    visitor.visit::<Rectangle>();

    visitor.visit::<Move>();
    visitor.visit::<Scale>();
    visitor.visit::<ScaleUniform>();
    visitor.visit::<Reflect>();
    visitor.visit::<ReflectX>();
    visitor.visit::<ReflectY>();
    visitor.visit::<ReflectZ>();
    visitor.visit::<Rotate>();
    visitor.visit::<RotateX>();
    visitor.visit::<RotateY>();
    visitor.visit::<RotateZ>();
    visitor.visit::<RevolveY>();
    visitor.visit::<ExtrudeZ>();
    visitor.visit::<LoftZ>();

    visitor.visit::<Union>();
    visitor.visit::<Blend>();
    visitor.visit::<Intersection>();
    visitor.visit::<Difference>();
    visitor.visit::<Inverse>();
}

////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod test {
    use super::*;
    use crate::types::eval_default_fn;
    use fidget_core::Context;

    #[test]
    fn circle_docstring() {
        assert_eq!(Circle::SHAPE.doc, &[" 2D circle"]);
    }

    #[test]
    fn transform_order() {
        let x = Tree::x();
        let moved: Tree = Move {
            shape: x,
            offset: Vec3::new(-1.0, 0.0, 0.0),
        }
        .into();
        let mut ctx = Context::new();
        let cm = ctx.import(&moved);
        assert_eq!(ctx.eval_xyz(cm, 0.0, 0.0, 0.0).unwrap(), 1.0);
        assert_eq!(ctx.eval_xyz(cm, 0.0, 1.0, 0.0).unwrap(), 1.0);
        assert_eq!(ctx.eval_xyz(cm, -1.0, 0.0, 0.0).unwrap(), 0.0);

        let rotated: Tree = RotateZ {
            shape: moved,
            angle: 90.0,
            center: Vec3::new(0.0, 0.0, 0.0),
        }
        .into();
        let cr = ctx.import(&rotated);
        assert_eq!(ctx.eval_xyz(cr, 0.0, 0.0, 0.0).unwrap(), 1.0);
        assert_eq!(ctx.eval_xyz(cr, 0.0, -1.0, 0.0).unwrap(), 0.0);
        assert_eq!(ctx.eval_xyz(cr, 0.0, 1.0, 0.0).unwrap(), 2.0);
    }

    #[test]
    fn scale_default_fn() {
        let facet::Type::User(facet::UserType::Struct(s)) = Scale::SHAPE.ty
        else {
            panic!();
        };
        for f in s.fields {
            if f.name == "scale" {
                let Some(facet::DefaultSource::Custom(f)) = f.default else {
                    panic!()
                };
                let v: Vec3 = unsafe { eval_default_fn(f) };
                assert_eq!(v.x, 1.0);
                assert_eq!(v.y, 1.0);
                assert_eq!(v.z, 1.0);
            } else {
                assert!(f.default.is_none());
            }
        }
    }

    struct ValidateVisitor;
    impl ShapeVisitor for ValidateVisitor {
        fn visit<
            T: Facet<'static> + Clone + Send + Sync + Into<Tree> + 'static,
        >(
            &mut self,
        ) {
            let facet::Type::User(facet::UserType::Struct(s)) = T::SHAPE.ty
            else {
                panic!("shape `{}` must be a struct", T::SHAPE.type_name());
            };
            for f in s.fields {
                if types::Type::try_from(f.shape().id).is_err() {
                    panic!(
                        "field `{}` in `{}` has unhandled type: {}",
                        f.name,
                        T::SHAPE.type_name(),
                        f.shape()
                    );
                }
                if let Some(d) = f.default {
                    assert!(
                        matches!(d, facet::DefaultSource::Custom(..)),
                        "default on field `{}` in `{}` must include value",
                        f.name,
                        T::SHAPE.type_name()
                    );
                }
            }
        }
    }

    #[test]
    fn validate_shapes() {
        let mut v = ValidateVisitor;
        visit_shapes(&mut v);
    }

    #[test]
    #[should_panic(
        expected = "field `uhoh` in `BadShape` has unhandled type: String"
    )]
    fn bad_shape_type() {
        #[derive(Clone, facet::Facet)]
        struct BadShape {
            uhoh: String,
        }
        impl From<BadShape> for Tree {
            fn from(_: BadShape) -> Tree {
                unimplemented!()
            }
        }
        let mut v = ValidateVisitor;
        v.visit::<BadShape>();
    }

    #[test]
    #[should_panic(
        expected = "default on field `center` in `BadShape` must include value"
    )]
    fn bad_shape_default() {
        #[derive(Clone, facet::Facet)]
        struct BadShape {
            #[facet(default)]
            center: f64,
        }
        impl From<BadShape> for Tree {
            fn from(_: BadShape) -> Tree {
                unimplemented!()
            }
        }
        let mut v = ValidateVisitor;
        v.visit::<BadShape>();
    }

    #[test]
    #[should_panic(expected = "shape `BadShapeEnum` must be a struct")]
    fn bad_shape_shape() {
        #[derive(Clone, facet::Facet)]
        #[repr(C)]
        enum BadShapeEnum {
            UhOh,
        }
        impl From<BadShapeEnum> for Tree {
            fn from(_: BadShapeEnum) -> Tree {
                unimplemented!()
            }
        }
        let mut v = ValidateVisitor;
        v.visit::<BadShapeEnum>();
    }
}
