use super::error::{RenderGraphError, Result};
use std::collections::HashMap;
use wgpu::{
    Buffer, BufferDescriptor, BufferUsages, Device, Extent3d, StoreOp, TextureDescriptor,
    TextureFormat, TextureUsages, TextureView, TextureViewDescriptor,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ResourceId(pub u32);

impl ResourceId {
    pub fn new(id: u32) -> Self {
        Self(id)
    }
}

#[derive(Debug, Clone)]
pub struct RenderGraphTextureDescriptor {
    pub format: TextureFormat,
    pub width: u32,
    pub height: u32,
    pub usage: TextureUsages,
    pub sample_count: u32,
    pub mip_level_count: u32,
    pub dimension: wgpu::TextureDimension,
    pub depth_or_array_layers: u32,
}

impl RenderGraphTextureDescriptor {
    pub fn to_wgpu_descriptor<'a>(&self, label: Option<&'a str>) -> TextureDescriptor<'a> {
        TextureDescriptor {
            label,
            size: Extent3d {
                width: self.width,
                height: self.height,
                depth_or_array_layers: self.depth_or_array_layers,
            },
            mip_level_count: self.mip_level_count,
            sample_count: self.sample_count,
            dimension: self.dimension,
            format: self.format,
            usage: self.usage,
            view_formats: &[],
        }
    }
}

#[derive(Debug, Clone)]
pub struct RenderGraphBufferDescriptor {
    pub size: u64,
    pub usage: BufferUsages,
    pub mapped_at_creation: bool,
}

impl RenderGraphBufferDescriptor {
    pub fn to_wgpu_descriptor<'a>(&self, label: Option<&'a str>) -> BufferDescriptor<'a> {
        BufferDescriptor {
            label,
            size: self.size,
            usage: self.usage,
            mapped_at_creation: self.mapped_at_creation,
        }
    }
}

#[derive(Debug, Clone)]
pub enum ResourceType {
    ExternalColor {
        clear_color: Option<wgpu::Color>,
        force_store: bool,
    },
    TransientColor {
        descriptor: RenderGraphTextureDescriptor,
        clear_color: Option<wgpu::Color>,
    },
    ExternalDepth {
        clear_depth: Option<f32>,
        force_store: bool,
    },
    TransientDepth {
        descriptor: RenderGraphTextureDescriptor,
        clear_depth: Option<f32>,
    },
    ExternalBuffer,
    TransientBuffer {
        descriptor: RenderGraphBufferDescriptor,
    },
}

#[derive(Debug, Clone)]
pub struct ResourceDescriptor {
    pub name: String,
    pub resource_type: ResourceType,
    pub is_external: bool,
}

pub enum ResourceHandle {
    ExternalTexture {
        view: TextureView,
        store_op: StoreOp,
        width: u32,
        height: u32,
    },
    TransientTexture {
        view: TextureView,
        store_op: StoreOp,
    },
    ExternalBuffer {
        buffer: Buffer,
    },
    TransientBuffer {
        buffer: Buffer,
    },
}

impl ResourceHandle {
    pub fn view(&self) -> &TextureView {
        match self {
            ResourceHandle::ExternalTexture { view, .. } => view,
            ResourceHandle::TransientTexture { view, .. } => view,
            _ => panic!("view() called on buffer resource"),
        }
    }

    pub fn store_op(&self) -> StoreOp {
        match self {
            ResourceHandle::ExternalTexture { store_op, .. } => *store_op,
            ResourceHandle::TransientTexture { store_op, .. } => *store_op,
            _ => panic!("store_op() called on buffer resource"),
        }
    }
}

pub struct RenderGraphResources {
    pub(super) descriptors: HashMap<ResourceId, ResourceDescriptor>,
    pub(super) handles: HashMap<ResourceId, ResourceHandle>,
    pub(super) versions: HashMap<ResourceId, u64>,
    pub(super) next_id: u32,
}

impl RenderGraphResources {
    pub fn new() -> Self {
        Self {
            descriptors: HashMap::new(),
            handles: HashMap::new(),
            versions: HashMap::new(),
            next_id: 0,
        }
    }

    pub fn get_version(&self, id: ResourceId) -> u64 {
        *self.versions.get(&id).unwrap_or(&0)
    }

    pub(super) fn increment_version(&mut self, id: ResourceId) {
        let version = self.versions.entry(id).or_insert(0);
        *version += 1;
    }

