use std::collections::HashMap;

use naga::proc::Layouter;

/// A borrowed, typed view into a region of the uniform buffer.
/// Supports indexing into compound types and writing scalar values.
pub struct UniformRef<'a> {
    pub module: &'a naga::Module,
    pub layouter: &'a Layouter,
    pub ty: naga::Handle<naga::Type>,
    pub data: &'a mut [u8],
}

#[derive(Debug, thiserror::Error)]
pub enum UniformError {
    #[error("can't be indexed (further)")]
    NotIndexable,
    #[error("index '{0}' out of bounds")]
    IndexOutOfBounds(usize),
    #[error("member '{0}' not found")]
    MemberNotFound(String),
    #[error("mismatched value types")]
    TypeMismatch,
    #[error("invalid number of OSC values")]
    SizeMismatch,
}

#[derive(Debug, thiserror::Error)]
pub enum BindError {
    #[error("binding point '{0}' not found")]
    NotFound(String),
    #[error("bindable resource '{0}' not found")]
    ResourceNotFound(String),
}

impl<'a> UniformRef<'a> {
    /// Navigate by path component: struct member by name, or swizzle/numeric index otherwise.
    pub fn field(self, name: &str) -> Result<UniformRef<'a>, UniformError> {
        let ty_inner = &self.module.types[self.ty].inner;
        if let naga::TypeInner::Struct { ref members, .. } = *ty_inner {
            for member in members {
                if member.name.as_deref() == Some(name) {
                    let offset = member.offset as usize;
                    let size = self.layouter[member.ty].size as usize;
                    return Ok(UniformRef {
                        module: self.module,
                        layouter: self.layouter,
                        ty: member.ty,
                        data: &mut self.data[offset..offset + size],
                    });
                }
            }
            return Err(UniformError::MemberNotFound(name.to_string()));
        }

