use rand::Rng;

use crate::bvh::{Aabb, Bvh};
use crate::vec3::{reflect, refract, schlick, Vec3};

pub const EPSILON: f64 = 1e-6;

/// Cosine-weighted direction in the hemisphere around the unit vector `n`.
/// Uses the branchless orthonormal basis of Duff et al. 2017 ("Building an
/// Orthonormal Basis, Revisited"); the result is already unit length.
pub fn cosine_hemisphere<R: Rng>(n: Vec3, rng: &mut R) -> Vec3 {
    let r1: f64 = rng.gen();
    let r2: f64 = rng.gen();
    let phi = 2.0 * std::f64::consts::PI * r1;
    let r = r2.sqrt();
    let (x, y, z) = (r * phi.cos(), r * phi.sin(), (1.0 - r2).max(0.0).sqrt());

    let sign = 1.0_f64.copysign(n.z);
    let a = -1.0 / (sign + n.z);
    let b = n.x * n.y * a;
    let t = Vec3::new(1.0 + sign * n.x * n.x * a, sign * b, -sign * n.x);
    let bt = Vec3::new(b, sign + n.y * n.y * a, -n.y);
    t * x + bt * y + n * z
}

#[derive(Clone, Copy, Debug)]
pub struct Ray {
    pub origin: Vec3,
    pub direction: Vec3,
}

impl Ray {
    pub fn new(origin: Vec3, direction: Vec3) -> Self {
        Ray {
            origin,
            direction: direction.normalize(),
        }
    }

    /// For directions already known to be unit length.
    pub fn from_unit(origin: Vec3, direction: Vec3) -> Self {
        Ray { origin, direction }
    }

