use crate::core::geometry::{spherical_triangle_area, SqrtExt, Tuple, VectorLike};
use crate::core::geometry::{
    Bounds3f, DirectionCone, Normal, Normal3f, Point2f, Point3f, Point3fi, Ray, Vector2f, Vector3,
    Vector3f,
};
use crate::core::interaction::{
    Interaction, InteractionBase, InteractionTrait, SimpleInteraction, SurfaceInteraction,
};
use crate::core::shape::{ShapeIntersection, ShapeSample, ShapeSampleContext, ShapeTrait};
use crate::shapes::mesh::TriangleMesh;
use crate::utils::math::{difference_of_products, square};
use crate::utils::sampling::{
    bilinear_pdf, invert_spherical_triangle_sample, sample_bilinear, sample_spherical_triangle,
    sample_uniform_triangle,
};
use crate::{gamma, Float, GVec, Ptr};

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct TriangleIntersection {
    b0: Float,
    b1: Float,
    b2: Float,
    t: Float,
}

impl TriangleIntersection {
    pub fn new(b0: Float, b1: Float, b2: Float, t: Float) -> Self {
        Self { b0, b1, b2, t }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct TriangleShape {
    pub mesh: Ptr<TriangleMesh>,
    pub tri_index: i32,
}

impl TriangleShape {
    pub const MIN_SPHERICAL_SAMPLE_AREA: Float = 3e-4;
    pub const MAX_SPHERICAL_SAMPLE_AREA: Float = 6.22;

    fn mesh(&self) -> &TriangleMesh {
        &*self.mesh
    }
    fn get_vertex_indices(&self) -> [usize; 3] {
        let mesh = self.mesh();
        let base = (self.tri_index as usize) * 3;
        [
            mesh.vertex_indices[base] as usize,
            mesh.vertex_indices[base + 1] as usize,
            mesh.vertex_indices[base + 2] as usize,
        ]
    }

    fn get_points(&self) -> [Point3f; 3] {
        let mesh = self.mesh();
        let [v0, v1, v2] = self.get_vertex_indices();
        [mesh.p[v0], mesh.p[v1], mesh.p[v2]]
    }

    fn get_shading_normals(&self) -> Option<[Normal3f; 3]> {
        let mesh = self.mesh();
        if mesh.n.is_empty() {
            return None;
        }
        let [v0, v1, v2] = self.get_vertex_indices();
        Some([mesh.n[v0], mesh.n[v1], mesh.n[v2]])
    }

    fn get_tangents(&self) -> Option<[Vector3f; 3]> {
        let mesh = self.mesh();
        if mesh.s.is_empty() {
            return None;
        }
        let [v0, v1, v2] = self.get_vertex_indices();
        Some([mesh.s[v0], mesh.s[v1], mesh.s[v2]])
    }


    fn get_uvs(&self) -> Option<[Point2f; 3]> {
        let mesh = self.mesh();
        if mesh.s.is_empty() {
            return None;
        }
        let [v0, v1, v2] = self.get_vertex_indices();
        Some([mesh.uv[v0], mesh.uv[v1], mesh.uv[v2]])
    }

    pub fn new(mesh: Ptr<TriangleMesh>, tri_index: i32) -> Self {
        Self { mesh, tri_index }
    }

    pub fn get_mesh(&self) -> Ptr<TriangleMesh> {
        self.mesh
    }

    pub fn solid_angle(&self, p: Point3f) -> Float {
        let [p0, p1, p2] = self.get_points();
        spherical_triangle_area(
            (p0 - p).normalize(),
            (p1 - p).normalize(),
            (p2 - p).normalize(),
        )
    }

    fn intersect_triangle(
        &self,
        _ray: &Ray,
        _t_max: Float,
        _p0: Point3f,
        _p1: Point3f,
        _p2: Point3f,
    ) -> Option<TriangleIntersection> {
        todo!()
    }

    fn interaction_from_intersection(
        &self,
        ti: TriangleIntersection,
        time: Float,
        wo: Vector3f,
    ) -> SurfaceInteraction {
        let [p0, p1, p2] = self.get_points();
        let uv = self.get_uvs().unwrap_or([
            Point2f::new(0.0, 0.0),
            Point2f::new(1.0, 0.0),
            Point2f::new(1.0, 1.0),
        ]);
        let duv02 = uv[0] - uv[2];
        let duv12 = uv[1] - uv[2];
        let dp02 = p0 - p2;
        let dp12 = p1 - p2;
        let determinant = difference_of_products(duv02[0], duv12[1], duv02[1], duv12[0]);
        let degenerate = determinant.abs() < 1e-9;
        let (mut dpdu, mut dpdv) = if !degenerate {
            let invdet = 1. / determinant;
            let ret0 = difference_of_products(duv12[1], dp02, duv02[1], dp12) * invdet;
            let ret1 = difference_of_products(duv02[0], dp12, duv12[0], dp02) * invdet;
            (ret0, ret1)
        } else {
            (Vector3f::zero(), Vector3f::zero())
        };

        if degenerate || dpdu.cross(dpdv).norm_squared() == 0. {
            let mut ng = (p2 - p0).cross(p1 - p0);
            if ng.norm_squared() == 0. {
                let v1 = p2 - p0;
                let v2 = p1 - p0;
                ng = v1.cast::<f64>().cross(v2.cast::<f64>()).cast::<Float>();
                assert!(ng.norm_squared() != 0.);
            }
            (dpdu, dpdv) = ng.normalize().coordinate_system();
        }

        let p0_vec = Vector3f::from(p0);
        let p1_vec = Vector3f::from(p1);
        let p2_vec = Vector3f::from(p2);
        let p_hit = Point3f::from(ti.b0 * p0_vec + ti.b1 * p1_vec + ti.b2 * p2_vec);
        let uv_hit = Point2f::from(
            ti.b0 * Vector2f::from(uv[0])
                + ti.b1 * Vector2f::from(uv[1])
                + ti.b2 * Vector2f::from(uv[2]),
        );

        let p_abs_sum = (ti.b0 * p0_vec).abs() + (ti.b1 * p1_vec).abs() + (ti.b2 * p2_vec).abs();
        let p_error = gamma(7) * p_abs_sum;
        let mut ng = Normal3f::from(dp02.cross(dp12).normalize());

        let flip_normal = self.mesh.reverse_orientation ^ self.mesh.transform_swaps_handedness;
        if flip_normal {
            ng = -ng;
        }

        let mut isect = SurfaceInteraction::new(
            Point3fi::new_with_error(p_hit, p_error),
            uv_hit,
            wo,
            dpdu,
            dpdv,
            Normal3f::zero(),
            Normal3f::zero(),
            time,
            flip_normal,
        );

        isect.face_index = if !self.mesh.face_indices.is_empty() {
            unsafe { *self.mesh.face_indices.as_ptr().add(self.tri_index as usize) }
        } else {
            0
        };

        isect.common.n = ng;
        isect.shading.n = ng;

        if !self.mesh.p.is_empty() || !self.mesh.s.is_empty() {
            self.compute_shading_geometry(&mut isect, &ti, uv, dpdu, determinant, degenerate);
        }
        isect
    }

    fn compute_shading_geometry(
        &self,
        isect: &mut SurfaceInteraction,
        ti: &TriangleIntersection,
        uv: [Point2f; 3],
        dpdu_geom: Vector3f,
        determinant: Float,
        degenerate_uv: bool,
    ) {
        // Interpolate vertex normals if they exist
        let ns = if let Some(normals) = self.get_shading_normals() {
            let n = ti.b0 * normals[0] + ti.b1 * normals[1] + ti.b2 * normals[2];
            if n.norm_squared() > 0.0 {
                n.normalize()
            } else {
                isect.n()
            }
        } else {
            isect.n()
        };

        // Interpolate tangents if they exist
        let mut ss = if let Some(tangents) = self.get_tangents() {
            let s = ti.b0 * tangents[0] + ti.b1 * tangents[1] + ti.b2 * tangents[2];
            if s.norm_squared() > 0.0 {
                s.normalize()
            } else {
                dpdu_geom
            }
        } else {
            dpdu_geom
        };

        // Ensure shading tangent (ss) is perpendicular to shading normal (ns)
        let mut ts = ns.cross(ss.into());
        if ts.norm_squared() > 0.0 {
            ss = ts.cross(ns.into()).into();
        } else {
            let (s, t) = ns.coordinate_system();
            ss = s.into();
            ts = t.into();
        }

        // How does the normal change as we move across UVs?
        let (dndu, dndv) = if let Some(normals) = self.get_shading_normals() {
            if degenerate_uv {
                let dn = (normals[2] - normals[0]).cross(normals[1] - normals[0]);
                if dn.norm_squared() == 0.0 {
                    (Normal3f::zero(), Normal3f::zero())
                } else {
                    let (dnu, dnv) = dn.coordinate_system();
                    (Normal3f::from(dnu), Normal3f::from(dnv))
                }
            } else {
                let dn1 = normals[0] - normals[2];
                let dn2 = normals[1] - normals[2];
                let duv02 = uv[0] - uv[2];
                let duv12 = uv[1] - uv[2];

                let inv_det = 1.0 / determinant;
                (
                    difference_of_products(duv12[1], dn1, duv02[1], dn2) * inv_det,
                    difference_of_products(duv02[0], dn2, duv12[0], dn1) * inv_det,
                )
            }
        } else {
            (Normal3f::zero(), Normal3f::zero())
        };

        isect.set_shading_geom(ns, ss, ts.into(), dndu, dndv, true);
    }
}

impl ShapeTrait for TriangleShape {
    fn bounds(&self) -> Bounds3f {
        let [p0, p1, p2] = self.get_points();
        Bounds3f::from_points(p0, p1).union_point(p2)
    }

    fn normal_bounds(&self) -> DirectionCone {
        let [p0, p1, p2] = self.get_points();
        let mut n: Normal3f = (p1 - p0).cross(p2 - p0).normalize().into();

        if let Some(normals) = self.get_shading_normals() {
            let [n0, n1, n2] = normals;
            let ns = n0 + n1 + n2;
            n = n.face_forward(ns);
        } else if self.mesh.reverse_orientation ^ self.mesh.transform_swaps_handedness {
            n = -n;
        }

        DirectionCone::new_from_vector(n.into())
    }

    fn area(&self) -> Float {
        let [p0, p1, p2] = self.get_points();
        0.5 * (p1 - p0).cross(p2 - p0).norm()
    }

    fn sample(&self, u: Point2f) -> Option<ShapeSample> {
        let [p0, p1, p2] = self.get_points();
        let b = sample_uniform_triangle(u);
        let p = p0 + b[1] * (p1 - p0) + b[2] * (p2 - p0);

        let mut n: Normal3f = (p1 - p0).cross(p2 - p0).normalize().into();

        if let Some(normals) = self.get_shading_normals() {
            let [n0, n1, n2] = normals;
            let ns = b[0] * n0 + b[1] * n1 + b[2] * n2; // b[2] is (1 - b0 - b1)
            n = n.face_forward(ns);
        } else if self.mesh.reverse_orientation ^ self.mesh.transform_swaps_handedness {
            n = -n;
        }

        let uv_sample = if let Some(uvs) = self.get_uvs() {
            let [uv0, uv1, uv2] = uvs;
            uv0 + b[1] * (uv1 - uv0) + b[2] * (uv2 - uv0)
        } else {
            let v = b[0] * Vector2f::new(0.0, 0.0)
                + b[1] * Vector2f::new(1.0, 0.0)
                + b[2] * Vector2f::new(1.0, 1.0);
            Point2f::from(v)
        };

        let p0_v = Vector3f::from(p0);
        let p1_v = Vector3f::from(p1);
        let p2_v = Vector3f::from(p2);
        let p_abs_sum = (b[0] * p0_v).abs() + (b[1] * p1_v).abs() + (b[2] * p2_v).abs();

        let p_error = Vector3f::from(p_abs_sum) * gamma(6);

        let intr_base = InteractionBase::new_surface_geom(
            Point3fi::new_with_error(p, p_error),
            n,
            uv_sample,
            Vector3f::default(),
            0.,
        );

        Some(ShapeSample {
            intr: Interaction::Simple(SimpleInteraction::new(intr_base)),
            pdf: 1.0 / self.area(),
        })
    }

    fn sample_from_context(&self, ctx: &ShapeSampleContext, mut u: Point2f) -> Option<ShapeSample> {
        let [p0, p1, p2] = self.get_points();

        let (b, tri_pdf) = sample_spherical_triangle(&[p0, p1, p2], ctx.p(), u)?;
        if tri_pdf == 0. {
            return None;
        }

        let solid_angle = self.solid_angle(ctx.p());
        if solid_angle < Self::MIN_SPHERICAL_SAMPLE_AREA
            || solid_angle > Self::MAX_SPHERICAL_SAMPLE_AREA
        {
            let mut ss = self.sample(u)?;
            ss.intr.get_common_mut().time = ctx.time;
            let mut wi: Normal3f = (ss.intr.p() - ctx.p()).into();
            if wi.norm_squared() == 0. {
                return None;
            }
            wi = wi.normalize();
            ss.pdf /= ss.intr.n().abs_dot(-wi) / ctx.p().distance_squared(ss.intr.p());
            if ss.pdf.is_infinite() {
                return None;
            }
            return Some(ss);
        }

        let mut pdf = 1.;
        if ctx.ns != Normal3f::zero() {
            let rp = ctx.p();
            let wi: [Vector3f; 3] = [
                (p0 - rp).normalize(),
                (p1 - rp).normalize(),
                (p2 - rp).normalize(),
            ];
            let w: [Float; 4] = [
                ctx.ns.abs_dot(wi[1].into()).max(0.01),
                ctx.ns.abs_dot(wi[1].into()).max(0.01),
                ctx.ns.abs_dot(wi[0].into()).max(0.01),
                ctx.ns.abs_dot(wi[2].into()).max(0.01),
            ];
            u = sample_bilinear(u, &w);
            pdf = bilinear_pdf(u, &w);
        }

        let p0_v = Vector3f::from(p0);
        let p1_v = Vector3f::from(p1);
        let p2_v = Vector3f::from(p2);
        let p_abs_sum = (b[0] * p0_v).abs() + (b[1] * p1_v).abs() + (b[2] * p2_v).abs();
        let mut n: Normal3f = (p1 - p0).cross(p2 - p0).normalize().into();

        if let Some(normals) = self.get_shading_normals() {
            let [n0, n1, n2] = normals;
            let ns = b[0] * n0 + b[1] * n1 + (1. - b[0] - b[1]) * n2;
            n = n.face_forward(ns);
        } else if self.mesh.reverse_orientation ^ self.mesh.transform_swaps_handedness {
            n = -n;
        }

        let uv_sample = if let Some(uvs) = self.get_uvs() {
            let [uv0, uv1, uv2] = uvs;
            uv0 + b[1] * (uv1 - uv0) + b[2] * (uv2 - uv0)
        } else {
            let v = b[0] * Vector2f::new(0.0, 0.0)
                + b[1] * Vector2f::new(1.0, 0.0)
                + b[2] * Vector2f::new(1.0, 1.0);
            Point2f::from(v)
        };

        let p = p0 + b[1] * (p1 - p0) + b[2] * (p2 - p0);
        let p_error = Vector3f::from(p_abs_sum) * gamma(6);
        let intr_base = InteractionBase::new_surface_geom(
            Point3fi::new_with_error(p, p_error),
            n,
            uv_sample,
            Vector3f::default(),
            0.,
        );

        Some(ShapeSample {
            intr: Interaction::Simple(SimpleInteraction::new(intr_base)),
            pdf,
        })
    }

    fn intersect(&self, ray: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let [p0, p1, p2] = self.get_points();
        let tri_isect = self.intersect_triangle(ray, t_max.unwrap_or(0.), p0, p1, p2)?;
        let intr = self.interaction_from_intersection(tri_isect, ray.time, -ray.d);
        Some(ShapeIntersection::new(intr, tri_isect.t))
    }

    fn intersect_p(&self, ray: &Ray, t_max: Option<Float>) -> bool {
        let [p0, p1, p2] = self.get_points();
        let tri_isect = self.intersect_triangle(ray, t_max.unwrap_or(0.), p0, p1, p2);
        tri_isect.is_some()
    }

    fn pdf(&self, _interaction: &Interaction) -> Float {
        1. / self.area()
    }

    fn pdf_from_context(&self, ctx: &ShapeSampleContext, wi: Vector3f) -> Float {
        let solid_angle = self.solid_angle(ctx.p());

        if solid_angle < Self::MIN_SPHERICAL_SAMPLE_AREA
            || solid_angle > Self::MAX_SPHERICAL_SAMPLE_AREA
        {
            let ray = ctx.spawn_ray(wi);
            let Some(isect) = self.intersect(&ray, None) else {
                return 0.;
            };

            let pdf = (1. / self.area())
                / (isect.intr.n().abs_dot(-Normal3f::from(wi))
                    / ctx.p().distance_squared(isect.intr.p()));

            if pdf.is_infinite() {
                return 0.;
            }
            return pdf;
        }

        let mut pdf = 1. / solid_angle;
        if ctx.ns != Normal3f::zero() {
            let [p0, p1, p2] = self.get_points();
            let u = invert_spherical_triangle_sample(&[p0, p1, p2], ctx.p(), wi)
                .expect("Could not calculate inverse sample");
            let rp = ctx.p();
            let wi = [
                (p0 - rp).normalize(),
                (p1 - rp).normalize(),
                (p2 - rp).normalize(),
            ];
            let w: [Float; 4] = [
                ctx.ns.abs_dot(wi[1].into()).max(0.01),
                ctx.ns.abs_dot(wi[1].into()).max(0.01),
                ctx.ns.abs_dot(wi[0].into()).max(0.01),
                ctx.ns.abs_dot(wi[2].into()).max(0.01),
            ];
            pdf *= bilinear_pdf(u, &w);
        }
        pdf
    }
}
