use std::borrow::Cow;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use winit::dpi::PhysicalSize;
use winit::window::Window;

use crate::shadertoy_adapter;

// Uniform buffer matching the shader's Uniforms struct
#[repr(C)]
#[derive(Copy, Clone, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct Uniforms {
    pub time: f32,
    pub _pad0: f32,
    pub resolution: [f32; 2],
    pub intensity: f32,
    pub _pad1: [f32; 3],
}

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

pub struct GpuState {
    surface: wgpu::Surface<'static>,
    device: wgpu::Device,
    queue: wgpu::Queue,
    config: wgpu::SurfaceConfiguration,

    // Pipelines
    fire_pipeline: wgpu::RenderPipeline,
    blur_pipeline: wgpu::RenderPipeline,
    composite_pipeline: wgpu::RenderPipeline,

    // Intermediate textures
    fire_texture: wgpu::Texture,
    fire_view: wgpu::TextureView,
    blur_texture: wgpu::Texture,
    blur_view: wgpu::TextureView,

    // Shared resources
    sampler: wgpu::Sampler,
    uniform_buffer: wgpu::Buffer,

    // Bind group layouts (needed for recreating bind groups on resize)
    #[allow(dead_code)]
    fire_bg_layout: wgpu::BindGroupLayout,
    blur_bg_layout: wgpu::BindGroupLayout,
    composite_bg_layout: wgpu::BindGroupLayout,

    // Bind groups
    fire_bind_group: wgpu::BindGroup,
    blur_bind_group: wgpu::BindGroup,
    composite_bind_group: wgpu::BindGroup,

    start_time: Instant,
    pub size: PhysicalSize<u32>,
}

