use nalgebra_glm::{Mat4, Vec3};

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SdfPrimitive {
    Sphere {
        radius: f32,
    },
    Box {
        half_extents: Vec3,
    },
    Cylinder {
        radius: f32,
        half_height: f32,
    },
    Torus {
        major_radius: f32,
        minor_radius: f32,
    },
    Capsule {
        radius: f32,
        half_height: f32,
    },
    Plane {
        normal: Vec3,
        offset: f32,
    },
}

impl SdfPrimitive {
    pub fn name(&self) -> &'static str {
        match self {
            Self::Sphere { .. } => "Sphere",
            Self::Box { .. } => "Box",
            Self::Cylinder { .. } => "Cylinder",
            Self::Torus { .. } => "Torus",
            Self::Capsule { .. } => "Capsule",
            Self::Plane { .. } => "Plane",
        }
    }

    pub fn evaluate(&self, point: Vec3) -> f32 {
        match *self {
            SdfPrimitive::Sphere { radius } => nalgebra_glm::length(&point) - radius,

            SdfPrimitive::Box { half_extents } => {
                let q = Vec3::new(
                    point.x.abs() - half_extents.x,
                    point.y.abs() - half_extents.y,
                    point.z.abs() - half_extents.z,
                );
                let outside =
                    nalgebra_glm::length(&Vec3::new(q.x.max(0.0), q.y.max(0.0), q.z.max(0.0)));
                let inside = q.x.max(q.y).max(q.z).min(0.0);
                outside + inside
            }

            SdfPrimitive::Cylinder {
                radius,
                half_height,
            } => {
                let d_xz = (point.x * point.x + point.z * point.z).sqrt() - radius;
                let d_y = point.y.abs() - half_height;
                let outside = Vec3::new(d_xz.max(0.0), d_y.max(0.0), 0.0);
                let inside = d_xz.max(d_y).min(0.0);
                nalgebra_glm::length(&outside.xy()) + inside
            }

            SdfPrimitive::Torus {
                major_radius,
                minor_radius,
            } => {
                let q_xz = (point.x * point.x + point.z * point.z).sqrt() - major_radius;
                let q = Vec3::new(q_xz, point.y, 0.0);
                nalgebra_glm::length(&q.xy()) - minor_radius
            }

            SdfPrimitive::Capsule {
                radius,
                half_height,
            } => {
                let clamped_y = point.y.clamp(-half_height, half_height);
                let closest = Vec3::new(0.0, clamped_y, 0.0);
                nalgebra_glm::length(&(point - closest)) - radius
            }

            SdfPrimitive::Plane { normal, offset } => nalgebra_glm::dot(&point, &normal) - offset,
        }
    }

    pub fn bounding_radius(&self) -> f32 {
        match *self {
            SdfPrimitive::Sphere { radius } => radius,
            SdfPrimitive::Box { half_extents } => nalgebra_glm::length(&half_extents),
            SdfPrimitive::Cylinder {
                radius,
                half_height,
            } => (radius * radius + half_height * half_height).sqrt(),
            SdfPrimitive::Torus {
                major_radius,
                minor_radius,
            } => major_radius + minor_radius,
            SdfPrimitive::Capsule {
                radius,
                half_height,
            } => half_height + radius,
            SdfPrimitive::Plane { .. } => f32::INFINITY,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Default)]
pub enum CsgOperation {
    #[default]
    Union,
    Subtraction,
    Intersection,
    SmoothUnion {
        smoothness: f32,
    },
    SmoothSubtraction {
        smoothness: f32,
    },
    SmoothIntersection {
        smoothness: f32,
    },
}

impl CsgOperation {
    pub fn smoothness(&self) -> f32 {
        match *self {
            CsgOperation::SmoothUnion { smoothness } => smoothness,
            CsgOperation::SmoothSubtraction { smoothness } => smoothness,
            CsgOperation::SmoothIntersection { smoothness } => smoothness,
            _ => 0.0,
        }
    }