    pub fn register_external_resource(
        &mut self,
        name: String,
        resource_type: ResourceType,
    ) -> ResourceId {
        let id = ResourceId::new(self.next_id);
        self.next_id += 1;
        self.descriptors.insert(
            id,
            ResourceDescriptor {
                name,
                resource_type,
                is_external: true,
            },
        );
        id
    }

    pub fn set_external_texture(
        &mut self,
        id: ResourceId,
        view: TextureView,
        width: u32,
        height: u32,
    ) {
        self.handles.insert(
            id,
            ResourceHandle::ExternalTexture {
                view,
                store_op: StoreOp::Store,
                width,
                height,
            },
        );
    }

    pub fn set_external_buffer(&mut self, id: ResourceId, buffer: Buffer) {
        self.handles
            .insert(id, ResourceHandle::ExternalBuffer { buffer });
    }

    pub fn register_transient_resource(
        &mut self,
        name: String,
        resource_type: ResourceType,
    ) -> ResourceId {
        let id = ResourceId::new(self.next_id);
        self.next_id += 1;
        self.descriptors.insert(
            id,
            ResourceDescriptor {
                name,
                resource_type,
                is_external: false,
            },
        );
        id
    }

    pub fn get_handle(&self, id: ResourceId) -> Option<&ResourceHandle> {
        self.handles.get(&id)
    }

    pub fn get_descriptor(&self, id: ResourceId) -> Option<&ResourceDescriptor> {
        self.descriptors.get(&id)
    }

    pub fn get_color_attachment(
        &self,
        id: ResourceId,
        node_index: petgraph::graph::NodeIndex,
        clear_ops: &std::collections::HashSet<(petgraph::graph::NodeIndex, ResourceId)>,
    ) -> Result<(&TextureView, wgpu::LoadOp<wgpu::Color>, StoreOp)> {
        let handle = self
            .get_handle(id)
            .ok_or_else(|| RenderGraphError::ResourceNotBound {
                resource: format!("color_attachment_{:?}", id),
                id,
            })?;
        let descriptor =
            self.get_descriptor(id)
                .ok_or_else(|| RenderGraphError::DescriptorNotFound {
                    resource: format!("color_attachment_{:?}", id),
                    id,
                })?;

        let load_op = match &descriptor.resource_type {
            ResourceType::ExternalColor { clear_color, .. }
            | ResourceType::TransientColor { clear_color, .. } => {
                if clear_color.is_some() && clear_ops.contains(&(node_index, id)) {
                    wgpu::LoadOp::Clear(*clear_color.as_ref().unwrap())
                } else {
                    wgpu::LoadOp::Load
                }
            }
            _ => {
                return Err(RenderGraphError::TypeMismatch {
                    operation: "get_color_attachment".to_string(),
                    actual_type: match &descriptor.resource_type {
                        ResourceType::ExternalDepth { .. }
                        | ResourceType::TransientDepth { .. } => "depth".to_string(),
                        ResourceType::ExternalBuffer | ResourceType::TransientBuffer { .. } => {
                            "buffer".to_string()
                        }
                        _ => "unknown".to_string(),
                    },
                    resource: descriptor.name.clone(),
                });
            }
        };

        Ok((handle.view(), load_op, handle.store_op()))
    }

    pub fn get_depth_attachment(
        &self,
        id: ResourceId,
        node_index: petgraph::graph::NodeIndex,
        clear_ops: &std::collections::HashSet<(petgraph::graph::NodeIndex, ResourceId)>,
    ) -> Result<(&TextureView, wgpu::LoadOp<f32>, StoreOp)> {
        let handle = self
            .get_handle(id)
            .ok_or_else(|| RenderGraphError::ResourceNotBound {
                resource: format!("depth_attachment_{:?}", id),
                id,
            })?;
        let descriptor =
            self.get_descriptor(id)
                .ok_or_else(|| RenderGraphError::DescriptorNotFound {
                    resource: format!("depth_attachment_{:?}", id),
                    id,
                })?;

        let load_op = match &descriptor.resource_type {
            ResourceType::ExternalDepth { clear_depth, .. }
            | ResourceType::TransientDepth { clear_depth, .. } => {
                if clear_depth.is_some() && clear_ops.contains(&(node_index, id)) {
                    wgpu::LoadOp::Clear(*clear_depth.as_ref().unwrap())
                } else {
                    wgpu::LoadOp::Load
                }
            }
            _ => {
                return Err(RenderGraphError::TypeMismatch {
                    operation: "get_depth_attachment".to_string(),
                    actual_type: match &descriptor.resource_type {
                        ResourceType::ExternalColor { .. }
                        | ResourceType::TransientColor { .. } => "color".to_string(),
                        ResourceType::ExternalBuffer | ResourceType::TransientBuffer { .. } => {
                            "buffer".to_string()
                        }
                        _ => "unknown".to_string(),
                    },
                    resource: descriptor.name.clone(),
                });
            }
        };

        Ok((handle.view(), load_op, handle.store_op()))
    }

