use std::ptr::null_mut;

use anyhow::Result;
use embree4_rs::{geometry::TriangleMeshGeometry, Device, Scene, SceneOptions};
use embree4_sys::{RTCFeatureFlags, RTCIntersectArguments};
use glam::Vec3;
use rand::{rngs::StdRng, Rng, SeedableRng};
use rayon::prelude::{IntoParallelIterator, ParallelIterator};

pub fn main() -> Result<()> {
    let device = Device::try_new(None)?;

    // Build a random triangle mesh by generating random vertices and indices
    let num_tris = 1_000_000;
    let mesh = build_triangle_mesh(&device, num_tris)?;

    // Construct a scene with high build quality and robust intersection tests
    let scene = Scene::try_new(
        device,
        SceneOptions {
            build_quality: embree4_sys::RTCBuildQuality::HIGH,
            flags: embree4_sys::RTCSceneFlags::ROBUST,
        },
    )?;
    scene.attach_geometry(&mesh)?;
    let scene = scene.commit()?;

    let bounds = scene.bounds()?;
    println!("Scene bounds: {:?}", bounds);

    // Generate random rays, with the origin at (0, 0) and a random direction
    let num_rays = 1_000_000;
    let rays = generate_rays(num_rays);

    let args = RTCIntersectArguments {
        flags: embree4_sys::RTCRayQueryFlags::INCOHERENT,
        feature_mask: RTCFeatureFlags::RTC_FEATURE_FLAG_ALL,
        context: null_mut(),
        filter: None,
        intersect: None,
    };

    // Perform intersection tests in parallel.
    // Count the number of hits and measure the time
    let t0 = std::time::Instant::now();
    let hits: usize = rays
        .into_par_iter()
        .map(|ray| {
            scene
                .intersect_1(ray, Some(&args))
                .unwrap()
                .map_or(0, |_| 1)
        })
        .sum();
    let elapsed = t0.elapsed();
    let rays_per_sec = (num_rays as f32 / elapsed.as_secs_f32()) as usize;

    println!("Traced {} rays in {:?}", num_rays, elapsed);
    let frac_hits = hits as f32 / num_rays as f32;
    println!("  {} hits ({:.3}%)", hits, 100.0 * frac_hits);
    println!("  ({} rays/s)", rays_per_sec);

    Ok(())
}

fn generate_rays(num_rays: i32) -> Vec<embree4_sys::RTCRay> {
    let mut rng = StdRng::seed_from_u64(0);
    let rays: Vec<_> = (0..num_rays)
        .map(|_| {
            let origin = Vec3::ZERO;
            let direction = 2.0 * rng.gen::<Vec3>() - 1.0;

            embree4_sys::RTCRay {
                org_x: origin.x,
                org_y: origin.y,
                org_z: origin.z,
                dir_x: direction.x,
                dir_y: direction.y,
                dir_z: direction.z,
                ..Default::default()
            }
        })
        .collect();
    rays
}

fn build_triangle_mesh(
    device: &Device,
    num_tris: usize,
) -> Result<TriangleMeshGeometry, anyhow::Error> {
    let mut rng = StdRng::seed_from_u64(0);

    let mut vertices = Vec::with_capacity(3 * num_tris);
    let mut indices = Vec::with_capacity(num_tris);
    for i in 0..num_tris as u32 {
        let pos = 1_000.0 * (2.0 * rng.gen::<Vec3>() - 1.0);

        let p = pos + rng.gen::<Vec3>();
        let q = pos + rng.gen::<Vec3>();
        let r = pos + rng.gen::<Vec3>();

        vertices.push((p.x, p.y, p.z));
        vertices.push((q.x, q.y, q.z));
        vertices.push((r.x, r.y, r.z));

        indices.push((3 * i, 3 * i + 1, 3 * i + 2));
    }
    let mesh = TriangleMeshGeometry::try_new(device, &vertices, &indices)?;
    Ok(mesh)
}