    pub fn apply(&self, distance_a: f32, distance_b: f32) -> f32 {
        match *self {
            CsgOperation::Union => distance_a.min(distance_b),
            CsgOperation::Subtraction => distance_a.max(-distance_b),
            CsgOperation::Intersection => distance_a.max(distance_b),
            CsgOperation::SmoothUnion { smoothness } => {
                smooth_min(distance_a, distance_b, smoothness)
            }
            CsgOperation::SmoothSubtraction { smoothness } => {
                -smooth_min(-distance_a, distance_b, smoothness)
            }
            CsgOperation::SmoothIntersection { smoothness } => {
                -smooth_min(-distance_a, -distance_b, smoothness)
            }
        }
    }

    pub fn apply_with_material(
        &self,
        distance_a: f32,
        material_a: u32,
        distance_b: f32,
        material_b: u32,
    ) -> (f32, u32) {
        self.apply_with_material_dithered(distance_a, material_a, distance_b, material_b, 0.5)
    }

    pub fn apply_with_material_dithered(
        &self,
        distance_a: f32,
        material_a: u32,
        distance_b: f32,
        material_b: u32,
        dither_threshold: f32,
    ) -> (f32, u32) {
        match *self {
            CsgOperation::Union => {
                if distance_a < distance_b {
                    (distance_a, material_a)
                } else {
                    (distance_b, material_b)
                }
            }
            CsgOperation::Subtraction => {
                if distance_a > -distance_b {
                    (distance_a, material_a)
                } else {
                    (-distance_b, material_a)
                }
            }
            CsgOperation::Intersection => {
                if distance_a > distance_b {
                    (distance_a, material_a)
                } else {
                    (distance_b, material_b)
                }
            }
            CsgOperation::SmoothUnion { smoothness } => {
                let h = (0.5 + 0.5 * (distance_b - distance_a) / smoothness).clamp(0.0, 1.0);
                let distance = lerp(distance_b, distance_a, h) - smoothness * h * (1.0 - h);
                let material = if h > dither_threshold {
                    material_a
                } else {
                    material_b
                };
                (distance, material)
            }
            CsgOperation::SmoothSubtraction { smoothness } => {
                let h = (0.5 - 0.5 * (distance_a + distance_b) / smoothness).clamp(0.0, 1.0);
                let distance = lerp(distance_a, -distance_b, h) + smoothness * h * (1.0 - h);
                (distance, material_a)
            }
            CsgOperation::SmoothIntersection { smoothness } => {
                let h = (0.5 - 0.5 * (distance_b - distance_a) / smoothness).clamp(0.0, 1.0);
                let distance = lerp(distance_b, distance_a, h) + smoothness * h * (1.0 - h);
                let material = if h > dither_threshold {
                    material_a
                } else {
                    material_b
                };
                (distance, material)
            }
        }
    }
}

fn smooth_min(a: f32, b: f32, k: f32) -> f32 {
    if k <= 0.0 {
        return a.min(b);
    }
    let h = (0.5 + 0.5 * (b - a) / k).clamp(0.0, 1.0);
    lerp(b, a, h) - k * h * (1.0 - h)
}

fn lerp(a: f32, b: f32, t: f32) -> f32 {
    a + (b - a) * t
}

#[derive(Debug, Clone)]
pub enum SdfEdit {
    Union {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
    },
    Subtraction {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
    },
    Intersection {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
    },
    SmoothUnion {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
        smoothness: f32,
    },
    SmoothSubtraction {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
        smoothness: f32,
    },
    SmoothIntersection {
        primitive: SdfPrimitive,
        transform: Mat4,
        inverse_transform: Mat4,
        uniform_scale: f32,
        material_id: u32,
        smoothness: f32,
    },
}

