use image::{ImageBuffer, ImageResult, Rgb};
use rand::{rngs::SmallRng, SeedableRng};
use ray_tracing_core::{
    camera::{BasicCamera, Camera},
    prelude::*,
    renderer::ClassicalRenderer,
    scene::Scene,
};
use ray_tracing_renderer::path_tracer_importance::PathTracerImportance;
use ray_tracing_scene::examples::{basic_cornell, ExampleScene};
use rayon::prelude::*;
use std::path::Path;

fn render_image<
    R: ClassicalRenderer<SmallRng, S, C> + Sync,
    S: Scene<SmallRng> + Sync,
    C: Camera<SmallRng> + Sync,
>(
    renderer: &R,
    scene: &S,
    camera: &C,
    outputfilename: impl AsRef<Path>,
    samples_per_pixel: usize,
) -> ImageResult<()> {
    let mut data = vec![Color::black(); (renderer.width() * renderer.height()) as usize];

    data.par_iter_mut().enumerate().for_each(|(i, c)| {
        let x = (i % renderer.width() as usize) as u32;
        let y = (i / renderer.width() as usize) as u32;

        let mut rng = SmallRng::seed_from_u64((x + y * renderer.width()) as u64);
        for _ in 0..samples_per_pixel {
            *c +=
                renderer.render_pixel(scene, camera, x, y, &mut rng) / (samples_per_pixel as Float);
        }
    });

    let img = ImageBuffer::from_fn(renderer.width(), renderer.height(), |x, y| {
        let c = data[(x + y * renderer.width()) as usize];
        Rgb::<Float>([c.r(), c.g(), c.b()])
    });

    img.save(outputfilename)
}

fn main() -> ImageResult<()> {
    // let s = BasicScene::new();
    let s = basic_cornell::BasicCornell::new();

    let c = BasicCamera::from_look_at(
        400,
        400,
        s.get_camera_pos(),
        s.get_camera_look_at(),
        s.get_camera_up(),
        s.get_horizontal_fov(),
    );

    let r = PathTracerImportance::new(400, 400);

    let s = s.get_scene();
    match &s {
        ray_tracing_scene::examples::ExampleSceneEnum::AccelerationStructureScene(s) => {
            render_image(&r, s, &c, "test.exr", 1048)
        }
        ray_tracing_scene::examples::ExampleSceneEnum::TriangleBVH(s) => {
            render_image(&r, s, &c, "test.exr", 1048)
        }
        ray_tracing_scene::examples::ExampleSceneEnum::BasicScene(s) => {
            render_image(&r, s, &c, "test.exr", 1048)
        }
    }
}
