mod camera;
mod scene;
mod vec3;

use std::f64::consts::PI;
use std::time::Instant;

use indicatif::{ProgressBar, ProgressStyle};
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use rayon::prelude::*;

use camera::Camera;
use scene::{cornell_box, Material, Ray, Scene, EPSILON};
use vec3::Vec3;

fn trace<R: Rng>(scene: &Scene, ray: &Ray, depth: u32, max_depth: u32, rng: &mut R) -> Vec3 {
    if depth >= max_depth {
        return Vec3::ZERO;
    }

    let Some(rec) = scene.hit(ray, EPSILON, f64::INFINITY) else {
        let t = 0.5 * (ray.direction.y + 1.0);
        return Vec3::ONE * (1.0 - t) + Vec3::new(0.5, 0.7, 1.0) * t;
    };

    let emitted = rec.material.emission();
    if let Material::Light { .. } = rec.material {
        return emitted;
    }

    let Some((attenuation, scattered)) = rec.material.scatter(ray, &rec, rng) else {
        return emitted;
    };

    let indirect = trace(scene, &scattered, depth + 1, max_depth, rng);
    let mut color = emitted + attenuation.mul_elem(indirect);

    if let Material::Diffuse { albedo } = rec.material {
        let mut direct = Vec3::ZERO;
        for light in &scene.lights {
            let (lo, hi) = (light.min, light.max);
            for _ in 0..2 {
                let face = (rng.gen::<f64>() * 6.0) as usize;
                let x = lo.x + (hi.x - lo.x) * rng.gen::<f64>();
                let y = lo.y + (hi.y - lo.y) * rng.gen::<f64>();
                let z = lo.z + (hi.z - lo.z) * rng.gen::<f64>();
                let lp = match face {
                    0 => Vec3::new(lo.x, y, z),
                    1 => Vec3::new(hi.x, y, z),
                    2 => Vec3::new(x, lo.y, z),
                    3 => Vec3::new(x, hi.y, z),
                    4 => Vec3::new(x, y, lo.z),
                    _ => Vec3::new(x, y, hi.z),
                };

                let to_light = lp - rec.point;
                let dist_sq = to_light.dot(to_light);
                if dist_sq < 1e-8 {
                    continue;
                }
                let dist = dist_sq.sqrt();
                let d_to_light = to_light / dist;
                let cos_at_rec = rec.normal.dot(d_to_light);
                if cos_at_rec <= 0.0 {
                    continue;
                }

                if scene.blocked(rec.point + EPSILON * rec.normal, d_to_light, dist - EPSILON) {
                    continue;
                }

                let light_area = 6.0 * (hi.x - lo.x) * (hi.y - lo.y) * (hi.z - lo.z);
                let pdf = light_area / (dist_sq * cos_at_rec);
                if pdf < 1e-12 {
                    continue;
                }
                let brdf = albedo / PI;
                direct += light.emission.mul_elem(brdf) * (cos_at_rec / pdf);
                break;
            }
        }

        if direct.dot(direct) > 0.0 {
            color = emitted + attenuation.mul_elem(indirect + direct);
        }
    }

    color
}

struct Settings {
    width: usize,
    height: usize,
    spp: u32,
    max_depth: u32,
    output: String,
}

fn render(scene: &Scene, s: &Settings) -> Vec<Vec3> {
    let (width, height) = (s.width, s.height);
    let cam = Camera::new(
        Vec3::new(278.0, 278.0, -800.0),
        Vec3::new(278.0, 278.0, 278.0),
        Vec3::new(0.0, 1.0, 0.0),
        40.0,
        width as f64 / height as f64,
        1078.0,
    );

    let pb = ProgressBar::new(height as u64);
    pb.set_style(
        ProgressStyle::with_template(
            "{spinner:.green} Rendering [{bar:50.bright_green/blue}] {percent:>3}% {pos}/{len} rows  {elapsed_precise}  ETA {eta_precise}",
        )
        .unwrap()
        .progress_chars("=> "),
    );

    let mut image = vec![Vec3::ZERO; width * height];
    image
        .par_chunks_mut(width)
        .enumerate()
        .for_each(|(j, row)| {
            for (i, px) in row.iter_mut().enumerate() {
                let mut rng = SmallRng::seed_from_u64((j * width + i) as u64 ^ 0x9E37_79B9_7F4A_7C15);
                let mut color = Vec3::ZERO;
                for _ in 0..s.spp {
                    let u = (i as f64 + rng.gen::<f64>()) / (width - 1) as f64;
                    let v = ((height - 1 - j) as f64 + rng.gen::<f64>()) / (height - 1) as f64;
                    let ray = cam.get_ray(u, v);
                    color += trace(scene, &ray, 0, s.max_depth, &mut rng);
                }
                color /= s.spp as f64;
                *px = color.map(|c| c.max(0.0).sqrt().min(1.0));
            }
            pb.inc(1);
        });
    pb.finish();

    image
}