impl GpuState {
    fn create_blur_bind_group(
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        uniform_buffer: &wgpu::Buffer,
        fire_view: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("blur_bind_group"),
            layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(fire_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    fn create_composite_bind_group(
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        uniform_buffer: &wgpu::Buffer,
        fire_view: &wgpu::TextureView,
        blur_view: &wgpu::TextureView,
        sampler: &wgpu::Sampler,
    ) -> wgpu::BindGroup {
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("composite_bind_group"),
            layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(fire_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::TextureView(blur_view),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: wgpu::BindingResource::Sampler(sampler),
                },
            ],
        })
    }

    pub async fn new(window: Arc<Window>) -> Self {
        let size = window.inner_size();

        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
            backends: wgpu::Backends::all(),
            ..Default::default()
        });

        let surface = instance.create_surface(window.clone()).unwrap();

        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                compatible_surface: Some(&surface),
                force_fallback_adapter: false,
            })
            .await
            .expect("Failed to find a suitable GPU adapter");

        log::info!("Using GPU: {:?}", adapter.get_info().name);

        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    label: Some("heatslop_device"),
                    required_features: wgpu::Features::empty(),
                    required_limits: wgpu::Limits::default(),
                    ..Default::default()
                },
                None,
            )
            .await
            .expect("Failed to create device");

        let surface_caps = surface.get_capabilities(&adapter);
        let surface_format = surface_caps
            .formats
            .iter()
            .find(|f| f.is_srgb())
            .copied()
            .unwrap_or(surface_caps.formats[0]);

        let config = wgpu::SurfaceConfiguration {
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            format: surface_format,
            width: size.width.max(1),
            height: size.height.max(1),
            present_mode: wgpu::PresentMode::AutoNoVsync,
            alpha_mode: surface_caps.alpha_modes[0],
            view_formats: vec![],
            desired_maximum_frame_latency: 2,
        };
        surface.configure(&device, &config);

        // ---- Shaders ----
        // Vertex shader: single WGSL module, override constant controls Y-flip
        let vert = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("vert"),
            source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!(
                "shaders/fullscreen.wgsl"
            ))),
        });
        let flip_off = HashMap::from([("0".to_string(), 0.0_f64)]);
        let flip_on = HashMap::from([("0".to_string(), 1.0_f64)]);

        // Fragment shaders: original Shadertoy code composed via adapter
        let fire_src = shadertoy_adapter::fire_frag();
        let fire_frag = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("fire_frag"),
            source: wgpu::ShaderSource::Glsl {
                shader: Cow::Owned(fire_src),
                stage: naga::ShaderStage::Fragment,
                defines: Default::default(),
            },
        });
        let blur_src = shadertoy_adapter::blur_frag();
        let blur_frag = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("blur_frag"),
            source: wgpu::ShaderSource::Glsl {
                shader: Cow::Owned(blur_src),
                stage: naga::ShaderStage::Fragment,
                defines: Default::default(),
            },
        });
        let composite_src = shadertoy_adapter::composite_frag();
        let composite_frag = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("composite_frag"),
            source: wgpu::ShaderSource::Glsl {
                shader: Cow::Owned(composite_src),
                stage: naga::ShaderStage::Fragment,
                defines: Default::default(),
            },
        });

        // ---- Uniform buffer ----
        let uniforms = Uniforms {
            time: 0.0,
            _pad0: 0.0,
            resolution: [size.width as f32, size.height as f32],
            intensity: 1.0,
            _pad1: [0.0; 3],
        };
        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("uniform_buffer"),
            size: std::mem::size_of::<Uniforms>() as u64,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        queue.write_buffer(&uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));

        // ---- Sampler ----
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("tex_sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            ..Default::default()
        });

        // ---- Intermediate textures ----
        let (fire_texture, fire_view) =
            create_intermediate_texture(&device, size.width, size.height, "fire_texture");
        let (blur_texture, blur_view) =
            create_intermediate_texture(&device, size.width, size.height, "blur_texture");

        // ---- Bind group layouts ----

        // Fire: just uniforms
        let fire_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("fire_bg_layout"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });

        // Blur: uniforms + input texture + sampler
        let blur_bg_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("blur_bg_layout"),
            entries: &[
                wgpu::BindGroupLayoutEntry {
                    binding: 0,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 1,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Texture {
                        sample_type: wgpu::TextureSampleType::Float { filterable: true },
                        view_dimension: wgpu::TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: None,
                },
                wgpu::BindGroupLayoutEntry {
                    binding: 2,
                    visibility: wgpu::ShaderStages::FRAGMENT,
                    ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                    count: None,
                },
            ],
        });

        // Composite: uniforms + fire texture + blur texture + sampler
        let composite_bg_layout =
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("composite_bg_layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: false,
                            min_binding_size: None,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: true },
                            view_dimension: wgpu::TextureViewDimension::D2,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 3,
                        visibility: wgpu::ShaderStages::FRAGMENT,
                        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                        count: None,
                    },
                ],
            });

        // ---- Bind groups ----
        let fire_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("fire_bind_group"),
            layout: &fire_bg_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: uniform_buffer.as_entire_binding(),
            }],
        });

        let blur_bind_group = Self::create_blur_bind_group(
            &device,
            &blur_bg_layout,
            &uniform_buffer,
            &fire_view,
            &sampler,
        );

        let composite_bind_group = Self::create_composite_bind_group(
            &device,
            &composite_bg_layout,
            &uniform_buffer,
            &fire_view,
            &blur_view,
            &sampler,
        );

        // ---- Pipeline layouts ----
        let fire_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("fire_pipeline_layout"),
            bind_group_layouts: &[&fire_bg_layout],
            push_constant_ranges: &[],
        });
        let blur_pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("blur_pipeline_layout"),
            bind_group_layouts: &[&blur_bg_layout],
            push_constant_ranges: &[],
        });
        let composite_pipeline_layout =
            device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("composite_pipeline_layout"),
                bind_group_layouts: &[&composite_bg_layout],
                push_constant_ranges: &[],
            });

        // ---- Render pipelines ----

        // Fire pipeline: renders to Rgba16Float intermediate texture
        let fire_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("fire_pipeline"),
            layout: Some(&fire_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vert,
                entry_point: Some("main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &flip_off,
                    ..Default::default()
                },
            },
            fragment: Some(wgpu::FragmentState {
                module: &fire_frag,
                entry_point: Some("main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: wgpu::TextureFormat::Rgba16Float,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                cull_mode: None,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        // Blur pipeline: renders to Rgba16Float intermediate texture
        let blur_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("blur_pipeline"),
            layout: Some(&blur_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vert,
                entry_point: Some("main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &flip_on,
                    ..Default::default()
                },
            },
            fragment: Some(wgpu::FragmentState {
                module: &blur_frag,
                entry_point: Some("main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: wgpu::TextureFormat::Rgba16Float,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                cull_mode: None,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        // Composite pipeline: renders to swapchain (sRGB surface)
        let composite_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("composite_pipeline"),
            layout: Some(&composite_pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vert,
                entry_point: Some("main"),
                buffers: &[],
                compilation_options: wgpu::PipelineCompilationOptions {
                    constants: &flip_on,
                    ..Default::default()
                },
            },
            fragment: Some(wgpu::FragmentState {
                module: &composite_frag,
                entry_point: Some("main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: surface_format,
                    blend: Some(wgpu::BlendState::REPLACE),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                cull_mode: None,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        Self {
            surface,
            device,
            queue,
            config,
            fire_pipeline,
            blur_pipeline,
            composite_pipeline,
            fire_texture,
            fire_view,
            blur_texture,
            blur_view,
            sampler,
            uniform_buffer,
            fire_bg_layout,
            blur_bg_layout,
            composite_bg_layout,
            fire_bind_group,
            blur_bind_group,
            composite_bind_group,
            start_time: Instant::now(),
            size,
        }
    }

    pub fn resize(&mut self, new_size: PhysicalSize<u32>) {
        if new_size.width > 0 && new_size.height > 0 {
            self.size = new_size;
            self.config.width = new_size.width;
            self.config.height = new_size.height;
            self.surface.configure(&self.device, &self.config);

            // Recreate intermediate textures at new size
            let (ft, fv) = create_intermediate_texture(
                &self.device,
                new_size.width,
                new_size.height,
                "fire_texture",
            );
            self.fire_texture = ft;
            self.fire_view = fv;

            let (bt, bv) = create_intermediate_texture(
                &self.device,
                new_size.width,
                new_size.height,
                "blur_texture",
            );
            self.blur_texture = bt;
            self.blur_view = bv;

            // Recreate bind groups that reference the textures
            self.blur_bind_group = Self::create_blur_bind_group(
                &self.device,
                &self.blur_bg_layout,
                &self.uniform_buffer,
                &self.fire_view,
                &self.sampler,
            );

            self.composite_bind_group = Self::create_composite_bind_group(
                &self.device,
                &self.composite_bg_layout,
                &self.uniform_buffer,
                &self.fire_view,
                &self.blur_view,
                &self.sampler,
            );
        }
    }

    pub fn render(&mut self) -> Result<(), wgpu::SurfaceError> {
        let elapsed = self.start_time.elapsed().as_secs_f32();

        // Update uniforms
        let uniforms = Uniforms {
            time: elapsed,
            _pad0: 0.0,
            resolution: [self.size.width as f32, self.size.height as f32],
            intensity: 1.0,
            _pad1: [0.0; 3],
        };
        self.queue
            .write_buffer(&self.uniform_buffer, 0, bytemuck::cast_slice(&[uniforms]));

        let output = self.surface.get_current_texture()?;
        let output_view = output
            .texture
            .create_view(&wgpu::TextureViewDescriptor::default());

        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("render_encoder"),
            });

        // ---- Pass 1: Fire generation → fire_texture ----
        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("fire_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &self.fire_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            pass.set_pipeline(&self.fire_pipeline);
            pass.set_bind_group(0, &self.fire_bind_group, &[]);
            pass.draw(0..3, 0..1);
        }

        // ---- Pass 2: Gaussian blur of fire → blur_texture ----
        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("blur_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &self.blur_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            pass.set_pipeline(&self.blur_pipeline);
            pass.set_bind_group(0, &self.blur_bind_group, &[]);
            pass.draw(0..3, 0..1);
        }

        // ---- Pass 3: Composite (fire + bloom) → screen ----
        {
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("composite_pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &output_view,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            pass.set_pipeline(&self.composite_pipeline);
            pass.set_bind_group(0, &self.composite_bind_group, &[]);
            pass.draw(0..3, 0..1);
        }

        self.queue.submit(std::iter::once(encoder.finish()));
        output.present();

        Ok(())
    }
}