use clap::Parser;
use indicatif::{ProgressBar, ProgressStyle};
use rand::{rngs::SmallRng, SeedableRng};
use ray_tracing_core::{
    camera::{BasicCamera, Camera},
    prelude::*,
    renderer::ClassicalRenderer,
    scene::Scene,
};
use ray_tracing_pbrt_scene::{parse_pbrt_v4, CameraType};
use ray_tracing_renderer::{
    depth_renderer::DepthRenderer, mis::MIS, next_event_estimation::NextEventEstimation,
    path_tracer::PathTracer, path_tracer_importance::PathTracerImportance,
};
use ray_tracing_scene::examples::example_scenes;
use rayon::prelude::*;
use std::net::TcpStream;
use tev_client::{PacketCreateImage, PacketUpdateImage, TevClient, TevError};

#[derive(Parser)]
struct Args {
    #[arg(long, default_value = "127.0.0.1:14158")]
    tev: String,
    #[arg(long)]
    renderers: Vec<String>,
    #[arg(long)]
    scenes: Vec<String>,
    #[arg(long, default_value_t = 400)]
    width: u32,
    #[arg(long, default_value_t = 400)]
    height: u32,
    #[arg(long, default_value_t = 1024)]
    samples_per_pixel: usize,
    #[arg(long)]
    pbrt_filename: Option<std::path::PathBuf>,
}

fn render_image<
    R: ClassicalRenderer<SmallRng, S, C> + Sync,
    S: Scene<SmallRng> + Sync,
    C: Camera<SmallRng> + Sync,
>(
    name: impl AsRef<str>,
    renderer: &R,
    scene: &S,
    camera: &C,
    samples_per_pixel: usize,
    tev: &mut TevClient,
) -> Result<(), TevError> {
    let mut data = vec![0.0; (renderer.width() * renderer.height() * 3) as usize];

    let bar = ProgressBar::new(renderer.width() as u64 * renderer.height() as u64);
    bar.set_style(
        ProgressStyle::with_template(
            "[{elapsed_precise}] [{wide_bar:.cyan/blue}] {percent_precise}% ({eta_precise})",
        )
        .unwrap()
        .progress_chars("#>-"),
    );

    let start = std::time::Instant::now();
    data.par_chunks_mut(3).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 {
            let r =
                renderer.render_pixel(scene, camera, x, y, &mut rng) / (samples_per_pixel as Float);
            if r.r().is_nan() || r.g().is_nan() || r.b().is_nan() {
                eprintln!("NAN x: {:?}, y: {:?}, r: {:?}", x, y, r);
            } else {
                c[0] += r.r();
                c[1] += r.g();
                c[2] += r.b();
            }
        }
        bar.inc(1);
    });

    bar.finish_and_clear();

    println!("Rendered \"{}\" in {:?}", name.as_ref(), start.elapsed());

    let channel_names = &["r", "g", "b"];

    tev.send(PacketCreateImage {
        image_name: name.as_ref(),
        grab_focus: false,
        width: renderer.width(),
        height: renderer.height(),
        channel_names,
    })?;

    tev.send(PacketUpdateImage {
        image_name: name.as_ref(),
        grab_focus: false,
        channel_names,
        channel_offsets: &[0, 1, 2],
        channel_strides: &[3, 3, 3],
        x: 0,
        y: 0,
        width: renderer.width(),
        height: renderer.height(),
        data: &data,
    })?;

    Ok(())
}

fn choose_renderer<S, C>(args: &Args, scene: &str, s: &S, c: &C, tev: &mut TevClient)
where
    S: Scene<SmallRng> + Sync,
    C: Camera<SmallRng> + Sync,
{
    if args.renderers.is_empty() || args.renderers.contains(&String::from("depth")) {
        let r = DepthRenderer::new(args.width, args.height);
        render_image(
            format!("{scene} - depth renderer"),
            &r,
            s,
            c,
            args.samples_per_pixel,
            tev,
        )
        .unwrap();
    }

    if args.renderers.is_empty() || args.renderers.contains(&String::from("path")) {
        let r = PathTracer::new(args.width, args.height);
        render_image(
            format!("{scene} - path tracer"),
            &r,
            s,
            c,
            args.samples_per_pixel,
            tev,
        )
        .unwrap();
    }

    if args.renderers.is_empty() || args.renderers.contains(&String::from("importance")) {
        let r = PathTracerImportance::new(args.width, args.height);
        render_image(
            format!("{scene} - path tracer importance"),
            &r,
            s,
            c,
            args.samples_per_pixel,
            tev,
        )
        .unwrap();
    }

    if args.renderers.is_empty() || args.renderers.contains(&String::from("nee")) {
        let r = NextEventEstimation::new(args.width, args.height);
        render_image(
            format!("{scene} - next event estimation"),
            &r,
            s,
            c,
            args.samples_per_pixel,
            tev,
        )
        .unwrap();
    }

    if args.renderers.is_empty() || args.renderers.contains(&String::from("mis")) {
        let r = MIS::new(args.width, args.height);
        render_image(
            format!("{scene} - mis"),
            &r,
            s,
            c,
            args.samples_per_pixel,
            tev,
        )
        .unwrap();
    }
}
fn main() {
    let args = Args::parse();

    let mut client = TevClient::wrap(TcpStream::connect(&args.tev).unwrap());

    let map = example_scenes::<SmallRng>();

    let scenes: Vec<&str> = if args.pbrt_filename.is_none() {
        if args.scenes.is_empty() {
            map.keys().copied().collect()
        } else {
            args.scenes.iter().map(|s| s.as_str()).collect()
        }
    } else {
        Vec::new()
    };

    for scene in scenes {
        let f = map.get(scene).unwrap();

        let s = f.get_scene();

        let c = BasicCamera::from_look_at(
            args.width,
            args.height,
            f.get_camera_pos(),
            f.get_camera_look_at(),
            f.get_camera_up(),
            f.get_horizontal_fov(),
        );
        match &s {
            ray_tracing_scene::examples::ExampleSceneEnum::AccelerationStructureScene(s) => {
                choose_renderer(&args, scene, s, &c, &mut client);
            }
            ray_tracing_scene::examples::ExampleSceneEnum::TriangleBVH(s) => {
                choose_renderer(&args, scene, s, &c, &mut client);
            }
            ray_tracing_scene::examples::ExampleSceneEnum::BasicScene(s) => {
                choose_renderer(&args, scene, s, &c, &mut client);
            }
        }
    }

    if let Some(pbrt_filename) = &args.pbrt_filename {
        let pbrt = match parse_pbrt_v4(pbrt_filename) {
            Ok(pbrt) => pbrt,
            Err(e) => {
                println!("{e:?}");
                std::process::exit(1);
            }
        };

        let fov = match pbrt.settings.camera.camera_type {
            CameraType::Perspective { fov, .. } => 2.0 * fov,
            _ => todo!(),
        };

        let c = BasicCamera::from_affine_transform(
            args.width,
            args.height,
            pbrt.settings.camera_ctm,
            Float::to_radians(fov),
        );

        // dbg!(&c);

        // let c = BasicCamera::from_look_at(
        //     args.width,
        //     args.height,
        //     Pos3::new(3.0, 5.0, -5.0),
        //     Pos3::new(0.0, 0.075, 0.0),
        //     Dir3::new(0.0, 1.0, 0.0),
        //     Float::to_radians(2.0 * 37.0),
        // );

        // dbg!(&c);

        let s = &pbrt.scene;

        choose_renderer(
            &args,
            &format!("pbrt \"{pbrt_filename:?}\""),
            s,
            &c,
            &mut client,
        );
    }
}