    pub fn get_texture_view(&self, id: ResourceId) -> Option<&TextureView> {
        self.get_handle(id).map(|handle| handle.view())
    }

    pub fn update_transient_descriptor(
        &mut self,
        id: ResourceId,
        width: u32,
        height: u32,
    ) -> Result<()> {
        let descriptor =
            self.get_descriptor(id)
                .ok_or_else(|| RenderGraphError::ResourceNotFound {
                    resource: format!("resource_{:?}", id),
                    id,
                })?;

        if descriptor.is_external {
            return Err(RenderGraphError::CannotResizeExternal {
                resource: descriptor.name.clone(),
            });
        }

        let name = descriptor.name.clone();

        let updated_descriptor = match &descriptor.resource_type {
            ResourceType::TransientColor {
                descriptor: tex_desc,
                clear_color,
            } => ResourceType::TransientColor {
                descriptor: RenderGraphTextureDescriptor {
                    width,
                    height,
                    ..tex_desc.clone()
                },
                clear_color: *clear_color,
            },
            ResourceType::TransientDepth {
                descriptor: tex_desc,
                clear_depth,
            } => ResourceType::TransientDepth {
                descriptor: RenderGraphTextureDescriptor {
                    width,
                    height,
                    ..tex_desc.clone()
                },
                clear_depth: *clear_depth,
            },
            ResourceType::TransientBuffer { .. } => {
                return Err(RenderGraphError::CannotResizeBuffer {
                    resource: name.clone(),
                });
            }
            _ => {
                return Err(RenderGraphError::CannotResizeNonTransient {
                    resource: name.clone(),
                });
            }
        };

        self.descriptors.insert(
            id,
            ResourceDescriptor {
                name,
                resource_type: updated_descriptor,
                is_external: false,
            },
        );

        Ok(())
    }

    pub fn allocate_transient_resources_with_aliasing(
        &mut self,
        device: &Device,
        store_ops: &HashMap<ResourceId, StoreOp>,
        aliasing_info: &mut super::graph::ResourceAliasingInfo,
    ) {
        for (pool_index, pool_slot) in aliasing_info.pools.iter_mut().enumerate() {
            if pool_slot.resource.is_some() {
                continue;
            }

            if let Some(descriptor_info) = &pool_slot.descriptor_info {
                let label = format!("pool_{}", pool_index);

                match descriptor_info {
                    super::graph::PoolDescriptorInfo::Texture(tex_desc) => {
                        let texture_descriptor = tex_desc.to_wgpu_descriptor(Some(&label));
                        let texture = device.create_texture(&texture_descriptor);
                        pool_slot.resource =
                            Some(super::graph::PooledResource::Texture { texture });
                    }
                    super::graph::PoolDescriptorInfo::Buffer(buf_desc) => {
                        let buffer_descriptor = buf_desc.to_wgpu_descriptor(Some(&label));
                        let buffer = device.create_buffer(&buffer_descriptor);
                        pool_slot.resource = Some(super::graph::PooledResource::Buffer { buffer });
                    }
                }
            }
        }

        let mut allocated_resources = Vec::new();

        for (resource_id, descriptor) in &self.descriptors {
            if descriptor.is_external || self.handles.contains_key(resource_id) {
                continue;
            }

            if let Some(&pool_index) = aliasing_info.aliases.get(resource_id)
                && let Some(pool_slot) = aliasing_info.pools.get(pool_index)
            {
                match &pool_slot.resource {
                    Some(super::graph::PooledResource::Texture { texture, .. }) => {
                        let view = texture.create_view(&TextureViewDescriptor::default());
                        let store_op = *store_ops.get(resource_id).unwrap_or(&StoreOp::Store);

                        self.handles.insert(
                            *resource_id,
                            ResourceHandle::TransientTexture { view, store_op },
                        );
                        allocated_resources.push(*resource_id);
                    }
                    Some(super::graph::PooledResource::Buffer { buffer, .. }) => {
                        self.handles.insert(
                            *resource_id,
                            ResourceHandle::TransientBuffer {
                                buffer: buffer.clone(),
                            },
                        );
                        allocated_resources.push(*resource_id);
                    }
                    None => {}
                }
            }
        }

        for resource_id in allocated_resources {
            self.increment_version(resource_id);
        }
    }
}

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