        let i = match name {
            "x" | "r" | "s" => 0,
            "y" | "g" | "t" => 1,
            "z" | "b" | "p" => 2,
            "w" | "a" | "q" => 3,
            s => s.parse::<usize>().map_err(|_| UniformError::NotIndexable)?,
        };
        self.index(i)
    }

    /// Index into a vector, matrix, or array by position.
    pub fn index(self, i: usize) -> Result<UniformRef<'a>, UniformError> {
        let ty_inner = &self.module.types[self.ty].inner;
        match *ty_inner {
            naga::TypeInner::Vector { size, scalar } => {
                if i >= size as usize {
                    return Err(UniformError::IndexOutOfBounds(i));
                }
                let w = scalar.width as usize;
                let scalar_ty = find_or_expect_scalar_type(self.module, scalar);
                Ok(UniformRef {
                    module: self.module,
                    layouter: self.layouter,
                    ty: scalar_ty,
                    data: &mut self.data[i * w..][..w],
                })
            }
            naga::TypeInner::Matrix {
                columns,
                rows,
                scalar,
            } => {
                if i >= columns as usize {
                    return Err(UniformError::IndexOutOfBounds(i));
                }
                let col_ty = find_or_expect_vector_type(self.module, rows, scalar);
                let col_layout = self.layouter[col_ty];
                let stride = col_layout.to_stride() as usize;
                let size = col_layout.size as usize;
                Ok(UniformRef {
                    module: self.module,
                    layouter: self.layouter,
                    ty: col_ty,
                    data: &mut self.data[i * stride..][..size],
                })
            }
            naga::TypeInner::Array { base, size, stride } => {
                let len = match size {
                    naga::ArraySize::Constant(n) => n.get() as usize,
                    _ => return Err(UniformError::NotIndexable),
                };
                if i >= len {
                    return Err(UniformError::IndexOutOfBounds(i));
                }
                let elem_size = self.layouter[base].size as usize;
                Ok(UniformRef {
                    module: self.module,
                    layouter: self.layouter,
                    ty: base,
                    data: &mut self.data[i * stride as usize..][..elem_size],
                })
            }
            _ => Err(UniformError::NotIndexable),
        }
    }

    /// Get the leaf scalar type of this uniform.
    pub fn leaf_scalar(&self) -> Result<naga::Scalar, UniformError> {
        leaf_scalar_of(self.module, self.ty)
    }

    /// Count the total number of leaf scalars in this type.
    fn leaf_count(&self) -> Result<usize, UniformError> {
        let ty_inner = &self.module.types[self.ty].inner;
        match *ty_inner {
            naga::TypeInner::Scalar(_) => Ok(1),
            naga::TypeInner::Vector { size, .. } => Ok(size as usize),
            naga::TypeInner::Matrix { columns, rows, .. } => Ok(columns as usize * rows as usize),
            naga::TypeInner::Array { base, size, .. } => {
                let len = match size {
                    naga::ArraySize::Constant(n) => n.get() as usize,
                    _ => return Err(UniformError::NotIndexable),
                };
                let inner = &self.module.types[base].inner;
                let per_elem = match *inner {
                    naga::TypeInner::Scalar(_) => 1,
                    naga::TypeInner::Vector { size, .. } => size as usize,
                    naga::TypeInner::Matrix { columns, rows, .. } => {
                        columns as usize * rows as usize
                    }
                    _ => return Err(UniformError::NotIndexable),
                };
                Ok(len * per_elem)
            }
            _ => Err(UniformError::NotIndexable),
        }
    }

    pub fn set_f32(self, values: &[f32]) -> Result<(), UniformError> {
        self.set_scalars(values, naga::ScalarKind::Float, 4)
    }

    pub fn set_i32(self, values: &[i32]) -> Result<(), UniformError> {
        self.set_scalars(values, naga::ScalarKind::Sint, 4)
    }

    pub fn set_u32(self, values: &[u32]) -> Result<(), UniformError> {
        self.set_scalars(values, naga::ScalarKind::Uint, 4)
    }

    fn set_scalars<T: Copy + bytemuck::NoUninit>(
        self,
        values: &[T],
        expected_kind: naga::ScalarKind,
        expected_width: u8,
    ) -> Result<(), UniformError> {
        let expected = self.leaf_count()?;
        if values.len() != expected {
            return Err(UniformError::SizeMismatch);
        }
        self.write_recursive(values, expected_kind, expected_width, &mut 0)
    }

    /// `val_idx` tracks position in the flat values array.
    fn write_recursive<T: Copy + bytemuck::NoUninit>(
        self,
        values: &[T],
        expected_kind: naga::ScalarKind,
        expected_width: u8,
        val_idx: &mut usize,
    ) -> Result<(), UniformError> {
        let ty_inner = &self.module.types[self.ty].inner;
        match *ty_inner {
            naga::TypeInner::Scalar(scalar) => {
                if scalar.kind != expected_kind || scalar.width != expected_width {
                    return Err(UniformError::TypeMismatch);
                }
                let bytes = bytemuck::bytes_of(&values[*val_idx]);
                self.data[..bytes.len()].copy_from_slice(bytes);
                *val_idx += 1;
                Ok(())
            }
            naga::TypeInner::Vector { size, scalar } => {
                if scalar.kind != expected_kind || scalar.width != expected_width {
                    return Err(UniformError::TypeMismatch);
                }
                let count = size as usize;
                let w = scalar.width as usize;
                for c in 0..count {
                    let offset = c * w;
                    let bytes = bytemuck::bytes_of(&values[*val_idx]);
                    self.data[offset..offset + w].copy_from_slice(bytes);
                    *val_idx += 1;
                }
                Ok(())
            }
            naga::TypeInner::Matrix {
                columns,
                rows,
                scalar,
            } => {
                if scalar.kind != expected_kind || scalar.width != expected_width {
                    return Err(UniformError::TypeMismatch);
                }
                let col_ty = find_or_expect_vector_type(self.module, rows, scalar);
                let col_stride = self.layouter[col_ty].to_stride() as usize;
                let w = scalar.width as usize;
                let row_count = rows as usize;
                for col in 0..columns as usize {
                    let col_offset = col * col_stride;
                    for row in 0..row_count {
                        let offset = col_offset + row * w;
                        let bytes = bytemuck::bytes_of(&values[*val_idx]);
                        self.data[offset..offset + w].copy_from_slice(bytes);
                        *val_idx += 1;
                    }
                }
                Ok(())
            }
            naga::TypeInner::Array { base, size, stride } => {
                let len = match size {
                    naga::ArraySize::Constant(n) => n.get() as usize,
                    _ => return Err(UniformError::NotIndexable),
                };
                let elem_size = self.layouter[base].size as usize;
                let stride = stride as usize;
                for i in 0..len {
                    let offset = i * stride;
                    let elem_ref = UniformRef {
                        module: self.module,
                        layouter: self.layouter,
                        ty: base,
                        data: &mut self.data[offset..offset + elem_size],
                    };
                    elem_ref.write_recursive(values, expected_kind, expected_width, val_idx)?;
                }
                Ok(())
            }
            _ => Err(UniformError::TypeMismatch),
        }
    }
}