fn save_png(image: &[Vec3], width: usize, height: usize, path: &str) -> image::ImageResult<()> {
    let bytes: Vec<u8> = image
        .iter()
        .flat_map(|c| [c.x, c.y, c.z])
        .map(|c| (c.clamp(0.0, 1.0) * 255.0) as u8)
        .collect();
    image::save_buffer(path, &bytes, width as u32, height as u32, image::ColorType::Rgb8)
}

fn region_mean(image: &[Vec3], width: usize, rows: std::ops::Range<usize>, cols: std::ops::Range<usize>) -> f64 {
    let n = rows.len() * cols.len();
    if n == 0 {
        return 0.0;
    }
    let sum: f64 = rows
        .flat_map(|j| cols.clone().map(move |i| image[j * width + i].mean()))
        .sum();
    sum / n as f64
}

fn print_region_table(image: &[Vec3], width: usize, height: usize) {
    // Same regions as the Python version (which slices the Floor/Ceiling/Center
    // columns by height-based indices).
    let (h_start, h_end) = (height / 3, 2 * height / 3);
    let clamp_cols = |r: std::ops::Range<usize>| r.start.min(width)..r.end.min(width);
    let regions = [
        ("Left Wall", 0..height, 0..width / 4),
        ("Right Wall", 0..height, 3 * width / 4..width),
        ("Floor", 2 * height / 3..height, clamp_cols(h_start..h_end)),
        ("Ceiling", 0..height / 3, clamp_cols(h_start..h_end)),
        ("Center", h_start..h_end, clamp_cols(h_start..h_end)),
    ];

    println!("Region Brightness Analysis");
    println!("{:<12} {:>10}  Visual", "Region", "Brightness");
    for (name, rows, cols) in regions {
        let brightness = region_mean(image, width, rows, cols) * 255.0;
        let filled = ((brightness / 255.0 * 20.0) as usize).min(20);
        let bar = format!("{}{}", "#".repeat(filled), "-".repeat(20 - filled));
        let color = if brightness > 80.0 { 32 } else if brightness > 40.0 { 33 } else { 31 };
        println!("{name:<12} {brightness:>10.1}  \x1b[{color}m{bar}\x1b[0m");
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let arg = |i: usize, default: &str| args.get(i).cloned().unwrap_or_else(|| default.to_string());
    let parse = |i: usize, default: &str| -> usize {
        arg(i, default).parse().unwrap_or_else(|_| {
            eprintln!("usage: pathtracer [width] [height] [samples_per_pixel] [max_bounces] [output]");
            std::process::exit(2);
        })
    };

    let settings = Settings {
        width: parse(1, "400"),
        height: parse(2, "300"),
        spp: parse(3, "50") as u32,
        max_depth: parse(4, "15") as u32,
        output: arg(5, "cornell_box.png"),
    };
    if settings.width < 2 || settings.height < 2 || settings.spp == 0 {
        eprintln!("width and height must be >= 2 and samples_per_pixel >= 1");
        std::process::exit(2);
    }

    println!("=== Cornell Box Path Tracer (Rust) ===");
    println!("Features: reflections, refraction/caustics, multiple bounces, NEE");
    println!();
    println!("Resolution:        {}x{}", settings.width, settings.height);
    println!("Samples per pixel: {}", settings.spp);
    println!("Max bounces:       {}", settings.max_depth);
    println!("Threads:           {}", rayon::current_num_threads());
    println!("Output:            {}", settings.output);
    println!();

    let scene = cornell_box();
    let start = Instant::now();
    let image = render(&scene, &settings);
    let elapsed = start.elapsed().as_secs_f64();

    if let Err(e) = save_png(&image, settings.width, settings.height, &settings.output) {
        eprintln!("failed to write {}: {e}", settings.output);
        std::process::exit(1);
    }

    println!();
    println!("Rendered {}", settings.output);
    println!(
        "Total time: {:.2}s | Speed: {:.2} MPix*samp/s",
        elapsed,
        (settings.width * settings.height) as f64 * settings.spp as f64 / elapsed / 1e6
    );
    println!();
    print_region_table(&image, settings.width, settings.height);
}