use image::{ExtendedColorType, ImageEncoder, codecs::png::PngEncoder};
use log::{error, info, trace};
use rayon::iter::IntoParallelIterator;
use rayon::prelude::*;
use std::{fs::File, sync::Arc, time::Instant};

use crate::{
    objects::{hit::Hit, traits::Hittable},
    ray::Ray,
    scenes::scene::Scene,
    vec3::Colour,
};

pub fn write_image(filename: &str, pixels: &[u8], width: u32, height: u32) {
    match File::create(filename) {
        Ok(out) => {
            let enc = PngEncoder::new(out);
            if let Err(e) = enc.write_image(pixels, width, height, ExtendedColorType::Rgb8) {
                error!("Failed writing image: {}", e);
            }
        }
        Err(e) => error!("Failed to create image file: {} ", e),
    }
}

pub fn render(scene: &mut Scene) {
    scene.init();

    // TODO (low priority): currently splits per vertical line, but could be more granular (per chunk)
    let mut pixels = vec![0_u8; (scene.get_image_width() * scene.get_image_height() * 3) as usize];
    let scanlines: Vec<(usize, &mut [u8])> = pixels
        .chunks_mut(scene.get_image_width() as usize * 3)
        .enumerate()
        .collect();

    let now = Instant::now();
    scanlines.into_par_iter().for_each(|(i, chunk)| {
        render_chunk(chunk, scene, i as u32);
    });
    info!(
        "rendering took {}.{} seconds.",
        now.elapsed().as_secs(),
        now.elapsed().subsec_nanos()
    );

    info!("Writing image file...");
    write_image(
        &scene.get_filename(),
        &pixels,
        scene.get_image_width() as u32,
        scene.get_image_height() as u32,
    );
}

pub fn render_chunk(chunk: &mut [u8], scene: &Scene, y: u32) {
    let camera = &scene.get_camera();
    let aa_rate = camera.get_anti_alias_rate();

    for i in 0..scene.get_image_width() {
        let pixel_tl = camera.get_pixel_tl(i, y);

        let mut pixel_colour = Colour::default();
        for y in 1..(aa_rate + 1) {
            for x in 1..(aa_rate + 1) {
                let r = camera.get_ray(pixel_tl, x, y);
                pixel_colour += ray_colour(&scene.get_objects(), &r, scene.get_max_depth());
            }
        }

        let (r, g, b) = (pixel_colour / (aa_rate * aa_rate) as f32).output();
        chunk[(i * 3) as usize] = r;
        chunk[(i * 3) as usize + 1] = g;
        chunk[(i * 3) as usize + 2] = b;
    }
    trace!("Line {}\tfinished.", y);
}

pub fn ray_colour(hittables: &Vec<Arc<dyn Hittable>>, r: &Ray, depth: u32) -> Colour {
    if depth == 0 {
        return Colour::default();
    }

    let closest = Hit::hit_list(hittables, r);
    if let Some(hit) = closest {
        if let Some((scattered, att)) = hit.mat().scatter(&hit, r) {
            if let Some(ray) = scattered {
                return att * ray_colour(hittables, &ray, depth - 1);
            }
            return att;
        }
        return Colour::default();
    }

    // background
    let unit_dir = r.dir().get_unit();
    let a = 0.5 * (unit_dir.y() + 1.);

    (1.0 - a) * Colour::new(1., 1., 1.) + a * Colour::new(0.5, 0.7, 1.0)
}