use super::error::{RenderGraphError, Result};
use super::resources::{
    RenderGraphBufferDescriptor, RenderGraphResources, RenderGraphTextureDescriptor,
    ResourceHandle, ResourceId, ResourceType,
};
use std::collections::HashMap;
use wgpu::{
    Buffer, BufferUsages, CommandEncoder, Device, TextureFormat, TextureUsages, TextureView,
};

pub struct SubGraphRunCommand<'a> {
    pub sub_graph_name: String,
    pub inputs: Vec<SlotValue<'a>>,
}

pub struct PassExecutionContext<'r, 'e, C = ()> {
    pub encoder: &'e mut CommandEncoder,
    pub resources: &'r RenderGraphResources,
    pub device: &'r Device,
    pub(super) slot_mappings: &'r HashMap<String, ResourceId>,
    pub configs: &'r C,
    pub(crate) sub_graph_commands: Vec<SubGraphRunCommand<'r>>,
    pub(super) node_index: petgraph::graph::NodeIndex,
    pub(super) clear_ops: &'r std::collections::HashSet<(petgraph::graph::NodeIndex, ResourceId)>,
    pub(super) pass_enabled: bool,
}

impl<'r, 'e, C> PassExecutionContext<'r, 'e, C> {
    pub fn is_pass_enabled(&self) -> bool {
        self.pass_enabled
    }

    pub fn get_slot(&self, slot: &str) -> Result<ResourceId> {
        self.slot_mappings
            .get(slot)
            .copied()
            .ok_or_else(|| RenderGraphError::SlotNotFound {
                slot: slot.to_string(),
                pass: "unknown".to_string(),
            })
    }

    pub fn get_texture_view(&self, slot: &str) -> Result<&'r wgpu::TextureView> {
        let resource_id = self.get_slot(slot)?;
        self.resources.get_texture_view(resource_id).ok_or_else(|| {
            RenderGraphError::ResourceNotBound {
                resource: slot.to_string(),
                id: resource_id,
            }
        })
    }

    pub fn get_color_attachment(
        &self,
        slot: &str,
    ) -> Result<(
        &'r wgpu::TextureView,
        wgpu::LoadOp<wgpu::Color>,
        wgpu::StoreOp,
    )> {
        let resource_id = self.get_slot(slot)?;
        self.resources
            .get_color_attachment(resource_id, self.node_index, self.clear_ops)
    }

    pub fn get_depth_attachment(
        &self,
        slot: &str,
    ) -> Result<(&'r wgpu::TextureView, wgpu::LoadOp<f32>, wgpu::StoreOp)> {
        let resource_id = self.get_slot(slot)?;
        self.resources
            .get_depth_attachment(resource_id, self.node_index, self.clear_ops)
    }

    pub fn get_buffer(&self, slot: &str) -> Result<&'r wgpu::Buffer> {
        let resource_id = self.get_slot(slot)?;
        let handle = self.resources.get_handle(resource_id).ok_or_else(|| {
            RenderGraphError::ResourceNotBound {
                resource: slot.to_string(),
                id: resource_id,
            }
        })?;

        match handle {
            ResourceHandle::ExternalBuffer { buffer }
            | ResourceHandle::TransientBuffer { buffer } => Ok(buffer),
            _ => Err(RenderGraphError::TypeMismatch {
                operation: "get_buffer".to_string(),
                actual_type: "texture".to_string(),
                resource: slot.to_string(),
            }),
        }
    }

    pub fn get_texture_size(&self, slot: &str) -> Result<(u32, u32)> {
        let resource_id = self.get_slot(slot)?;
        let descriptor = self.resources.get_descriptor(resource_id).ok_or_else(|| {
            RenderGraphError::DescriptorNotFound {
                resource: slot.to_string(),
                id: resource_id,
            }
        })?;

        match &descriptor.resource_type {
            ResourceType::TransientColor {
                descriptor: texture_desc,
                ..
            }
            | ResourceType::TransientDepth {
                descriptor: texture_desc,
                ..
            } => Ok((texture_desc.width, texture_desc.height)),
            ResourceType::ExternalColor { .. } | ResourceType::ExternalDepth { .. } => {
                let handle = self.resources.get_handle(resource_id).ok_or_else(|| {
                    RenderGraphError::ResourceNotBound {
                        resource: slot.to_string(),
                        id: resource_id,
                    }
                })?;

                match handle {
                    ResourceHandle::ExternalTexture { width, height, .. } => Ok((*width, *height)),
                    _ => Err(RenderGraphError::TypeMismatch {
                        operation: "get_texture_size".to_string(),
                        actual_type: "transient_texture".to_string(),
                        resource: slot.to_string(),
                    }),
                }
            }
            _ => Err(RenderGraphError::TypeMismatch {
                operation: "get_texture_size".to_string(),
                actual_type: "buffer".to_string(),
                resource: slot.to_string(),
            }),
        }
    }

    pub fn run_sub_graph(&mut self, sub_graph_name: String, inputs: Vec<SlotValue<'r>>) {
        self.sub_graph_commands.push(SubGraphRunCommand {
            sub_graph_name,
            inputs,
        });
    }

    pub fn into_sub_graph_commands(self) -> Vec<SubGraphRunCommand<'r>> {
        self.sub_graph_commands
    }
}

