use crate::hittable::Hit;
use crate::ray::{NextRay, Ray};
use crate::utils::{random_in_unit_sphere, random_unit_vector, reflect, refract, schlick};
use cgmath::prelude::*;
use cgmath::{dot, vec3, Vector3};
use rand::prelude::*;

pub trait Material {
    fn scatter(&self, rng: &mut ThreadRng, ray: &Ray, hit: &Hit) -> Option<NextRay>;
}

pub struct Lambertian {
    pub albedo: Vector3<f64>,
}

impl Lambertian {
    pub fn new(albedo: Vector3<f64>) -> Lambertian {
        Lambertian { albedo }
    }
}

impl Material for Lambertian {
    fn scatter(&self, rng: &mut ThreadRng, ray: &Ray, hit: &Hit) -> Option<NextRay> {
        let scatter_direction = hit.normal + random_unit_vector(rng);
        Some(NextRay::new(
            self.albedo,
            Ray::new(hit.p, scatter_direction, ray.time),
        ))
    }
}

pub struct Metal {
    pub albedo: Vector3<f64>,
    pub fuzz: f64,
}

impl Metal {
    pub fn new(albedo: Vector3<f64>, f: f64) -> Metal {
        let fuzz = if f < 1.0 { f } else { 1.0 };
        Metal { albedo, fuzz }
    }
}

impl Material for Metal {
    fn scatter(&self, rng: &mut ThreadRng, ray: &Ray, hit: &Hit) -> Option<NextRay> {
        let reflected = reflect(&ray.dir.normalize(), &hit.normal);
        let scattered = Ray::new(
            hit.p,
            reflected + self.fuzz * random_in_unit_sphere(rng),
            ray.time,
        );
        if dot(scattered.dir, hit.normal) > 0.0 {
            Some(NextRay::new(self.albedo, scattered))
        } else {
            None
        }
    }
}

pub struct Dielectric {
    pub ref_idx: f64,
}

impl Dielectric {
    pub fn new(ref_idx: f64) -> Dielectric {
        Dielectric { ref_idx }
    }
}

impl Material for Dielectric {
    fn scatter(&self, rng: &mut ThreadRng, ray: &Ray, hit: &Hit) -> Option<NextRay> {
        let attenuation = vec3(1.0, 1.0, 1.0);
        let etai_over_etat = if hit.front_face {
            1.0 / self.ref_idx
        } else {
            self.ref_idx
        };

        let unit_direction = ray.dir.normalize();
        let cos_theta = f64::min(dot(-unit_direction, hit.normal), 1.0);
        let sin_theta = (1.0 - cos_theta * cos_theta).sqrt();

        if etai_over_etat * sin_theta > 1.0 {
            let reflected = reflect(&unit_direction, &hit.normal);
            return Some(NextRay::new(
                attenuation,
                Ray::new(hit.p, reflected, ray.time),
            ));
        }

        let reflect_prob = schlick(cos_theta, etai_over_etat);
        if rng.gen::<f64>() < reflect_prob {
            let reflected = reflect(&unit_direction, &hit.normal);
            return Some(NextRay::new(
                attenuation,
                Ray::new(hit.p, reflected, ray.time),
            ));
        }

        let refracted = refract(&unit_direction, &hit.normal, etai_over_etat);
        Some(NextRay::new(
            attenuation,
            Ray::new(hit.p, refracted, ray.time),
        ))
    }
}