use rand::Rng;
use rand_distr::StandardNormal;

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

pub const EPSILON: f64 = 1e-6;

pub fn random_unit_vector<R: Rng>(rng: &mut R) -> Vec3 {
    loop {
        let v = Vec3::new(
            rng.sample(StandardNormal),
            rng.sample(StandardNormal),
            rng.sample(StandardNormal),
        );
        let n = v.length();
        if n > 1e-12 {
            return v / n;
        }
    }
}

#[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(),
        }
    }

    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 } => {
                // normal + uniform unit vector is cosine-distributed, so the
                // Lambertian BRDF * cos / pdf reduces to just the albedo.
                let mut target = rec.normal + random_unit_vector(rng);
                if target.length() < 1e-8 {
                    target = rec.normal;
                }
                Some((albedo, Ray::new(rec.point, target)))
            }
            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 },
}

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])))
            }
        }
    }
}

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

/// An axis-aligned box light, sampled explicitly for next event estimation.
#[derive(Clone, Copy, Debug)]
pub struct BoxLight {
    pub min: Vec3,
    pub max: Vec3,
    pub emission: Vec3,
}

#[derive(Default)]
pub struct Scene {
    pub objects: Vec<Object>,
    pub lights: Vec<BoxLight>,
}

impl Scene {
    pub fn add(&mut self, shape: Shape, material: Material) {
        self.objects.push(Object { shape, material });
    }

    pub fn add_box_light(&mut self, min: Vec3, max: Vec3, emission: Vec3) {
        self.lights.push(BoxLight { min, max, emission });
        self.add(Shape::Box { min, max }, Material::Light { emission });
    }

    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.objects {
            if let Some((t, outward)) = obj.shape.hit(r, t_min, min_t) {
                if t < min_t {
                    min_t = t;
                    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::new(origin, direction);
        self.objects
            .iter()
            .any(|obj| obj.shape.hit(&r, EPSILON, t_max).is_some())
    }
}

pub fn cornell_box() -> Scene {
    let mut scene = Scene::default();

    let white = Material::Diffuse { albedo: Vec3::new(0.73, 0.73, 0.73) };
    let gray = Material::Diffuse { albedo: Vec3::new(0.5, 0.5, 0.5) };
    let red = Material::Diffuse { albedo: Vec3::new(0.65, 0.05, 0.05) };
    let green = Material::Diffuse { albedo: Vec3::new(0.12, 0.45, 0.15) };
    let mirror = Material::Mirror { albedo: Vec3::new(0.95, 0.95, 0.95) };
    let glass = Material::Dielectric { ref_idx: 1.5, albedo: Vec3::ONE };

    let plane = |p: Vec3, n: Vec3| Shape::Plane { point: p, normal: n.normalize() };
    scene.add(plane(Vec3::new(0.0, 0.0, 0.0), Vec3::new(0.0, 1.0, 0.0)), gray);
    scene.add(plane(Vec3::new(0.0, 500.0, 0.0), Vec3::new(0.0, -1.0, 0.0)), white);
    scene.add(plane(Vec3::new(0.0, 0.0, 500.0), Vec3::new(0.0, 0.0, -1.0)), white);
    scene.add(plane(Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 0.0, 0.0)), red);
    scene.add(plane(Vec3::new(500.0, 0.0, 0.0), Vec3::new(-1.0, 0.0, 0.0)), green);

    scene.add_box_light(
        Vec3::new(0.0, 180.0, 180.0),
        Vec3::new(25.0, 220.0, 320.0),
        Vec3::splat(50.0),
    );

    let sphere = |c: Vec3, r: f64| Shape::Sphere { center: c, radius: r };
    scene.add(sphere(Vec3::new(250.0, 180.0, 250.0), 80.0), glass);
    scene.add(sphere(Vec3::new(150.0, 180.0, 380.0), 80.0), mirror);
    scene.add(sphere(Vec3::new(350.0, 160.0, 300.0), 60.0), red);

    scene
}