use crate::core::image::Image;
use crate::core::{options::PBRTOptions, sampler::get_camera_sample};
use indicatif::{ProgressBar, ProgressStyle};
use std::io::Write;
use std::path::Path;

use super::*;

struct PbrtProgress {
    bar: ProgressBar,
}

impl PbrtProgress {
    fn new(total_work: u64, description: &str, quiet: bool) -> Self {
        if quiet {
            return Self {
                bar: ProgressBar::hidden(),
            };
        }

        let bar = ProgressBar::new(total_work);

        bar.set_style(
            ProgressStyle::default_bar()
                .template("[{elapsed_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}")
                .unwrap()
                .progress_chars("=>-"),
        );

        bar.set_message(description.to_string());

        Self { bar }
    }

    fn update(&self, amount: u64) {
        self.bar.inc(amount);
    }

    fn done(&self) {
        self.bar.finish_with_message("Done");
    }

    fn elapsed_seconds(&self) -> f32 {
        self.bar.elapsed().as_secs_f32()
    }
}

fn generate_tiles(bounds: Bounds2i) -> Vec<Bounds2i> {
    let mut tiles = Vec::new();
    const TILE_SIZE: i32 = 16;
    for y in (bounds.p_min.y()..bounds.p_max.y()).step_by(TILE_SIZE as usize) {
        for x in (bounds.p_min.x()..bounds.p_max.x()).step_by(TILE_SIZE as usize) {
            let p_min = Point2i::new(x, y);
            let p_max = Point2i::new(
                (x + TILE_SIZE).min(bounds.p_max.x()),
                (y + TILE_SIZE).min(bounds.p_max.y()),
            );
            tiles.push(Bounds2i::from_points(p_min, p_max));
        }
    }
    tiles
}