fn leaf_scalar_of(
    module: &naga::Module,
    ty: naga::Handle<naga::Type>,
) -> Result<naga::Scalar, UniformError> {
    match module.types[ty].inner {
        naga::TypeInner::Scalar(scalar) => Ok(scalar),
        naga::TypeInner::Vector { scalar, .. } => Ok(scalar),
        naga::TypeInner::Matrix { scalar, .. } => Ok(scalar),
        naga::TypeInner::Array { base, .. } => leaf_scalar_of(module, base),
        _ => Err(UniformError::TypeMismatch),
    }
}

/// Ensure that all component sub-types (scalars for vectors, vectors for matrix columns)
/// exist as standalone entries in the module's type arena.
/// naga doesn't always create these when they're only used implicitly.
fn ensure_component_types(module: &mut naga::Module) {
    let mut needed_scalars = Vec::new();
    let mut needed_vectors = Vec::new();

    for (_handle, ty) in module.types.iter() {
        match ty.inner {
            naga::TypeInner::Vector { scalar, .. } => {
                needed_scalars.push(scalar);
            }
            naga::TypeInner::Matrix { rows, scalar, .. } => {
                needed_scalars.push(scalar);
                needed_vectors.push((rows, scalar));
            }
            _ => {}
        }
    }

    for scalar in needed_scalars {
        let exists = module
            .types
            .iter()
            .any(|(_, ty)| ty.inner == naga::TypeInner::Scalar(scalar));
        if !exists {
            module.types.insert(
                naga::Type {
                    name: None,
                    inner: naga::TypeInner::Scalar(scalar),
                },
                naga::Span::UNDEFINED,
            );
        }
    }

    for (size, scalar) in needed_vectors {
        let exists = module
            .types
            .iter()
            .any(|(_, ty)| ty.inner == naga::TypeInner::Vector { size, scalar });
        if !exists {
            module.types.insert(
                naga::Type {
                    name: None,
                    inner: naga::TypeInner::Vector { size, scalar },
                },
                naga::Span::UNDEFINED,
            );
        }
    }
}

fn find_or_expect_scalar_type(
    module: &naga::Module,
    scalar: naga::Scalar,
) -> naga::Handle<naga::Type> {
    for (handle, ty) in module.types.iter() {
        if ty.inner == naga::TypeInner::Scalar(scalar) {
            return handle;
        }
    }
    panic!("scalar type {:?} not found in module type arena", scalar);
}

fn find_or_expect_vector_type(
    module: &naga::Module,
    size: naga::VectorSize,
    scalar: naga::Scalar,
) -> naga::Handle<naga::Type> {
    for (handle, ty) in module.types.iter() {
        if ty.inner == (naga::TypeInner::Vector { size, scalar }) {
            return handle;
        }
    }
    panic!(
        "vector type vec{}<{:?}> not found in module type arena",
        size as u8, scalar
    );
}

/// Compare two types from potentially different naga Modules for structural equivalence.
pub fn types_compatible(
    a_mod: &naga::Module,
    a_ty: naga::Handle<naga::Type>,
    b_mod: &naga::Module,
    b_ty: naga::Handle<naga::Type>,
) -> bool {
    let a = &a_mod.types[a_ty].inner;
    let b = &b_mod.types[b_ty].inner;

    // Uniform types (Scalar, Vector, Matrix) contain no handles or pointers,
    // so TypeInner's derived PartialEq works directly. Only Array and Struct
    // need recursion because they contain Handle<Type> from different arenas.
    match (a, b) {
        (
            naga::TypeInner::Array {
                base: ab,
                size: asize,
                stride: astride,
            },
            naga::TypeInner::Array {
                base: bb,
                size: bsize,
                stride: bstride,
            },
        ) => asize == bsize && astride == bstride && types_compatible(a_mod, *ab, b_mod, *bb),
        (
            naga::TypeInner::Struct {
                members: am,
                span: asp,
            },
            naga::TypeInner::Struct {
                members: bm,
                span: bsp,
            },
        ) => {
            asp == bsp
                && am.len() == bm.len()
                && am.iter().zip(bm.iter()).all(|(a, b)| {
                    a.offset == b.offset && types_compatible(a_mod, a.ty, b_mod, b.ty)
                })
        }
        _ => a == b,
    }
}

