use crate::ecs::asset_id::TextureId;
use crate::ecs::generational_registry::GenerationalRegistry;
use std::collections::HashMap;
use wgpu::util::DeviceExt;

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

pub struct TextureCache {
    pub registry: GenerationalRegistry<TextureEntry>,
    pub pending_references: HashMap<String, usize>,
}

impl Default for TextureCache {
    fn default() -> Self {
        Self {
            registry: GenerationalRegistry::new(),
            pending_references: HashMap::new(),
        }
    }
}

impl TextureCache {
    pub fn load_texture_from_raw_rgba(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        name: String,
        rgba_data: &[u8],
        dimensions: (u32, u32),
    ) -> Result<TextureId, String> {
        texture_cache_load_from_raw_rgba(self, device, queue, name, rgba_data, dimensions)
    }

    pub fn get(&self, name: &str) -> Option<&TextureEntry> {
        let index = self.registry.name_to_index.get(name)?;
        self.registry.entries[*index as usize].as_ref()
    }
}

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

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::Rgba8Unorm,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        },
        wgpu::util::TextureDataOrder::LayerMajor,
        &pixels,
    )
}

pub fn texture_cache_load_from_raw_rgba(
    cache: &mut TextureCache,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: String,
    rgba_data: &[u8],
    dimensions: (u32, u32),
) -> Result<TextureId, String> {
    texture_cache_load_from_raw_rgba_with_format(
        cache,
        device,
        queue,
        name,
        rgba_data,
        dimensions,
        wgpu::TextureFormat::Rgba8Unorm,
    )
}

pub fn texture_cache_load_from_raw_rgba_linear(
    cache: &mut TextureCache,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: String,
    rgba_data: &[u8],
    dimensions: (u32, u32),
) -> Result<TextureId, String> {
    texture_cache_load_from_raw_rgba_with_format(
        cache,
        device,
        queue,
        name,
        rgba_data,
        dimensions,
        wgpu::TextureFormat::Rgba8Unorm,
    )
}

pub fn texture_cache_load_from_raw_rgba_with_format(
    cache: &mut TextureCache,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: String,
    rgba_data: &[u8],
    dimensions: (u32, u32),
    format: wgpu::TextureFormat,
) -> Result<TextureId, String> {
    if let Some((index, generation)) = cache.registry.lookup_index(&name) {
        return Ok(TextureId { index, generation });
    }

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

    let size = wgpu::Extent3d {
        width: dimensions.0,
        height: dimensions.1,
        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,
            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::Nearest,
        min_filter: wgpu::FilterMode::Nearest,
        mipmap_filter: wgpu::FilterMode::Nearest,
        ..Default::default()
    });

    let entry = TextureEntry {
        texture,
        view,
        sampler,
    };

    let (index, generation) = cache.registry.insert(name, entry);
    cache.registry.reference_counts[index as usize] = 1 + pending_refs;

    Ok(TextureId { index, generation })
}

pub fn texture_cache_lookup_id(cache: &TextureCache, name: &str) -> Option<TextureId> {
    cache
        .registry
        .lookup_index(name)
        .map(|(index, generation)| TextureId { index, generation })
}

pub fn texture_cache_add_reference(cache: &mut TextureCache, name: &str) {
    if let Some(&index) = cache.registry.name_to_index.get(name) {
        cache.registry.add_reference(index);
    } else {
        *cache
            .pending_references
            .entry(name.to_string())
            .or_insert(0) += 1;
    }
}

pub fn texture_cache_remove_reference(cache: &mut TextureCache, name: &str) {
    if let Some(&index) = cache.registry.name_to_index.get(name) {
        cache.registry.remove_reference(index);
    } else if let Some(count) = cache.pending_references.get_mut(name) {
        *count = count.saturating_sub(1);
        if *count == 0 {
            cache.pending_references.remove(name);
        }
    }
}

pub fn texture_cache_remove_unused(cache: &mut TextureCache) -> Vec<String> {
    cache.registry.remove_unused()
}

#[cfg(not(target_arch = "wasm32"))]
pub fn texture_cache_load_from_bytes(
    cache: &mut TextureCache,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: String,
    bytes: &[u8],
) -> Result<TextureId, 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();

    texture_cache_load_from_raw_rgba(cache, device, queue, name, &rgba, dimensions)
}

pub fn texture_cache_ensure_dummy<'a>(
    cache: &'a mut TextureCache,
    device: &wgpu::Device,
    queue: &wgpu::Queue,
    name: &str,
) -> &'a TextureEntry {
    if !cache.registry.name_to_index.contains_key(name) {
        let texture = 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::Nearest,
            min_filter: wgpu::FilterMode::Nearest,
            mipmap_filter: wgpu::FilterMode::Nearest,
            ..Default::default()
        });
        let entry = TextureEntry {
            texture,
            view,
            sampler,
        };
        let (index, _) = cache.registry.insert(name.to_string(), entry);
        cache.registry.reference_counts[index as usize] = 1;
    }
    let index = cache.registry.name_to_index.get(name).unwrap();
    cache.registry.entries[*index as usize].as_ref().unwrap()
}