    pub fn at(&self, t: f64) -> Vec3 {
        self.origin + t * self.direction
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Material {
    Diffuse { albedo: Vec3 },
    Mirror { albedo: Vec3 },
    Dielectric { ref_idx: f64, albedo: Vec3 },
    Light { emission: Vec3 },
}

pub struct HitRecord {
    pub point: Vec3,
    pub normal: Vec3,
    pub front_face: bool,
    pub material: Material,
}

impl HitRecord {
    /// Orients the normal against the incoming ray, as in the Python version.
    fn new(r: &Ray, t: f64, outward: Vec3, material: Material) -> Self {
        let front_face = r.direction.dot(outward) < 0.0;
        HitRecord {
            point: r.at(t),
            normal: if front_face { outward } else { -outward },
            front_face,
            material,
        }
    }
}

impl Material {
    /// Returns (attenuation, scattered ray), or None if the ray is absorbed.
    pub fn scatter<R: Rng>(&self, ray_in: &Ray, rec: &HitRecord, rng: &mut R) -> Option<(Vec3, Ray)> {
        match *self {
            Material::Diffuse { albedo } => {
                // Cosine-weighted sampling: the Lambertian BRDF * cos / pdf
                // reduces to just the albedo.
                let direction = cosine_hemisphere(rec.normal, rng);
                Some((albedo, Ray::from_unit(rec.point, direction)))
            }
            Material::Mirror { albedo } => {
                let reflected = reflect(ray_in.direction, rec.normal);
                if reflected.dot(rec.normal) <= 0.0 {
                    return None;
                }
                Some((albedo, Ray::new(rec.point, reflected)))
            }
            Material::Dielectric { ref_idx, albedo } => {
                let ratio = if rec.front_face { 1.0 / ref_idx } else { ref_idx };
                let d = ray_in.direction;
                let cos_theta = (-d.dot(rec.normal)).min(1.0);
                let sin_theta = (1.0 - cos_theta * cos_theta).max(0.0).sqrt();
                let cannot_refract = ratio * sin_theta > 1.0;
                let direction = if cannot_refract || schlick(cos_theta, ratio) > rng.gen::<f64>() {
                    reflect(d, rec.normal)
                } else {
                    refract(d, rec.normal, ratio).unwrap_or_else(|| reflect(d, rec.normal))
                };
                Some((albedo, Ray::new(rec.point, direction)))
            }
            Material::Light { .. } => None,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Shape {
    Sphere { center: Vec3, radius: f64 },
    Plane { point: Vec3, normal: Vec3 },
    Box { min: Vec3, max: Vec3 },
    /// Axis-aligned rectangle spanning min..max, flat along the axis of
    /// `normal` (which must be ±x, ±y or ±z). As a light it emits only on the
    /// `normal` side.
    Rect { min: Vec3, max: Vec3, normal: Vec3 },
}

impl Shape {
    fn hit(&self, r: &Ray, t_min: f64, t_max: f64) -> Option<(f64, Vec3)> {
        match *self {
            Shape::Sphere { center, radius } => {
                let oc = r.origin - center;
                let a = r.direction.dot(r.direction);
                let b = oc.dot(r.direction);
                let c = oc.dot(oc) - radius * radius;
                let disc = b * b - a * c;
                if disc < 0.0 {
                    return None;
                }
                let sq = disc.sqrt();
                let mut root = (-b - sq) / a;
                if root < t_min || root > t_max {
                    root = (-b + sq) / a;
                    if root < t_min || root > t_max {
                        return None;
                    }
                }
                let outward = (r.at(root) - center).normalize();
                Some((root, outward))
            }
            Shape::Plane { point, normal } => {
                let denom = r.direction.dot(normal);
                if denom.abs() < 1e-12 {
                    return None;
                }
                let t = (point - r.origin).dot(normal) / denom;
                if t < t_min || t > t_max {
                    return None;
                }
                Some((t, normal))
            }
            Shape::Box { min, max } => {
                let mut t_lo = t_min;
                let mut t_hi = t_max;
                // Python indexes with -1 when no slab raises t_lo, which means the z axis.
                let mut hit_axis = 2;
                let mut inv_ds = [0.0; 3];
                for axis in 0..3 {
                    let d = r.direction[axis];
                    let inv_d = if d.abs() > 1e-12 { 1.0 / d } else { 1e12 };
                    inv_ds[axis] = inv_d;
                    let mut t0 = (min[axis] - r.origin[axis]) * inv_d;
                    let mut t1 = (max[axis] - r.origin[axis]) * inv_d;
                    if inv_d < 0.0 {
                        std::mem::swap(&mut t0, &mut t1);
                    }
                    if t0 > t_lo {
                        t_lo = t0;
                        hit_axis = axis;
                    }
                    t_hi = t_hi.min(t1);
                    if t_hi < t_lo {
                        return None;
                    }
                }
                if t_lo >= t_max || t_lo <= t_min {
                    return None;
                }
                // A ray travelling +axis enters through the min face, whose outward normal is -axis.
                let sign = if inv_ds[hit_axis] > 0.0 { -1.0 } else { 1.0 };
                let mut n = [0.0; 3];
                n[hit_axis] = sign;
                Some((t_lo, Vec3::new(n[0], n[1], n[2])))
            }
            Shape::Rect { min, max, normal } => {
                let denom = r.direction.dot(normal);
                if denom.abs() < 1e-12 {
                    return None;
                }
                let t = (min - r.origin).dot(normal) / denom;
                if t < t_min || t > t_max {
                    return None;
                }
                let p = r.at(t);
                let tol = 1e-9;
                let inside = (0..3).all(|a| p[a] >= min[a] - tol && p[a] <= max[a] + tol);
                inside.then_some((t, normal))
            }
        }
    }
}

impl Shape {
    /// None for unbounded shapes (planes), which are tested outside the BVH.
    fn bounds(&self) -> Option<Aabb> {
        match *self {
            Shape::Sphere { center, radius } => Some(Aabb {
                min: center - Vec3::splat(radius),
                max: center + Vec3::splat(radius),
            }),
            Shape::Box { min, max } => Some(Aabb { min, max }),
            Shape::Rect { min, max, .. } => Some(Aabb {
                min: min - Vec3::splat(1e-9),
                max: max + Vec3::splat(1e-9),
            }),
            Shape::Plane { .. } => None,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Object {
    pub shape: Shape,
    pub material: Material,
}

/// An emitter that next event estimation can sample: a box, sphere or rectangle.
#[derive(Clone, Copy, Debug)]
pub struct Light {
    shape: Shape,
    pub emission: Vec3,
    /// Total emitted power (up to a constant): luminance * area.
    power: f64,
    center: Vec3,
    /// Half the bounding box diagonal: how close a point can be to the centre
    /// while still outside the light.
    extent: f64,
}

impl Light {
    fn new(shape: Shape, emission: Vec3) -> Self {
        let bounds = shape.bounds().expect("planes cannot be lights");
        let mut light = Light {
            shape,
            emission,
            power: 0.0,
            center: bounds.centroid(),
            extent: 0.5 * (bounds.max - bounds.min).length(),
        };
        light.power = emission.luminance() * light.area();
        light
    }

    pub fn area(&self) -> f64 {
        match self.shape {
            Shape::Box { min, max } => {
                let d = max - min;
                2.0 * (d.x * d.y + d.y * d.z + d.x * d.z)
            }
            Shape::Sphere { radius, .. } => 4.0 * std::f64::consts::PI * radius * radius,
            Shape::Rect { min, max, normal } => {
                let d = max - min;
                if normal.x != 0.0 {
                    d.y * d.z
                } else if normal.y != 0.0 {
                    d.x * d.z
                } else {
                    d.x * d.y
                }
            }
            Shape::Plane { .. } => unreachable!("planes cannot be lights"),
        }
    }

    /// Rough estimate of how much this light contributes at `point`: its
    /// power over the squared distance, clamped so points near or inside the
    /// light's bounds don't blow up.
    fn importance(&self, point: Vec3) -> f64 {
        let d = point - self.center;
        self.power / d.dot(d).max(self.extent * self.extent)
    }

    /// A uniformly distributed point on the surface (pdf 1 / area) and the
    /// outward normal there.
    pub fn sample<R: Rng>(&self, rng: &mut R) -> (Vec3, Vec3) {
        match self.shape {
            Shape::Box { min, max } => {
                // Pick a face with probability proportional to its area, then
                // a uniform point on it.
                let size = max - min;
                let face_areas = [size.y * size.z, size.x * size.z, size.x * size.y];
                let u = rng.gen::<f64>() * (face_areas[0] + face_areas[1] + face_areas[2]);
                let axis = if u < face_areas[0] {
                    0
                } else if u < face_areas[0] + face_areas[1] {
                    1
                } else {
                    2
                };
                let positive_side = rng.gen::<f64>() < 0.5;
                let mut p = [
                    min.x + size.x * rng.gen::<f64>(),
                    min.y + size.y * rng.gen::<f64>(),
                    min.z + size.z * rng.gen::<f64>(),
                ];
                p[axis] = if positive_side { max[axis] } else { min[axis] };
                let mut n = [0.0; 3];
                n[axis] = if positive_side { 1.0 } else { -1.0 };
                (Vec3::new(p[0], p[1], p[2]), Vec3::new(n[0], n[1], n[2]))
            }
            Shape::Sphere { center, radius } => {
                let z = 1.0 - 2.0 * rng.gen::<f64>();
                let r = (1.0 - z * z).max(0.0).sqrt();
                let phi = 2.0 * std::f64::consts::PI * rng.gen::<f64>();
                let n = Vec3::new(r * phi.cos(), r * phi.sin(), z);
                (center + radius * n, n)
            }
            Shape::Rect { min, max, normal } => {
                let size = max - min;
                let p = Vec3::new(
                    min.x + size.x * rng.gen::<f64>(),
                    min.y + size.y * rng.gen::<f64>(),
                    min.z + size.z * rng.gen::<f64>(),
                );
                (p, normal)
            }
            Shape::Plane { .. } => unreachable!("planes cannot be lights"),
        }
    }
}

/// What a ray that leaves the scene sees.
#[derive(Clone, Copy, Debug)]
pub enum Background {
    /// White-to-blue gradient, as in raytracer.py.
    Sky,
    Color(Vec3),
}

impl Background {
    pub fn radiance(&self, direction: Vec3) -> Vec3 {
        match *self {
            Background::Sky => {
                let t = 0.5 * (direction.y + 1.0);
                Vec3::ONE * (1.0 - t) + Vec3::new(0.5, 0.7, 1.0) * t
            }
            Background::Color(c) => c,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct View {
    pub lookfrom: Vec3,
    pub lookat: Vec3,
    pub vfov: f64,
}

pub struct SceneBuilder {
    objects: Vec<Object>,
    planes: Vec<Object>,
    lights: Vec<Light>,
    background: Background,
    view: View,
}

impl SceneBuilder {
    pub fn new(view: View, background: Background) -> Self {
        SceneBuilder {
            objects: Vec::new(),
            planes: Vec::new(),
            lights: Vec::new(),
            background,
            view,
        }
    }

    pub fn add(&mut self, shape: Shape, material: Material) {
        let obj = Object { shape, material };
        match shape {
            Shape::Plane { .. } => self.planes.push(obj),
            _ => self.objects.push(obj),
        }
    }

    pub fn plane(&mut self, point: Vec3, normal: Vec3, material: Material) {
        self.add(Shape::Plane { point, normal: normal.normalize() }, material);
    }

    pub fn sphere(&mut self, center: Vec3, radius: f64, material: Material) {
        self.add(Shape::Sphere { center, radius }, material);
    }

    pub fn cuboid(&mut self, min: Vec3, max: Vec3, material: Material) {
        self.add(Shape::Box { min, max }, material);
    }

    /// Every emitter must be added through a *_light method: after a diffuse
    /// bounce, trace() only counts light found by sampling these.
    pub fn box_light(&mut self, min: Vec3, max: Vec3, emission: Vec3) {
        self.add_light(Shape::Box { min, max }, emission);
    }

    pub fn sphere_light(&mut self, center: Vec3, radius: f64, emission: Vec3) {
        self.add_light(Shape::Sphere { center, radius }, emission);
    }

    /// One-sided rectangular light emitting towards `normal` (±x, ±y or ±z).
    pub fn rect_light(&mut self, min: Vec3, max: Vec3, normal: Vec3, emission: Vec3) {
        self.add_light(Shape::Rect { min, max, normal }, emission);
    }

    fn add_light(&mut self, shape: Shape, emission: Vec3) {
        self.lights.push(Light::new(shape, emission));
        self.add(shape, Material::Light { emission });
    }

    pub fn build(self) -> Scene {
        let bounds: Vec<Aabb> = self.objects.iter().map(|o| o.shape.bounds().unwrap()).collect();
        let spheres: Vec<Option<(Vec3, f64)>> = self
            .objects
            .iter()
            .map(|o| match o.shape {
                Shape::Sphere { center, radius } => Some((center, radius)),
                _ => None,
            })
            .collect();
        Scene {
            bvh: Bvh::build(&bounds, &spheres),
            objects: self.objects,
            planes: self.planes,
            lights: self.lights,
            background: self.background,
            view: self.view,
        }
    }
}

pub struct Scene {
    objects: Vec<Object>,
    planes: Vec<Object>,
    bvh: Bvh,
    lights: Vec<Light>,
    pub background: Background,
    pub view: View,
}

impl Scene {
    pub fn object_count(&self) -> usize {
        self.objects.len() + self.planes.len()
    }

    pub fn bvh_depth(&self) -> usize {
        self.bvh.depth()
    }

    pub fn uses_avx512(&self) -> bool {
        self.bvh.uses_avx512()
    }

    pub fn light_count(&self) -> usize {
        self.lights.len()
    }

    /// Picks one light for shading `point`, with probability proportional to
    /// its estimated contribution there (power / distance^2), so a scene with
    /// many lights costs one shadow ray per shading point and mostly samples
    /// nearby lights. Returns the light and the probability of picking it.
    /// This is O(lights) per call; thousands of lights would want a light tree.
    pub fn pick_light<R: Rng>(&self, point: Vec3, rng: &mut R) -> Option<(&Light, f64)> {
        match self.lights.len() {
            0 => None,
            1 => Some((&self.lights[0], 1.0)),
            _ => {
                let total: f64 = self.lights.iter().map(|l| l.importance(point)).sum();
                let u = rng.gen::<f64>() * total;
                let mut acc = 0.0;
                for light in &self.lights {
                    let w = light.importance(point);
                    acc += w;
                    if u < acc {
                        return Some((light, w / total));
                    }
                }
                // Rounding left u just past the last bucket.
                let last = self.lights.last().unwrap();
                Some((last, last.importance(point) / total))
            }
        }
    }

    pub fn hit(&self, r: &Ray, t_min: f64, t_max: f64) -> Option<HitRecord> {
        let mut closest: Option<(f64, Vec3, Material)> = None;
        let mut min_t = t_max;
        for obj in &self.planes {
            if let Some((t, outward)) = obj.shape.hit(r, t_min, min_t) {
                min_t = t;
                closest = Some((t, outward, obj.material));
            }
        }
        // Spheres are intersected inside the BVH (8 at a time); other shapes
        // come back through this callback, which remembers their normal.
        let mut other_normal = Vec3::ZERO;
        let bvh_hit = self.bvh.closest(r, t_min, min_t, |i, t_max| {
            let (t, outward) = self.objects[i].shape.hit(r, t_min, t_max)?;
            other_normal = outward;
            Some(t)
        });
        if let Some((i, t, is_sphere)) = bvh_hit {
            let obj = &self.objects[i];
            let outward = match obj.shape {
                Shape::Sphere { center, .. } if is_sphere => (r.at(t) - center).normalize(),
                _ => other_normal,
            };
            closest = Some((t, outward, obj.material));
        }
        closest.map(|(t, outward, material)| HitRecord::new(r, t, outward, material))
    }

    pub fn blocked(&self, origin: Vec3, direction: Vec3, t_max: f64) -> bool {
        let r = Ray::from_unit(origin, direction);
        self.planes.iter().any(|obj| obj.shape.hit(&r, EPSILON, t_max).is_some())
            || self
                .bvh
                .any(&r, EPSILON, t_max, |i| self.objects[i].shape.hit(&r, EPSILON, t_max).is_some())
    }
}