/// A pooled TSV texture resource.
pub struct TextureResource {
    id: String,
    tsv_name: Option<String>,
    pub tsv_registered: bool,
    texture: wgpu::Texture,
    view: wgpu::TextureView,
    format: wgpu::TextureFormat,
    width: u32,
    height: u32,
    depth_or_array_layers: u32,
    view_dimension: wgpu::TextureViewDimension,
}

impl TextureResource {
    pub fn new(device: &wgpu::Device, id: String, tsv_name: String) -> Self {
        let format = wgpu::TextureFormat::Rgba8Unorm;
        let view_dimension = wgpu::TextureViewDimension::D2;
        let dimension = view_dimension_to_texture_dimension(view_dimension);
        let (texture, view) =
            create_input_texture(device, &id, 1, 1, 1, dimension, view_dimension, format);
        Self {
            id,
            tsv_name: Some(tsv_name),
            tsv_registered: false,
            texture,
            view,
            format,
            width: 1,
            height: 1,
            depth_or_array_layers: 1,
            view_dimension,
        }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn tsv_name(&self) -> Option<&str> {
        self.tsv_name.as_deref()
    }

    pub fn set_tsv_name(&mut self, name: String) {
        if self.tsv_name.as_deref() != Some(&name) {
            self.tsv_name = Some(name);
            self.tsv_registered = false;
        }
    }

    pub fn texture(&self) -> &wgpu::Texture {
        &self.texture
    }

    pub fn view(&self) -> &wgpu::TextureView {
        &self.view
    }

    pub fn view_dimension(&self) -> wgpu::TextureViewDimension {
        self.view_dimension
    }

    pub fn width(&self) -> u32 {
        self.width
    }

    pub fn height(&self) -> u32 {
        self.height
    }

    /// Recreate the texture if size or format changed. Returns true if recreated.
    pub fn resize(
        &mut self,
        device: &wgpu::Device,
        width: u32,
        height: u32,
        depth_or_array_layers: u32,
        view_dimension: wgpu::TextureViewDimension,
        format: wgpu::TextureFormat,
    ) -> bool {
        if self.width == width
            && self.height == height
            && self.depth_or_array_layers == depth_or_array_layers
            && self.view_dimension == view_dimension
            && self.format == format
        {
            return false;
        }
        let dimension = view_dimension_to_texture_dimension(view_dimension);
        let (texture, view) = create_input_texture(
            device,
            &self.id,
            width,
            height,
            depth_or_array_layers,
            dimension,
            view_dimension,
            format,
        );
        self.texture = texture;
        self.view = view;
        self.format = format;
        self.width = width;
        self.height = height;
        self.depth_or_array_layers = depth_or_array_layers;
        self.view_dimension = view_dimension;
        true
    }
}

fn view_dimension_to_texture_dimension(
    view_dim: wgpu::TextureViewDimension,
) -> wgpu::TextureDimension {
    match view_dim {
        wgpu::TextureViewDimension::D1 => wgpu::TextureDimension::D1,
        wgpu::TextureViewDimension::D3 => wgpu::TextureDimension::D3,
        // D2, D2Array, Cube, CubeArray all use D2 textures
        _ => wgpu::TextureDimension::D2,
    }
}

fn create_input_texture(
    device: &wgpu::Device,
    label: &str,
    width: u32,
    height: u32,
    depth_or_array_layers: u32,
    dimension: wgpu::TextureDimension,
    view_dimension: wgpu::TextureViewDimension,
    format: wgpu::TextureFormat,
) -> (wgpu::Texture, wgpu::TextureView) {
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some(label),
        size: wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension,
        format,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });
    let view = texture.create_view(&wgpu::TextureViewDescriptor {
        dimension: Some(view_dimension),
        ..Default::default()
    });
    (texture, view)
}

pub struct SamplerResource {
    id: String,
    sampler: wgpu::Sampler,
}