#[cfg(not(target_arch = "wasm32"))]
pub trait PassNode<C = ()>: Send + Sync + std::any::Any {
    fn name(&self) -> &str;
    fn reads(&self) -> Vec<&str>;
    fn writes(&self) -> Vec<&str>;
    fn reads_writes(&self) -> Vec<&str> {
        Vec::new()
    }
    fn prepare(&mut self, _device: &Device, _queue: &wgpu::Queue, _configs: &C) {}
    fn invalidate_bind_groups(&mut self) {}
    fn execute<'r, 'e>(
        &mut self,
        context: PassExecutionContext<'r, 'e, C>,
    ) -> Result<Vec<SubGraphRunCommand<'r>>>;
}

#[cfg(target_arch = "wasm32")]
pub trait PassNode<C = ()>: std::any::Any {
    fn name(&self) -> &str;
    fn reads(&self) -> Vec<&str>;
    fn writes(&self) -> Vec<&str>;
    fn reads_writes(&self) -> Vec<&str> {
        Vec::new()
    }
    fn prepare(&mut self, _device: &Device, _queue: &wgpu::Queue, _configs: &C) {}
    fn invalidate_bind_groups(&mut self) {}
    fn execute<'r, 'e>(
        &mut self,
        context: PassExecutionContext<'r, 'e, C>,
    ) -> Result<Vec<SubGraphRunCommand<'r>>>;
}

pub struct GraphNode<C> {
    pub name: String,
    pub reads: Vec<ResourceId>,
    pub writes: Vec<ResourceId>,
    pub reads_writes: Vec<ResourceId>,
    pub pass: Box<dyn PassNode<C> + 'static>,
    pub enabled: bool,
}

pub enum SlotValue<'a> {
    TextureView {
        view: &'a TextureView,
        width: u32,
        height: u32,
    },
    Buffer(&'a Buffer),
}

#[derive(Clone)]
pub struct SubGraphInputSlot {
    pub name: String,
}

pub struct ColorTextureBuilder<'a, C = ()> {
    pub(super) graph: &'a mut super::graph::RenderGraph<C>,
    pub(super) name: String,
    pub(super) descriptor: RenderGraphTextureDescriptor,
    pub(super) clear_color: Option<wgpu::Color>,
    pub(super) force_store: bool,
}

impl<'a, C> ColorTextureBuilder<'a, C> {
    pub fn format(mut self, format: TextureFormat) -> Self {
        self.descriptor.format = format;
        self
    }

    pub fn size(mut self, width: u32, height: u32) -> Self {
        self.descriptor.width = width;
        self.descriptor.height = height;
        self
    }

    pub fn usage(mut self, usage: TextureUsages) -> Self {
        self.descriptor.usage = usage;
        self
    }

    pub fn sample_count(mut self, count: u32) -> Self {
        self.descriptor.sample_count = count;
        self
    }

    pub fn mip_levels(mut self, levels: u32) -> Self {
        self.descriptor.mip_level_count = levels;
        self
    }

    pub fn clear_color(mut self, color: wgpu::Color) -> Self {
        self.clear_color = Some(color);
        self
    }

    pub fn no_store(mut self) -> Self {
        self.force_store = false;
        self
    }

    pub fn external(self) -> ResourceId {
        self.graph.resources.register_external_resource(
            self.name,
            ResourceType::ExternalColor {
                clear_color: self.clear_color,
                force_store: self.force_store,
            },
        )
    }

    pub fn transient(self) -> ResourceId {
        self.graph.resources.register_transient_resource(
            self.name,
            ResourceType::TransientColor {
                descriptor: self.descriptor,
                clear_color: self.clear_color,
            },
        )
    }
}

pub struct DepthTextureBuilder<'a, C = ()> {
    pub(super) graph: &'a mut super::graph::RenderGraph<C>,
    pub(super) name: String,
    pub(super) descriptor: RenderGraphTextureDescriptor,
    pub(super) clear_depth: Option<f32>,
    pub(super) force_store: bool,
}

impl<'a, C> DepthTextureBuilder<'a, C> {
    pub fn format(mut self, format: TextureFormat) -> Self {
        self.descriptor.format = format;
        self
    }

    pub fn size(mut self, width: u32, height: u32) -> Self {
        self.descriptor.width = width;
        self.descriptor.height = height;
        self
    }

    pub fn usage(mut self, usage: TextureUsages) -> Self {
        self.descriptor.usage = usage;
        self
    }

    pub fn sample_count(mut self, count: u32) -> Self {
        self.descriptor.sample_count = count;
        self
    }

