use std::sync::Arc;

use image::EncodableLayout;
use log::info;
use wgpu::{Extent3d, ImageCopyTexture, Origin3d, Texture, TextureUsages};

use crate::{
    app::App,
    diff::Differ,
    export::Exporter,
    helper::write_texture,
    paint::{PaintUniforms, Painter},
};

pub struct Approximator {
    app: Arc<App>,
    size: Extent3d,
    differ: Differ,
    painter: Painter,
    exporter: Exporter,

    tex_approx: Texture,
    tex_savestate: Texture,
    tex_target: Texture,
}

impl Approximator {
    pub fn new(app: &Arc<App>, tex_target: Texture, size: Extent3d) -> Self {
        let App { device, queue, .. } = app.as_ref();
        let tex_approx = device.create_texture(&wgpu::TextureDescriptor {
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: TextureUsages::COPY_DST
                | TextureUsages::TEXTURE_BINDING
                | TextureUsages::STORAGE_BINDING
                | TextureUsages::COPY_SRC,
            label: None,
        });
        let tex_savestate = device.create_texture(&wgpu::TextureDescriptor {
            size,
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: TextureUsages::COPY_DST | TextureUsages::COPY_SRC,
            label: None,
        });

        let img_initial = image::open("a/initial.png").unwrap().into_rgba8();
        write_texture(queue, &tex_approx, img_initial.as_bytes(), size);

        let differ = Differ::new(&app, size, &tex_approx, &tex_target);
        let painter = Painter::new(&app, size, &tex_approx);
        let exporter = Exporter::new(&app, size);

        Approximator {
            app: app.clone(),
            size,
            differ,
            tex_savestate,
            tex_approx,
            tex_target,
            painter,
            exporter,
        }
    }

    pub fn save(&self) {
        self.app
            .copy_texture(&self.tex_approx, &self.tex_savestate, self.size);
    }
    pub fn restore(&self) {
        self.app
            .copy_texture(&self.tex_savestate, &self.tex_approx, self.size);
    }

    pub async fn run(&mut self) {
        let mut current_diff = self.differ.run().await;
        let mut params = PaintUniforms {
            x: 128.0,
            y: 128.0,
            r: 0.5,
            g: 0.5,
            b: 0.5,
            radius: 64.0,
        };

        self.optimize_param(&mut current_diff, &mut params, "", |p| p.b += 0.1);

        self.exporter.run(&self.tex_approx, "a/approx.png").await;
    }

    pub fn optimize_param<F>(
        &self,
        current_diff: &mut u32,
        params: &mut PaintUniforms,
        label: &'static str,
        f: F,
    ) where
        F: Fn(&mut PaintUniforms) -> (),
    {
        let mut p = params.clone();
        loop {
            self.save();
            info!("apply '{label:?}'");
            f(&mut p);
            self.painter.run(p);
            let diff = pollster::block_on(self.differ.run());
            self.restore();
            if diff >= *current_diff {
                break;
            }
            *current_diff = diff;
        }
    }
}