impl SamplerResource {
    pub fn new(
        device: &wgpu::Device,
        id: String,
        filter: wgpu::FilterMode,
        address_mode: wgpu::AddressMode,
    ) -> Self {
        Self {
            sampler: create_sampler(device, &id, filter, address_mode),
            id,
        }
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn sampler(&self) -> &wgpu::Sampler {
        &self.sampler
    }

    /// Recreate the sampler with new filter and address modes.
    pub fn configure(
        &mut self,
        device: &wgpu::Device,
        filter: wgpu::FilterMode,
        address_mode: wgpu::AddressMode,
    ) {
        self.sampler = create_sampler(device, &self.id, filter, address_mode);
    }
}

fn create_sampler(
    device: &wgpu::Device,
    label: &str,
    filter: wgpu::FilterMode,
    address_mode: wgpu::AddressMode,
) -> wgpu::Sampler {
    device.create_sampler(&wgpu::SamplerDescriptor {
        label: Some(label),
        mag_filter: filter,
        min_filter: filter,
        mipmap_filter: match filter {
            wgpu::FilterMode::Linear => wgpu::MipmapFilterMode::Linear,
            wgpu::FilterMode::Nearest => wgpu::MipmapFilterMode::Nearest,
        },
        address_mode_u: address_mode,
        address_mode_v: address_mode,
        address_mode_w: address_mode,
        ..Default::default()
    })
}

struct TextureBindingSlot {
    name: String,
    binding: u32,
    view_dimension: wgpu::TextureViewDimension,
    sample_type: wgpu::TextureSampleType,
    multisampled: bool,
    resource_id: Option<String>,
}

struct SamplerBindingSlot {
    name: String,
    binding: u32,
    binding_type: wgpu::SamplerBindingType,
    resource_id: Option<String>,
}

fn image_view_dimension(dim: naga::ImageDimension, arrayed: bool) -> wgpu::TextureViewDimension {
    match (dim, arrayed) {
        (naga::ImageDimension::D1, false) => wgpu::TextureViewDimension::D1,
        (naga::ImageDimension::D2, false) => wgpu::TextureViewDimension::D2,
        (naga::ImageDimension::D2, true) => wgpu::TextureViewDimension::D2Array,
        (naga::ImageDimension::D3, false) => wgpu::TextureViewDimension::D3,
        (naga::ImageDimension::Cube, false) => wgpu::TextureViewDimension::Cube,
        (naga::ImageDimension::Cube, true) => wgpu::TextureViewDimension::CubeArray,
        _ => wgpu::TextureViewDimension::D2,
    }
}

fn image_sample_type(class: naga::ImageClass) -> wgpu::TextureSampleType {
    match class {
        naga::ImageClass::Sampled { kind, .. } => match kind {
            naga::ScalarKind::Float => wgpu::TextureSampleType::Float { filterable: true },
            naga::ScalarKind::Sint => wgpu::TextureSampleType::Sint,
            naga::ScalarKind::Uint => wgpu::TextureSampleType::Uint,
            _ => wgpu::TextureSampleType::Float { filterable: true },
        },
        naga::ImageClass::Depth { .. } => wgpu::TextureSampleType::Depth,
        naga::ImageClass::Storage { .. } | naga::ImageClass::External => {
            wgpu::TextureSampleType::Float { filterable: false }
        }
    }
}

fn image_multisampled(class: naga::ImageClass) -> bool {
    matches!(
        class,
        naga::ImageClass::Sampled { multi: true, .. } | naga::ImageClass::Depth { multi: true }
    )
}

/// Manages uniform buffer data and texture/sampler bindings for all globals in a shader.
/// Also owns the pooled texture and sampler resources.
pub struct UniformCache {
    pub module: naga::Module,
    pub layouter: Layouter,
    uniforms: HashMap<String, UniformMember>,
    buffers: Vec<BufferState>,
    texture_pool: HashMap<String, TextureResource>,
    sampler_pool: HashMap<String, SamplerResource>,
    texture_slots: Vec<TextureBindingSlot>,
    sampler_slots: Vec<SamplerBindingSlot>,
    pub bind_group_layout: Option<wgpu::BindGroupLayout>,
    pub bind_group: Option<wgpu::BindGroup>,
    dirty: bool,
}

struct UniformMember {
    buffer_idx: usize,
    ty: naga::Handle<naga::Type>,
    offset: usize,
    size: usize,
}

struct BufferState {
    binding: u32,
    data: Vec<u8>,
    gpu_buffer: Option<wgpu::Buffer>,
}

impl Default for UniformCache {
    fn default() -> Self {
        Self::new()
    }
}

impl UniformCache {
    pub fn new() -> Self {
        Self {
            module: naga::Module::default(),
            layouter: Layouter::default(),
            uniforms: HashMap::new(),
            buffers: Vec::new(),
            texture_pool: HashMap::new(),
            sampler_pool: HashMap::new(),
            texture_slots: Vec::new(),
            sampler_slots: Vec::new(),
            bind_group_layout: None,
            bind_group: None,
            dirty: false,
        }
    }