pub fn render<T>(
    integrator: &T,
    _base: &IntegratorBase,
    camera: &Camera,
    sampler_prototype: &Sampler,
) where
    T: RayIntegratorTrait,
{
    let options = get_options();
    if let Some((p_pixel, sample_index)) = options.debug_start {
        let s_index = sample_index as usize;
        let mut tile_sampler = sampler_prototype.clone();

        tile_sampler.start_pixel_sample(p_pixel, s_index, None);

        evaluate_pixel_sample(integrator, camera, &mut tile_sampler, p_pixel, s_index);
        return;
    }

    let pixel_bounds = camera.get_film().pixel_bounds();
    let spp = sampler_prototype.samples_per_pixel();
    let total_work = (pixel_bounds.area() as u64) * (spp as u64);
    let progress = PbrtProgress::new(total_work, "Rendering", options.quiet);
    let mut wave_start = 0;
    let mut wave_end = 1;
    let mut next_wave_size = 1;

    let mut reference_image: Option<Image> = None;
    let mut mse_out_file: Option<std::fs::File> = None;

    if let Some(ref_path) = &options.mse_reference_image {
        let image_and_metadata =
            Image::read(Path::new(&ref_path), None).expect("Could not load image");
        let image = image_and_metadata.image;
        let metadata = image_and_metadata.metadata;
        let resolution = image.resolution();
        // reference_image = Some(image);
        let mse_pixel_bounds = metadata
            .pixel_bounds
            .unwrap_or_else(|| Bounds2i::from_points(Point2i::new(0, 0), resolution));

        if !mse_pixel_bounds.overlaps(&pixel_bounds) {
            panic!("Output pixel bounds dont fit inside the reference image");
        }

        let crop_p_min = Point2i::from(pixel_bounds.p_min - mse_pixel_bounds.p_min);
        let crop_p_max = Point2i::from(pixel_bounds.p_max - mse_pixel_bounds.p_min);

        let crop_bounds = Bounds2i::from_points(crop_p_min, crop_p_max);

        let cropped_image = image.crop(crop_bounds).clone();
        let cropped_resolution = cropped_image.resolution();

        let expected_res = Point2i::new(
            pixel_bounds.p_max.x() - pixel_bounds.p_min.x(),
            pixel_bounds.p_max.y() - pixel_bounds.p_min.y(),
        );

        reference_image = Some(cropped_image);

        assert_eq!(
            cropped_resolution, expected_res,
            "Cropped reference image resolution mismatch"
        );

        if let Some(out_path) = &options.mse_reference_output {
            mse_out_file = Some(
                std::fs::File::create(out_path)
                    .expect(&format!("Failed to create MSE output file: {}", out_path)),
            );
        }
    }

    let tiles = generate_tiles(pixel_bounds);
    while wave_start < spp {
        tiles.par_iter().for_each(|tile_bounds| {
            let mut sampler = sampler_prototype.clone();

            for p_pixel in tile_bounds {
                for sample_index in wave_start..wave_end {
                    sampler.start_pixel_sample(*p_pixel, sample_index, None);
                    evaluate_pixel_sample(integrator, camera, &mut sampler, *p_pixel, sample_index);
                }
            }

            let work_done = (tile_bounds.area() as u64) * ((wave_end - wave_start) as u64);
            progress.update(work_done);
        });

        wave_start = wave_end;
        wave_end = (wave_end + next_wave_size).min(spp);

        if reference_image.is_none() {
            next_wave_size = (2 * next_wave_size).min(64);
        }

        if wave_start == spp {
            progress.done();
        }

        if wave_start == spp || options.write_partial_images || reference_image.is_some() {
            let mut metadata = ImageMetadata {
                render_time_seconds: Some(progress.elapsed_seconds()),
                samples_per_pixel: Some(wave_start as i32),
                ..Default::default()
            };

            if wave_start == spp || options.write_partial_images {
                camera.init_metadata(&mut metadata);
                camera
                    .get_film()
                    .write_image(&metadata, 1.0 / wave_start as Float);
            }

            if let Some(ref_img) = &reference_image {
                let splat_scale = 1.0 / (wave_start as Float);

                let film_metadata = ImageMetadata::default();
                let film_image = camera.get_film().get_image(&film_metadata, splat_scale);

                let (mse_values, _mse_debug_img) =
                    film_image.mse(film_image.all_channels_desc(), ref_img, false);
                let mse_avg = mse_values.average();

                if let Some(file) = &mut mse_out_file {
                    writeln!(file, "{}, {:.9}", wave_start, mse_avg).ok();
                    file.flush().ok();
                }

                metadata.mse = Some(mse_avg);
            }
        }
    }
}

pub fn evaluate_pixel_sample<T: RayIntegratorTrait>(
    integrator: &T,
    camera: &Camera,
    sampler: &mut Sampler,
    pixel: Point2i,
    _sample_index: usize,
) {
    let mut lu = sampler.get1d();
    if get_options().disable_wavelength_jitter {
        lu = 0.5;
    }
    let lambda = camera.get_film().sample_wavelengths(lu);
    let film = camera.get_film();
    let filter = film.get_filter();
    let camera_sample = get_camera_sample(sampler, pixel, filter);
    if let Some(mut camera_ray) = camera.generate_ray_differential(camera_sample, &lambda) {
        debug_assert!(camera_ray.ray.d.norm() > 0.999);
        debug_assert!(camera_ray.ray.d.norm() < 1.001);
        let ray_diff_scale = (sampler.samples_per_pixel() as Float).sqrt().max(0.125);
        if get_options().disable_pixel_jitter {
            camera_ray.ray.scale_differentials(ray_diff_scale);
        }

        let initialize_visible_surface = film.uses_visible_surface();
        let (mut l, visible_surface) =
            integrator.li(camera_ray.ray, &lambda, sampler, initialize_visible_surface);
        l *= camera_ray.weight;

        if l.has_nans() || l.y(&lambda).is_infinite() {
            l = SampledSpectrum::new(0.);
        }

        film.add_sample(
            pixel,
            l,
            &lambda,
            visible_surface.as_ref(),
            camera_sample.filter_weight,
        );
    }
}