impl SdfEdit {
    pub fn union(primitive: SdfPrimitive, transform: Mat4, material_id: u32) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::Union {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
        }
    }

    pub fn subtraction(primitive: SdfPrimitive, transform: Mat4, material_id: u32) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::Subtraction {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
        }
    }

    pub fn intersection(primitive: SdfPrimitive, transform: Mat4, material_id: u32) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::Intersection {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
        }
    }

    pub fn smooth_union(
        primitive: SdfPrimitive,
        transform: Mat4,
        material_id: u32,
        smoothness: f32,
    ) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::SmoothUnion {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
            smoothness,
        }
    }

    pub fn smooth_subtraction(
        primitive: SdfPrimitive,
        transform: Mat4,
        material_id: u32,
        smoothness: f32,
    ) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::SmoothSubtraction {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
            smoothness,
        }
    }

    pub fn smooth_intersection(
        primitive: SdfPrimitive,
        transform: Mat4,
        material_id: u32,
        smoothness: f32,
    ) -> Self {
        let (inverse_transform, uniform_scale) = compute_transform_data(&transform);
        Self::SmoothIntersection {
            primitive,
            transform,
            inverse_transform,
            uniform_scale,
            material_id,
            smoothness,
        }
    }

    pub fn from_operation(
        primitive: SdfPrimitive,
        operation: CsgOperation,
        transform: Mat4,
        material_id: u32,
    ) -> Self {
        match operation {
            CsgOperation::Union => Self::union(primitive, transform, material_id),
            CsgOperation::Subtraction => Self::subtraction(primitive, transform, material_id),
            CsgOperation::Intersection => Self::intersection(primitive, transform, material_id),
            CsgOperation::SmoothUnion { smoothness } => {
                Self::smooth_union(primitive, transform, material_id, smoothness)
            }
            CsgOperation::SmoothSubtraction { smoothness } => {
                Self::smooth_subtraction(primitive, transform, material_id, smoothness)
            }
            CsgOperation::SmoothIntersection { smoothness } => {
                Self::smooth_intersection(primitive, transform, material_id, smoothness)
            }
        }
    }

    pub fn with_identity(primitive: SdfPrimitive, material_id: u32) -> Self {
        Self::union(primitive, Mat4::identity(), material_id)
    }

    pub fn primitive(&self) -> &SdfPrimitive {
        match self {
            Self::Union { primitive, .. }
            | Self::Subtraction { primitive, .. }
            | Self::Intersection { primitive, .. }
            | Self::SmoothUnion { primitive, .. }
            | Self::SmoothSubtraction { primitive, .. }
            | Self::SmoothIntersection { primitive, .. } => primitive,
        }
    }

    pub fn transform(&self) -> &Mat4 {
        match self {
            Self::Union { transform, .. }
            | Self::Subtraction { transform, .. }
            | Self::Intersection { transform, .. }
            | Self::SmoothUnion { transform, .. }
            | Self::SmoothSubtraction { transform, .. }
            | Self::SmoothIntersection { transform, .. } => transform,
        }
    }

    pub fn inverse_transform(&self) -> &Mat4 {
        match self {
            Self::Union {
                inverse_transform, ..
            }
            | Self::Subtraction {
                inverse_transform, ..
            }
            | Self::Intersection {
                inverse_transform, ..
            }
            | Self::SmoothUnion {
                inverse_transform, ..
            }
            | Self::SmoothSubtraction {
                inverse_transform, ..
            }
            | Self::SmoothIntersection {
                inverse_transform, ..
            } => inverse_transform,
        }
    }

    pub fn uniform_scale(&self) -> f32 {
        match self {
            Self::Union { uniform_scale, .. }
            | Self::Subtraction { uniform_scale, .. }
            | Self::Intersection { uniform_scale, .. }
            | Self::SmoothUnion { uniform_scale, .. }
            | Self::SmoothSubtraction { uniform_scale, .. }
            | Self::SmoothIntersection { uniform_scale, .. } => *uniform_scale,
        }
    }

    pub fn material_id(&self) -> u32 {
        match self {
            Self::Union { material_id, .. }
            | Self::Subtraction { material_id, .. }
            | Self::Intersection { material_id, .. }
            | Self::SmoothUnion { material_id, .. }
            | Self::SmoothSubtraction { material_id, .. }
            | Self::SmoothIntersection { material_id, .. } => *material_id,
        }
    }

    pub fn smoothness(&self) -> f32 {
        match self {
            Self::SmoothUnion { smoothness, .. }
            | Self::SmoothSubtraction { smoothness, .. }
            | Self::SmoothIntersection { smoothness, .. } => *smoothness,
            _ => 0.0,
        }
    }

    pub fn operation(&self) -> CsgOperation {
        match self {
            Self::Union { .. } => CsgOperation::Union,
            Self::Subtraction { .. } => CsgOperation::Subtraction,
            Self::Intersection { .. } => CsgOperation::Intersection,
            Self::SmoothUnion { smoothness, .. } => CsgOperation::SmoothUnion {
                smoothness: *smoothness,
            },
            Self::SmoothSubtraction { smoothness, .. } => CsgOperation::SmoothSubtraction {
                smoothness: *smoothness,
            },
            Self::SmoothIntersection { smoothness, .. } => CsgOperation::SmoothIntersection {
                smoothness: *smoothness,
            },
        }
    }

    pub fn operation_name(&self) -> &'static str {
        match self {
            Self::Union { .. } => "Union",
            Self::Subtraction { .. } => "Subtraction",
            Self::Intersection { .. } => "Intersection",
            Self::SmoothUnion { .. } => "Smooth Union",
            Self::SmoothSubtraction { .. } => "Smooth Subtraction",
            Self::SmoothIntersection { .. } => "Smooth Intersection",
        }
    }

    pub fn position(&self) -> Vec3 {
        let transform = self.transform();
        Vec3::new(transform[(0, 3)], transform[(1, 3)], transform[(2, 3)])
    }

    pub fn primitive_name(&self) -> &'static str {
        self.primitive().name()
    }

    pub fn set_transform(&mut self, new_transform: Mat4) {
        let (new_inverse, new_scale) = compute_transform_data(&new_transform);
        match self {
            Self::Union {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            }
            | Self::Subtraction {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            }
            | Self::Intersection {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            }
            | Self::SmoothUnion {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            }
            | Self::SmoothSubtraction {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            }
            | Self::SmoothIntersection {
                transform,
                inverse_transform,
                uniform_scale,
                ..
            } => {
                *transform = new_transform;
                *inverse_transform = new_inverse;
                *uniform_scale = new_scale;
            }
        }
    }

    pub fn evaluate(&self, world_point: Vec3) -> f32 {
        let local_point = (self.inverse_transform()
            * nalgebra_glm::vec4(world_point.x, world_point.y, world_point.z, 1.0))
        .xyz();

        self.primitive().evaluate(local_point) * self.uniform_scale()
    }

    pub fn world_bounds(&self) -> Aabb {
        let primitive = self.primitive();
        let transform = self.transform();

        match primitive {
            SdfPrimitive::Plane { normal, offset } => {
                let center = (transform * nalgebra_glm::vec4(0.0, 0.0, 0.0, 1.0)).xyz();
                let world_normal =
                    (transform * nalgebra_glm::vec4(normal.x, normal.y, normal.z, 0.0)).xyz();
                let world_normal = nalgebra_glm::normalize(&world_normal);

                let plane_point = center + world_normal * *offset;

                const SLAB_THICKNESS: f32 = 2.0;
                const PLANE_EXTENT: f32 = 20.0;

                let abs_normal = Vec3::new(
                    world_normal.x.abs(),
                    world_normal.y.abs(),
                    world_normal.z.abs(),
                );

                let extent = Vec3::new(
                    if abs_normal.x > 0.9 {
                        SLAB_THICKNESS
                    } else {
                        PLANE_EXTENT
                    },
                    if abs_normal.y > 0.9 {
                        SLAB_THICKNESS
                    } else {
                        PLANE_EXTENT
                    },
                    if abs_normal.z > 0.9 {
                        SLAB_THICKNESS
                    } else {
                        PLANE_EXTENT
                    },
                );

                Aabb {
                    min: plane_point - extent,
                    max: plane_point + extent,
                }
            }
            _ => {
                let local_radius = primitive.bounding_radius();
                let scale = extract_scale(transform);
                let max_scale = scale.x.max(scale.y).max(scale.z);
                let world_radius = local_radius * max_scale;

                let center = (transform * nalgebra_glm::vec4(0.0, 0.0, 0.0, 1.0)).xyz();

                Aabb {
                    min: center - Vec3::new(world_radius, world_radius, world_radius),
                    max: center + Vec3::new(world_radius, world_radius, world_radius),
                }
            }
        }
    }

    pub fn effective_bounds(&self) -> Aabb {
        let bounds = self.world_bounds();
        let smoothness = self.smoothness();
        if smoothness > 0.0 {
            bounds.expand(smoothness)
        } else {
            bounds
        }
    }
}