    /// Rebuild from a new shader module, transferring compatible uniform values
    /// and texture associations.
    pub fn refresh(&mut self, mut new_module: naga::Module, device: &wgpu::Device) {
        ensure_component_types(&mut new_module);
        let mut new_layouter = Layouter::default();
        new_layouter.update(new_module.to_ctx()).unwrap();

        let old_module = std::mem::take(&mut self.module);
        let old_uniforms = std::mem::take(&mut self.uniforms);
        let old_buffers = std::mem::take(&mut self.buffers);
        let old_textures = std::mem::take(&mut self.texture_slots);
        let old_samplers = std::mem::take(&mut self.sampler_slots);

        let mut new_uniforms: HashMap<String, UniformMember> = HashMap::new();
        let mut new_buffers: Vec<BufferState> = Vec::new();
        let mut new_textures: Vec<TextureBindingSlot> = Vec::new();
        let mut new_samplers: Vec<SamplerBindingSlot> = Vec::new();

        for (_handle, var) in new_module.global_variables.iter() {
            let binding = match &var.binding {
                Some(b) => b,
                None => continue,
            };

            if let Some(ref name) = var.name {
                log::info!("found uniform '{name}'");
            }

            match var.space {
                naga::AddressSpace::Uniform => {
                    let layout = new_layouter[var.ty];
                    let buffer_idx = new_buffers.len();

                    let data = vec![0u8; layout.size as usize];
                    new_buffers.push(BufferState {
                        binding: binding.binding,
                        data,
                        gpu_buffer: None,
                    });

                    if let Some(ref name) = var.name {
                        new_uniforms.insert(
                            name.clone(),
                            UniformMember {
                                buffer_idx,
                                ty: var.ty,
                                offset: 0,
                                size: layout.size as usize,
                            },
                        );
                    }
                }
                naga::AddressSpace::Handle => {
                    let ty_inner = &new_module.types[var.ty].inner;
                    match *ty_inner {
                        naga::TypeInner::Image {
                            dim,
                            arrayed,
                            class,
                        } => {
                            let name = var.name.clone().unwrap_or_default();
                            let view_dimension = image_view_dimension(dim, arrayed);
                            let sample_type = image_sample_type(class);
                            let multisampled = image_multisampled(class);

                            let mut slot = TextureBindingSlot {
                                name: name.clone(),
                                binding: binding.binding,
                                view_dimension,
                                sample_type,
                                multisampled,
                                resource_id: None,
                            };

                            if let Some(old) = old_textures.iter().find(|t| t.name == name) {
                                if old.view_dimension == slot.view_dimension
                                    && old.sample_type == slot.sample_type
                                    && old.multisampled == slot.multisampled
                                {
                                    slot.resource_id = old.resource_id.clone();
                                }
                            }

                            new_textures.push(slot);
                        }
                        naga::TypeInner::Sampler { comparison } => {
                            let binding_type = if comparison {
                                wgpu::SamplerBindingType::Comparison
                            } else {
                                wgpu::SamplerBindingType::Filtering
                            };
                            let name = var.name.clone().unwrap_or_default();
                            let mut slot = SamplerBindingSlot {
                                name,
                                binding: binding.binding,
                                binding_type,
                                resource_id: None,
                            };
                            if let Some(old) = old_samplers.iter().find(|s| {
                                s.name == slot.name && s.binding_type == slot.binding_type
                            }) {
                                slot.resource_id = old.resource_id.clone();
                            }
                            new_samplers.push(slot);
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }

        // Transfer compatible buffer data from old shader
        for (name, new_member) in &new_uniforms {
            if let Some(old_member) = old_uniforms.get(name) {
                if old_member.size == new_member.size
                    && types_compatible(&old_module, old_member.ty, &new_module, new_member.ty)
                {
                    let old_data = &old_buffers[old_member.buffer_idx].data
                        [old_member.offset..old_member.offset + old_member.size];
                    new_buffers[new_member.buffer_idx].data
                        [new_member.offset..new_member.offset + new_member.size]
                        .copy_from_slice(old_data);
                }
            }
        }

        // Create GPU buffers
        for buf in &mut new_buffers {
            let gpu_buffer = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some(&format!("uniform@{}", buf.binding)),
                size: buf.data.len() as u64,
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
            buf.gpu_buffer = Some(gpu_buffer);
        }

        self.module = new_module;
        self.layouter = new_layouter;
        self.uniforms = new_uniforms;
        self.buffers = new_buffers;
        self.texture_slots = new_textures;
        self.sampler_slots = new_samplers;
        self.dirty = true;
    }

    // --- Resource pool operations ---

    pub fn create_texture(&mut self, device: &wgpu::Device, id: &str, tsv_name: &str) {
        log::info!("creating texture {id} (tsv {tsv_name})");
        self.texture_pool.insert(
            id.to_string(),
            TextureResource::new(device, id.to_string(), tsv_name.to_string()),
        );
        self.build_bind_group(device);
    }

    pub fn destroy_texture(&mut self, device: &wgpu::Device, id: &str) {
        log::info!("destroying texture {id}");
        let removed = self.texture_pool.remove(id).is_some();
        let unbound = self.unbind_texture_resource(id);
        if removed || unbound {
            self.build_bind_group(device);
        }
    }

    pub fn textures_mut(&mut self) -> impl Iterator<Item = &mut TextureResource> {
        self.texture_pool.values_mut()
    }

    pub fn create_sampler(
        &mut self,
        device: &wgpu::Device,
        id: &str,
        filter: wgpu::FilterMode,
        address_mode: wgpu::AddressMode,
    ) {
        log::info!("creating sampler {id} ({filter:?} {address_mode:?})");
        self.sampler_pool.insert(
            id.to_string(),
            SamplerResource::new(device, id.to_string(), filter, address_mode),
        );
        self.build_bind_group(device);
    }

    pub fn destroy_sampler(&mut self, device: &wgpu::Device, id: &str) {
        log::info!("destroying sampler {id}");
        let removed = self.sampler_pool.remove(id).is_some();
        let unbound = self.unbind_sampler_resource(id);
        if removed || unbound {
            self.build_bind_group(device);
        }
    }

    pub fn bind_texture(
        &mut self,
        device: &wgpu::Device,
        slot_name: &str,
        resource_id: &str,
    ) -> Result<(), BindError> {
        if !self.texture_pool.contains_key(resource_id) {
            return Err(BindError::ResourceNotFound(resource_id.to_string()));
        }
        let slot = self
            .texture_slots
            .iter_mut()
            .find(|t| t.name == slot_name)
            .ok_or(BindError::NotFound(slot_name.to_string()))?;
        slot.resource_id = Some(resource_id.to_string());
        self.build_bind_group(device);
        log::info!("bound texture {resource_id} to slot {slot_name}");
        Ok(())
    }

    pub fn bind_sampler(
        &mut self,
        device: &wgpu::Device,
        slot_name: &str,
        resource_id: &str,
    ) -> Result<(), BindError> {
        if !self.sampler_pool.contains_key(resource_id) {
            return Err(BindError::ResourceNotFound(resource_id.to_string()));
        }
        let slot = self
            .sampler_slots
            .iter_mut()
            .find(|s| s.name == slot_name)
            .ok_or(BindError::NotFound(slot_name.to_string()))?;
        slot.resource_id = Some(resource_id.to_string());
        self.build_bind_group(device);
        log::info!("bound sampler {resource_id} to slot {slot_name}");
        Ok(())
    }

    fn unbind_texture_resource(&mut self, id: &str) -> bool {
        let mut changed = false;
        for slot in &mut self.texture_slots {
            if slot.resource_id.as_deref() == Some(id) {
                slot.resource_id = None;
                changed = true;
            }
        }
        changed
    }

    fn unbind_sampler_resource(&mut self, id: &str) -> bool {
        let mut changed = false;
        for slot in &mut self.sampler_slots {
            if slot.resource_id.as_deref() == Some(id) {
                slot.resource_id = None;
                changed = true;
            }
        }
        changed
    }

    /// Rebuild the bind group after external changes (e.g. texture resize from TSV).
    pub fn rebuild_bind_group(&mut self, device: &wgpu::Device) {
        self.build_bind_group(device);
    }

    fn build_bind_group(&mut self, device: &wgpu::Device) {
        let has_bindings = !self.buffers.is_empty()
            || !self.texture_slots.is_empty()
            || !self.sampler_slots.is_empty();

        if !has_bindings {
            self.bind_group_layout = None;
            self.bind_group = None;
            return;
        }

        let mut layout_entries: Vec<wgpu::BindGroupLayoutEntry> = Vec::new();

        for buf in &self.buffers {
            layout_entries.push(wgpu::BindGroupLayoutEntry {
                binding: buf.binding,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            });
        }

        for tex in &self.texture_slots {
            layout_entries.push(wgpu::BindGroupLayoutEntry {
                binding: tex.binding,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    sample_type: tex.sample_type,
                    view_dimension: tex.view_dimension,
                    multisampled: tex.multisampled,
                },
                count: None,
            });
        }

        for smp in &self.sampler_slots {
            layout_entries.push(wgpu::BindGroupLayoutEntry {
                binding: smp.binding,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler(smp.binding_type),
                count: None,
            });
        }

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("uniform_bind_group_layout"),
            entries: &layout_entries,
        });

        let mut bg_entries: Vec<wgpu::BindGroupEntry> = Vec::new();
        let fallback_textures_and_views: Vec<_> = self
            .texture_slots
            .iter()
            .map(|tex| create_placeholder_texture(device, tex))
            .collect();
        let fallback_samplers: Vec<_> = self
            .sampler_slots
            .iter()
            .map(|smp| {
                if smp.binding_type == wgpu::SamplerBindingType::Comparison {
                    device.create_sampler(&wgpu::SamplerDescriptor {
                        label: Some("comparison_sampler_fallback"),
                        compare: Some(wgpu::CompareFunction::LessEqual),
                        ..Default::default()
                    })
                } else {
                    device.create_sampler(&wgpu::SamplerDescriptor {
                        label: Some("sampler_fallback"),
                        ..Default::default()
                    })
                }
            })
            .collect();

        for buf in &self.buffers {
            bg_entries.push(wgpu::BindGroupEntry {
                binding: buf.binding,
                resource: buf.gpu_buffer.as_ref().unwrap().as_entire_binding(),
            });
        }

        for (idx, tex) in self.texture_slots.iter().enumerate() {
            let resource = tex
                .resource_id
                .as_deref()
                .and_then(|id| self.texture_pool.get(id))
                .filter(|res| texture_resource_matches(tex, res));
            let texture_view = if let Some(resource) = resource {
                resource.view()
            } else {
                &fallback_textures_and_views[idx].1
            };
            bg_entries.push(wgpu::BindGroupEntry {
                binding: tex.binding,
                resource: wgpu::BindingResource::TextureView(texture_view),
            });
        }

        for (idx, smp) in self.sampler_slots.iter().enumerate() {
            let sampler = if smp.binding_type == wgpu::SamplerBindingType::Comparison {
                &fallback_samplers[idx]
            } else if let Some(resource) = smp
                .resource_id
                .as_deref()
                .and_then(|id| self.sampler_pool.get(id))
            {
                resource.sampler()
            } else {
                &fallback_samplers[idx]
            };
            bg_entries.push(wgpu::BindGroupEntry {
                binding: smp.binding,
                resource: wgpu::BindingResource::Sampler(sampler),
            });
        }

        let bg = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("uniform_bind_group"),
            layout: &bgl,
            entries: &bg_entries,
        });

