use std::collections::HashMap;
use wgpu::util::DeviceExt;

#[cfg(not(target_arch = "wasm32"))]
use image::GenericImageView;

#[derive(Default)]
pub struct TextureCache {
    textures: HashMap<String, TextureEntry>,
    reference_counts: HashMap<String, usize>,
    pending_references: HashMap<String, usize>,
}

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

impl TextureCache {
    pub fn new() -> Self {
        Self {
            textures: HashMap::new(),
            reference_counts: HashMap::new(),
            pending_references: HashMap::new(),
        }
    }

    fn create_dummy_texture(device: &wgpu::Device, queue: &wgpu::Queue) -> wgpu::Texture {
        let width = 64;
        let height = 64;
        let checker_size = 8;

        let mut pixels = Vec::new();
        for y in 0..height {
            for x in 0..width {
                let checker_x = (x / checker_size) % 2;
                let checker_y = (y / checker_size) % 2;
                let is_purple = (checker_x + checker_y) % 2 == 0;

                if is_purple {
                    pixels.extend_from_slice(&[128, 0, 128, 255]);
                } else {
                    pixels.extend_from_slice(&[0, 0, 0, 255]);
                }
            }
        }

        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };

        device.create_texture_with_data(
            queue,
            &wgpu::TextureDescriptor {
                label: Some("Dummy Texture"),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::Rgba8UnormSrgb,
                usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
                view_formats: &[],
            },
            wgpu::util::TextureDataOrder::LayerMajor,
            &pixels,
        )
    }

    pub fn load_texture_from_raw_rgba(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        name: String,
        rgba_data: &[u8],
        width: u32,
        height: u32,
    ) -> Result<(), String> {
        if self.textures.contains_key(&name) {
            return Ok(());
        }

        let pending_refs = self.pending_references.remove(&name).unwrap_or(0);

        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };

        let texture = device.create_texture_with_data(
            queue,
            &wgpu::TextureDescriptor {
                label: Some(&name),
                size,
                mip_level_count: 1,
                sample_count: 1,
                dimension: wgpu::TextureDimension::D2,
                format: wgpu::TextureFormat::Rgba8UnormSrgb,
                usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
                view_formats: &[],
            },
            wgpu::util::TextureDataOrder::LayerMajor,
            rgba_data,
        );

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

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some(&format!("{} Sampler", name)),
            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::Linear,
            ..Default::default()
        });

        self.textures.insert(
            name.clone(),
            TextureEntry {
                texture,
                view,
                sampler,
            },
        );
        self.reference_counts.insert(name, 1 + pending_refs);

        Ok(())
    }

    pub fn add_reference(&mut self, name: &str) {
        if let Some(count) = self.reference_counts.get_mut(name) {
            *count += 1;
        } else {
            *self.pending_references.entry(name.to_string()).or_insert(0) += 1;
        }
    }

    pub fn remove_reference(&mut self, name: &str) {
        if let Some(count) = self.reference_counts.get_mut(name) {
            *count = count.saturating_sub(1);
        } else if let Some(count) = self.pending_references.get_mut(name) {
            *count = count.saturating_sub(1);
            if *count == 0 {
                self.pending_references.remove(name);
            }
        }
    }

    pub fn get_reference_count(&self, name: &str) -> usize {
        self.reference_counts.get(name).copied().unwrap_or(0)
    }

    pub fn remove_unused(&mut self) -> Vec<String> {
        let unused: Vec<String> = self
            .reference_counts
            .iter()
            .filter(|&(_, count)| *count == 0)
            .map(|(name, _)| name.clone())
            .collect();

        for name in &unused {
            self.textures.remove(name);
            self.reference_counts.remove(name);
        }

        unused
    }

    pub fn contains(&self, name: &str) -> bool {
        self.textures.contains_key(name)
    }

    #[cfg(not(target_arch = "wasm32"))]
    pub fn load_texture_from_bytes(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        name: String,
        bytes: &[u8],
    ) -> Result<(), String> {
        let img =
            image::load_from_memory(bytes).map_err(|e| format!("Failed to load image: {}", e))?;
        let rgba = img.to_rgba8();
        let dimensions = img.dimensions();

        self.load_texture_from_raw_rgba(device, queue, name, &rgba, dimensions.0, dimensions.1)
    }

    pub fn get_or_create_dummy(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        name: &str,
    ) -> &TextureEntry {
        if !self.textures.contains_key(name) {
            let texture = Self::create_dummy_texture(device, queue);
            let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
            let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
                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::Linear,
                ..Default::default()
            });
            self.textures.insert(
                name.to_string(),
                TextureEntry {
                    texture,
                    view,
                    sampler,
                },
            );
            self.reference_counts.insert(name.to_string(), 1);
        }
        self.textures.get(name).unwrap()
    }

    pub fn get(&self, name: &str) -> Option<&TextureEntry> {
        self.textures.get(name)
    }
}