    pub fn mip_levels(mut self, levels: u32) -> Self {
        self.descriptor.mip_level_count = levels;
        self
    }

    pub fn array_layers(mut self, layers: u32) -> Self {
        self.descriptor.depth_or_array_layers = layers;
        self
    }

    pub fn clear_depth(mut self, depth: f32) -> Self {
        self.clear_depth = Some(depth);
        self
    }

    pub fn no_store(mut self) -> Self {
        self.force_store = false;
        self
    }

    pub fn external(self) -> ResourceId {
        self.graph.resources.register_external_resource(
            self.name,
            ResourceType::ExternalDepth {
                clear_depth: self.clear_depth,
                force_store: self.force_store,
            },
        )
    }

    pub fn transient(self) -> ResourceId {
        self.graph.resources.register_transient_resource(
            self.name,
            ResourceType::TransientDepth {
                descriptor: self.descriptor,
                clear_depth: self.clear_depth,
            },
        )
    }
}

pub struct BufferBuilder<'a, C = ()> {
    pub(super) graph: &'a mut super::graph::RenderGraph<C>,
    pub(super) name: String,
    pub(super) descriptor: RenderGraphBufferDescriptor,
}

impl<'a, C> BufferBuilder<'a, C> {
    pub fn size(mut self, size: u64) -> Self {
        self.descriptor.size = size;
        self
    }

    pub fn usage(mut self, usage: BufferUsages) -> Self {
        self.descriptor.usage = usage;
        self
    }

    pub fn mapped_at_creation(mut self, mapped: bool) -> Self {
        self.descriptor.mapped_at_creation = mapped;
        self
    }

    pub fn external(self) -> ResourceId {
        self.graph
            .resources
            .register_external_resource(self.name, ResourceType::ExternalBuffer)
    }

    pub fn transient(self) -> ResourceId {
        self.graph.resources.register_transient_resource(
            self.name,
            ResourceType::TransientBuffer {
                descriptor: self.descriptor,
            },
        )
    }
}

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

impl ResourceTemplate {
    pub fn new(format: TextureFormat, width: u32, height: u32) -> Self {
        Self {
            format,
            width,
            height,
            usage: TextureUsages::RENDER_ATTACHMENT | TextureUsages::TEXTURE_BINDING,
            sample_count: 1,
            mip_level_count: 1,
            dimension: wgpu::TextureDimension::D2,
            depth_or_array_layers: 1,
        }
    }

    pub fn usage(mut self, usage: TextureUsages) -> Self {
        self.usage = usage;
        self
    }

    pub fn sample_count(mut self, count: u32) -> Self {
        self.sample_count = count;
        self
    }

    pub fn mip_levels(mut self, levels: u32) -> Self {
        self.mip_level_count = levels;
        self
    }

    pub fn cube_map(mut self) -> Self {
        self.dimension = wgpu::TextureDimension::D2;
        self.depth_or_array_layers = 6;
        self
    }

    pub fn array_layers(mut self, layers: u32) -> Self {
        self.depth_or_array_layers = layers;
        self
    }

    pub fn dimension_3d(mut self, depth: u32) -> Self {
        self.dimension = wgpu::TextureDimension::D3;
        self.depth_or_array_layers = depth;
        self
    }
}

pub struct PassBuilder<'a, C: 'static = ()> {
    pub(super) graph: &'a mut super::graph::RenderGraph<C>,
    pub(super) pass: Option<Box<dyn PassNode<C>>>,
    pub(super) slots: Vec<(&'static str, ResourceId)>,
}

impl<'a, C: 'static> PassBuilder<'a, C> {
    pub fn read(mut self, slot: &'static str, resource: ResourceId) -> Self {
        self.slots.push((slot, resource));
        self
    }

    pub fn write(mut self, slot: &'static str, resource: ResourceId) -> Self {
        self.slots.push((slot, resource));
        self
    }

    pub fn slot(mut self, slot: &'static str, resource: ResourceId) -> Self {
        self.slots.push((slot, resource));
        self
    }
}

impl<'a, C: 'static> Drop for PassBuilder<'a, C> {
    fn drop(&mut self) {
        if let Some(pass) = self.pass.take() {
            let result = self.graph.add_pass(pass, &self.slots);
            if let Err(e) = result {
                panic!("Failed to add render pass: {}", e);
            }
        }
    }
}

pub struct ResourcePool<'a, C = ()> {
    pub(super) graph: &'a mut super::graph::RenderGraph<C>,
    pub(super) template: ResourceTemplate,
}

impl<'a, C: 'static> ResourcePool<'a, C> {
    pub fn transient(&mut self, name: &str) -> ResourceId {
        self.graph
            .transient_color_from_template(name, &self.template)
    }

    pub fn transient_many(&mut self, names: &[&str]) -> Vec<ResourceId> {
        names.iter().map(|name| self.transient(name)).collect()
    }

    pub fn external(&mut self, name: &str) -> ResourceId {
        self.graph.external_color(name)
    }
}