        self.bind_group_layout = Some(bgl);
        self.bind_group = Some(bg);
    }

    pub fn get(&mut self, name: &str) -> Option<UniformRef<'_>> {
        // Split the borrow: get member info first, then borrow buffer data
        let member_info = self.uniforms.get(name)?;
        let buffer_idx = member_info.buffer_idx;
        let ty = member_info.ty;
        let offset = member_info.offset;
        let size = member_info.size;

        let buf = &mut self.buffers[buffer_idx];
        self.dirty = true;
        Some(UniformRef {
            module: &self.module,
            layouter: &self.layouter,
            ty,
            data: &mut buf.data[offset..offset + size],
        })
    }

    pub fn flush(&mut self, queue: &wgpu::Queue) {
        if !self.dirty {
            return;
        }
        for buf in &self.buffers {
            if let Some(ref gpu_buffer) = buf.gpu_buffer {
                queue.write_buffer(gpu_buffer, 0, &buf.data);
            }
        }
        self.dirty = false;
    }
}

fn texture_resource_matches(slot: &TextureBindingSlot, resource: &TextureResource) -> bool {
    slot.view_dimension == resource.view_dimension()
}

fn create_placeholder_texture(
    device: &wgpu::Device,
    slot: &TextureBindingSlot,
) -> (wgpu::Texture, wgpu::TextureView) {
    let format = match slot.sample_type {
        wgpu::TextureSampleType::Float { .. } => wgpu::TextureFormat::Rgba8Unorm,
        wgpu::TextureSampleType::Sint => wgpu::TextureFormat::Rgba8Sint,
        wgpu::TextureSampleType::Uint => wgpu::TextureFormat::Rgba8Uint,
        wgpu::TextureSampleType::Depth => wgpu::TextureFormat::Depth32Float,
    };
    create_input_texture(
        device,
        &format!("placeholder_{}", slot.name),
        1,
        1,
        1,
        view_dimension_to_texture_dimension(slot.view_dimension),
        slot.view_dimension,
        format,
    )
}