fn compute_transform_data(transform: &Mat4) -> (Mat4, f32) {
    let inverse_transform = nalgebra_glm::inverse(transform);
    let scale = extract_scale(transform);
    let uniform_scale = (scale.x * scale.y * scale.z).powf(1.0 / 3.0);
    (inverse_transform, uniform_scale)
}

fn extract_scale(matrix: &Mat4) -> Vec3 {
    Vec3::new(
        nalgebra_glm::length(&Vec3::new(matrix[(0, 0)], matrix[(1, 0)], matrix[(2, 0)])),
        nalgebra_glm::length(&Vec3::new(matrix[(0, 1)], matrix[(1, 1)], matrix[(2, 1)])),
        nalgebra_glm::length(&Vec3::new(matrix[(0, 2)], matrix[(1, 2)], matrix[(2, 2)])),
    )
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Aabb {
    pub min: Vec3,
    pub max: Vec3,
}

impl Aabb {
    pub fn new(min: Vec3, max: Vec3) -> Self {
        Self { min, max }
    }

    pub fn from_center_half_extents(center: Vec3, half_extents: Vec3) -> Self {
        Self {
            min: center - half_extents,
            max: center + half_extents,
        }
    }

    pub fn union(&self, other: &Aabb) -> Aabb {
        Aabb {
            min: Vec3::new(
                self.min.x.min(other.min.x),
                self.min.y.min(other.min.y),
                self.min.z.min(other.min.z),
            ),
            max: Vec3::new(
                self.max.x.max(other.max.x),
                self.max.y.max(other.max.y),
                self.max.z.max(other.max.z),
            ),
        }
    }

    pub fn intersects(&self, other: &Aabb) -> bool {
        self.min.x <= other.max.x
            && self.max.x >= other.min.x
            && self.min.y <= other.max.y
            && self.max.y >= other.min.y
            && self.min.z <= other.max.z
            && self.max.z >= other.min.z
    }

    pub fn contains_point(&self, point: Vec3) -> bool {
        point.x >= self.min.x
            && point.x <= self.max.x
            && point.y >= self.min.y
            && point.y <= self.max.y
            && point.z >= self.min.z
            && point.z <= self.max.z
    }

    pub fn expand(&self, amount: f32) -> Aabb {
        Aabb {
            min: self.min - Vec3::new(amount, amount, amount),
            max: self.max + Vec3::new(amount, amount, amount),
        }
    }

    pub fn center(&self) -> Vec3 {
        (self.min + self.max) * 0.5
    }

    pub fn half_extents(&self) -> Vec3 {
        (self.max - self.min) * 0.5
    }
}

impl Default for Aabb {
    fn default() -> Self {
        Self {
            min: Vec3::new(f32::INFINITY, f32::INFINITY, f32::INFINITY),
            max: Vec3::new(f32::NEG_INFINITY, f32::NEG_INFINITY, f32::NEG_INFINITY),
        }
    }
}
