use anyhow::Result;
use exr::prelude::{ReadChannels, ReadLayers};
use half::f16;

pub struct DitherTextures
{
    pub texture_array: wgpu::Texture,
    pub view: wgpu::TextureView,
    pub sampler: wgpu::Sampler,
}

pub struct FlowmapTexture
{
    pub texture: wgpu::Texture,
    pub view: wgpu::TextureView,
    pub sampler: wgpu::Sampler,
}

impl DitherTextures
{
    pub fn load_octaves(device: &wgpu::Device, queue: &wgpu::Queue) -> Result<Self>
    {
        let octave_paths = [
            "textures/dither/octave_0.png",
            "textures/dither/octave_1.png",
            "textures/dither/octave_2.png",
            "textures/dither/octave_3.png",
        ];

        let mut images = Vec::new();
        let mut texture_size = 0;

        for path in &octave_paths
        {
            let img = image::open(path)?.to_luma8();
            let (width, height) = img.dimensions();

            if texture_size == 0
            {
                texture_size = width;
            }
            else if width != texture_size || height != texture_size
            {
                return Err(anyhow::anyhow!(
                    "All dither textures must be the same size. Expected {}x{}, got {}x{}",
                    texture_size,
                    texture_size,
                    width,
                    height
                ));
            }

            if width != height
            {
                return Err(anyhow::anyhow!(
                    "Dither textures must be square. Got {}x{}",
                    width,
                    height
                ));
            }

            images.push(img);
        }

        let texture_array = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("Dither Texture Array"),
            size: wgpu::Extent3d {
                width: texture_size,
                height: texture_size,
                depth_or_array_layers: 4,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::R8Unorm,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });

        for (i, img) in images.iter().enumerate()
        {
            queue.write_texture(
                wgpu::TexelCopyTextureInfo {
                    texture: &texture_array,
                    mip_level: 0,
                    origin: wgpu::Origin3d {
                        x: 0,
                        y: 0,
                        z: i as u32,
                    },
                    aspect: wgpu::TextureAspect::All,
                },
                img.as_raw(),
                wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(texture_size),
                    rows_per_image: Some(texture_size),
                },
                wgpu::Extent3d {
                    width: texture_size,
                    height: texture_size,
                    depth_or_array_layers: 1,
                },
            );
        }

        let view = texture_array.create_view(&wgpu::TextureViewDescriptor {
            label: Some("Dither Texture Array View"),
            dimension: Some(wgpu::TextureViewDimension::D2Array),
            ..Default::default()
        });

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Dither Sampler"),
            address_mode_u: wgpu::AddressMode::Repeat,
            address_mode_v: wgpu::AddressMode::Repeat,
            address_mode_w: wgpu::AddressMode::Repeat,
            mag_filter: wgpu::FilterMode::Nearest,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });

        Ok(Self {
            texture_array,
            view,
            sampler,
        })
    }
}

impl FlowmapTexture
{
    pub fn load(device: &wgpu::Device, queue: &wgpu::Queue, path: &str) -> Result<Self>
    {
        let image = exr::prelude::read()
            .no_deep_data()
            .largest_resolution_level()
            .all_channels()
            .all_layers()
            .all_attributes()
            .from_file(path)?;

        let layer = &image.layer_data[0];
        let width = layer.size.width();
        let height = layer.size.height();

        if width != height
        {
            return Err(anyhow::anyhow!(
                "Flowmap texture must be square. Got {}x{}",
                width,
                height
            ));
        }

        let mut rgba_data: Vec<f32> = vec![1.0; width * height * 4];

        for channel in &layer.channel_data.list
        {
            let channel_name = channel.name.to_string();
            let values: Vec<f32> = channel.sample_data.values_as_f32().collect();

            let target_channel = match channel_name.as_str()
            {
                "R" => 0,
                "G" => 1,
                "B" => 2,
                "A" => 3,
                _ => continue,
            };

            for (i, &value) in values.iter().enumerate()
            {
                rgba_data[i * 4 + target_channel] = value;
            }
        }

        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("Flowmap Texture"),
            size: wgpu::Extent3d {
                width: width as u32,
                height: height as u32,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba16Float,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });

        let rgba_data_f16: Vec<u16> = rgba_data
            .iter()
            .map(|&f| f16::from_f32(f).to_bits())
            .collect();

        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            bytemuck::cast_slice(&rgba_data_f16),
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(8 * width as u32),
                rows_per_image: Some(height as u32),
            },
            wgpu::Extent3d {
                width: width as u32,
                height: height as u32,
                depth_or_array_layers: 1,
            },
        );

        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("Flowmap Sampler"),
            address_mode_u: wgpu::AddressMode::Repeat,
            address_mode_v: wgpu::AddressMode::Repeat,
            address_mode_w: wgpu::AddressMode::Repeat,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });

        Ok(Self {
            texture,
            view,
            sampler,
        })
    }
}