use std::collections::HashMap;

pub const SPRITE_ATLAS_TOTAL_SLOTS: u32 = 128;
pub const SPRITE_ATLAS_SLOT_SIZE: (u32, u32) = (512, 512);

pub struct SpriteTextureAtlas {
    pub texture: wgpu::Texture,
    pub view: wgpu::TextureView,
    pub sampler: wgpu::Sampler,
    pub atlas_size: (u32, u32),
    pub slot_size: (u32, u32),
    pub slots_per_row: u32,
    pub total_slots: u32,
    occupied_slots: HashMap<u32, bool>,
}

impl SpriteTextureAtlas {
    pub fn new(device: &wgpu::Device, total_slots: u32, slot_size: (u32, u32)) -> Self {
        let slots_per_row = (total_slots as f32).sqrt().ceil() as u32;
        let atlas_size = (slots_per_row * slot_size.0, slots_per_row * slot_size.1);

        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("Sprite Texture Atlas"),
            size: wgpu::Extent3d {
                width: atlas_size.0,
                height: atlas_size.1,
                depth_or_array_layers: 1,
            },
            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: &[],
        });

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

        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            mipmap_filter: wgpu::FilterMode::Linear,
            ..Default::default()
        });

        Self {
            texture,
            view,
            sampler,
            atlas_size,
            slot_size,
            slots_per_row,
            total_slots,
            occupied_slots: HashMap::new(),
        }
    }

    pub fn upload_texture(
        &mut self,
        queue: &wgpu::Queue,
        slot: u32,
        rgba_data: &[u8],
        width: u32,
        height: u32,
    ) {
        if slot >= self.total_slots {
            return;
        }

        let row = slot / self.slots_per_row;
        let col = slot % self.slots_per_row;

        let x_offset = col * self.slot_size.0;
        let y_offset = row * self.slot_size.1;

        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &self.texture,
                mip_level: 0,
                origin: wgpu::Origin3d {
                    x: x_offset,
                    y: y_offset,
                    z: 0,
                },
                aspect: wgpu::TextureAspect::All,
            },
            rgba_data,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(width * 4),
                rows_per_image: Some(height),
            },
            wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
        );

        self.occupied_slots.insert(slot, true);
    }

    pub fn get_uv_transform(&self, slot: u32) -> (f32, f32, f32, f32) {
        let row = slot / self.slots_per_row;
        let col = slot % self.slots_per_row;

        let u_min = (col * self.slot_size.0) as f32 / self.atlas_size.0 as f32;
        let v_min = (row * self.slot_size.1) as f32 / self.atlas_size.1 as f32;
        let u_max = ((col + 1) * self.slot_size.0) as f32 / self.atlas_size.0 as f32;
        let v_max = ((row + 1) * self.slot_size.1) as f32 / self.atlas_size.1 as f32;

        (u_min, v_min, u_max, v_max)
    }
}
