use std::f32::consts::PI;

use glam::Vec3;
use rand_distr::{Distribution, Uniform, UnitSphere};

use super::Reflector;

#[derive(Debug, Clone)]
pub struct DiscReflector {
	pub roughness: f32,
}

impl Reflector for DiscReflector {
	fn brdf(&self, normal: Vec3, incident: Vec3, reflected: Vec3) -> f32 {
		let specular = incident.reflect(normal);
		let convergence = reflected.dot(specular);
		if convergence <= 1. - self.roughness {
			return 0.;
		}
		let c = -incident.dot(normal) * reflected.dot(normal);
		if c < 0. {
			return 0.;
		}
		0.5 / (self.roughness * PI * (specular.dot(normal) * reflected.dot(normal)).sqrt())
	}

	fn reflect(&self, rgen: &mut impl rand::Rng, normal: Vec3, incident: Vec3) -> Vec3 {
		let specular = incident.reflect(normal);

		let dist_theta = Uniform::new(0., (1. - self.roughness).acos()).unwrap();
		let dist_sphere = UnitSphere;

		loop {
			let theta: f32 = dist_theta.sample(rgen);
			let (xy, z) = theta.sin_cos();

			let spherical: Vec3 = dist_sphere.sample(rgen).into();
			let off_dir = spherical.reject_from_normalized(specular).normalize();

			let reflected = xy * off_dir + z * specular;
			if reflected.dot(normal) >= 0. {
				return reflected;
			}
		}
	}
}