//! Temporal Super-Resolution (TSR)
//!
//! Implements temporal upscaling for improved image quality at reduced rendering cost.
//! Key features:
//! - Jittered projection matrices (Halton sequence)
//! - Motion vector tracking for temporal reprojection
//! - History buffer accumulation with AABB clamping
//! - Configurable upscale factors (1.33x, 1.5x, 2.0x)

use ash::vk;
use std::sync::Arc;

use crate::vulkan::VulkanDevice;
use crate::Result;

/// Upscale quality presets
#[derive(Clone, Copy, Debug, Default)]
pub enum VsrQuality {
    /// 2.0x upscale (50% render resolution) - Fastest
    Performance,
    /// 1.5x upscale (67% render resolution) - Balanced
    #[default]
    Balanced,
    /// 1.33x upscale (75% render resolution) - Quality
    Quality,
    /// 1.0x (native) - No upscaling, TAA only
    Native,
}

impl VsrQuality {
    /// Get the upscale factor
    pub fn factor(&self) -> f32 {
        match self {
            VsrQuality::Performance => 2.0,
            VsrQuality::Balanced => 1.5,
            VsrQuality::Quality => 1.33,
            VsrQuality::Native => 1.0,
        }
    }

    /// Get render resolution from display resolution
    pub fn render_size(&self, display_width: u32, display_height: u32) -> (u32, u32) {
        let factor = self.factor();
        (
            (display_width as f32 / factor) as u32,
            (display_height as f32 / factor) as u32,
        )
    }
}

/// Halton sequence for QMC jittering
pub struct HaltonSequence {
    u: u32, // current index
    b2: Vec<f32>,
    b3: Vec<f32>,
}

impl HaltonSequence {
    pub fn new(samples: usize) -> Self {
        let mut b2 = Vec::with_capacity(samples);
        let mut b3 = Vec::with_capacity(samples);

        for i in 1..=samples {
            b2.push(Self::phi(i as u32, 2));
            b3.push(Self::phi(i as u32, 3));
        }

        Self { u: 0, b2, b3 }
    }

    // Corput radical inverse in base b
    fn phi(mut i: u32, b: u32) -> f32 {
        let mut r = 0.0;
        let mut f = 1.0 / b as f32;
        while i > 0 {
            r += f * (i % b) as f32;
            i /= b;
            f /= b as f32;
        }
        r
    }

    /// Next jitter sample in [-0.5, 0.5]
    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self) -> (f32, f32) {
        let idx = self.u as usize % self.b2.len();
        self.u = self.u.wrapping_add(1);
        (self.b2[idx] - 0.5, self.b3[idx] - 0.5)
    }

    pub fn reset(&mut self) {
        self.u = 0;
    }
}

/// Push constants for TSR upscale shader
#[repr(C)]
#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]
pub struct TsrPushConstants {
    /// Current jitter offset (applied to projection)
    pub jitter: [f32; 2],
    /// Screen dimensions (render resolution)
    pub render_size: [f32; 2],
    /// Screen dimensions (display resolution)
    pub display_size: [f32; 2],
    /// History blend factor (0.9 = 90% history)
    pub history_weight: f32,
    /// Frame index for temporal variation
    pub frame_index: u32,
}

/// Temporal Super-Resolution pass
pub struct VsrPass {
    device: Arc<ash::Device>,

    // Render-res buffers
    motion_img: vk::Image,
    motion_alloc: Option<vk_mem::Allocation>,
    motion_v: vk::ImageView,

    // Display-res history (ping-pong)
    history_imgs: [vk::Image; 2],
    history_allocs: [Option<vk_mem::Allocation>; 2],
    history_vs: [vk::ImageView; 2],

    upscale_pl: vk::Pipeline,
    upscale_layout: vk::PipelineLayout,

    descriptor_pool: vk::DescriptorPool,
    desc_layout: vk::DescriptorSetLayout,
    desc_sets: [vk::DescriptorSet; 2],

    sampler: vk::Sampler,

    // Dimensions
    render_w: u32,
    render_h: u32,
    display_w: u32,
    display_h: u32,

    quality: VsrQuality,
    halton: HaltonSequence,
    frame_idx: u32,

    initialized: bool,
}

impl VsrPass {
    pub fn new(device: Arc<ash::Device>) -> Self {
        Self {
            device,
            motion_img: vk::Image::null(),
            motion_alloc: None,
            motion_v: vk::ImageView::null(),
            history_imgs: [vk::Image::null(); 2],
            history_allocs: [None, None],
            history_vs: [vk::ImageView::null(); 2],
            upscale_pl: vk::Pipeline::null(),
            upscale_layout: vk::PipelineLayout::null(),
            descriptor_pool: vk::DescriptorPool::null(),
            desc_layout: vk::DescriptorSetLayout::null(),
            desc_sets: [vk::DescriptorSet::null(); 2],
            sampler: vk::Sampler::null(),
            render_w: 0,
            render_h: 0,
            display_w: 0,
            display_h: 0,
            quality: VsrQuality::default(),
            halton: HaltonSequence::new(16),
            frame_idx: 0,
            initialized: false,
        }
    }

    /// # Safety
    /// Allocator must be valid.
    pub unsafe fn init(
        &mut self,
        alloc: &vk_mem::Allocator,
        _vulkan_device: &VulkanDevice,
        w: u32,
        h: u32,
        quality: VsrQuality,
    ) {
        if self.initialized {
            return;
        }

        self.display_w = w;
        self.display_h = h;
        self.quality = quality;

        let (rw, rh) = quality.render_size(w, h);
        self.render_w = rw;
        self.render_h = rh;

        // VSR resources are mandatory for TSR rendering paths
        self.create_motion_image(alloc)
            .expect("TSR: Motion buffer failed");
        self.create_history_images(alloc)
            .expect("TSR: History buffer failed");
        self.create_sampler().expect("TSR: Sampler failed");
        self.create_descriptors().expect("TSR: Descriptors failed");
        self.create_pipeline().expect("TSR: Pipeline failed");

        self.initialized = true;
    }

    unsafe fn create_motion_image(&mut self, alloc: &vk_mem::Allocator) -> Result<()> {
        use vk_mem::Alloc;

        let info = vk::ImageCreateInfo::default()
            .image_type(vk::ImageType::TYPE_2D)
            .format(vk::Format::R16G16_SFLOAT)
            .extent(vk::Extent3D {
                width: self.render_w,
                height: self.render_h,
                depth: 1,
            })
            .mip_levels(1)
            .array_layers(1)
            .samples(vk::SampleCountFlags::TYPE_1)
            .tiling(vk::ImageTiling::OPTIMAL)
            .usage(vk::ImageUsageFlags::COLOR_ATTACHMENT | vk::ImageUsageFlags::SAMPLED)
            .sharing_mode(vk::SharingMode::EXCLUSIVE);

        let alloc_info = vk_mem::AllocationCreateInfo {
            usage: vk_mem::MemoryUsage::AutoPreferDevice,
            ..Default::default()
        };

        let (img, allocation) = alloc
            .create_image(&info, &alloc_info)
            .map_err(|e| crate::AshError::VulkanError(format!("Motion image: {e:?}")))?;

        self.motion_img = img;
        self.motion_alloc = Some(allocation);

        let view_info = vk::ImageViewCreateInfo::default()
            .image(self.motion_img)
            .view_type(vk::ImageViewType::TYPE_2D)
            .format(vk::Format::R16G16_SFLOAT)
            .subresource_range(
                vk::ImageSubresourceRange::default()
                    .aspect_mask(vk::ImageAspectFlags::COLOR)
                    .level_count(1)
                    .layer_count(1),
            );

        self.motion_v = self.device.create_image_view(&view_info, None)?;
        Ok(())
    }

    unsafe fn create_history_images(&mut self, alloc: &vk_mem::Allocator) -> Result<()> {
        use vk_mem::Alloc;

        for i in 0..2 {
            let info = vk::ImageCreateInfo::default()
                .image_type(vk::ImageType::TYPE_2D)
                .format(vk::Format::R16G16B16A16_SFLOAT)
                .extent(vk::Extent3D {
                    width: self.display_w,
                    height: self.display_h,
                    depth: 1,
                })
                .mip_levels(1)
                .array_layers(1)
                .samples(vk::SampleCountFlags::TYPE_1)
                .tiling(vk::ImageTiling::OPTIMAL)
                .usage(
                    vk::ImageUsageFlags::STORAGE
                        | vk::ImageUsageFlags::SAMPLED
                        | vk::ImageUsageFlags::TRANSFER_DST,
                )
                .sharing_mode(vk::SharingMode::EXCLUSIVE);

            let alloc_info = vk_mem::AllocationCreateInfo {
                usage: vk_mem::MemoryUsage::AutoPreferDevice,
                ..Default::default()
            };

            let (img, allocation) = alloc
                .create_image(&info, &alloc_info)
                .map_err(|e| crate::AshError::VulkanError(format!("History image {i}: {e:?}")))?;

            self.history_imgs[i] = img;
            self.history_allocs[i] = Some(allocation);

            let view_info = vk::ImageViewCreateInfo::default()
                .image(self.history_imgs[i])
                .view_type(vk::ImageViewType::TYPE_2D)
                .format(vk::Format::R16G16B16A16_SFLOAT)
                .subresource_range(
                    vk::ImageSubresourceRange::default()
                        .aspect_mask(vk::ImageAspectFlags::COLOR)
                        .level_count(1)
                        .layer_count(1),
                );

            self.history_vs[i] = self.device.create_image_view(&view_info, None)?;
        }

        Ok(())
    }

    unsafe fn create_sampler(&mut self) -> Result<()> {
        let info = vk::SamplerCreateInfo::default()
            .mag_filter(vk::Filter::LINEAR)
            .min_filter(vk::Filter::LINEAR)
            .mipmap_mode(vk::SamplerMipmapMode::NEAREST)
            .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
            .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)
            .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_EDGE);

        self.sampler = self.device.create_sampler(&info, None)?;
        Ok(())
    }

    unsafe fn create_descriptors(&mut self) -> Result<()> {
        // [0..3]: Color, Motion, Depth, History
        // [4]: Output
        let bindings = [
            vk::DescriptorSetLayoutBinding::default()
                .binding(0)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::COMPUTE),
            vk::DescriptorSetLayoutBinding::default()
                .binding(1)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::COMPUTE),
            vk::DescriptorSetLayoutBinding::default()
                .binding(2)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::COMPUTE),
            vk::DescriptorSetLayoutBinding::default()
                .binding(3)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::COMPUTE),
            vk::DescriptorSetLayoutBinding::default()
                .binding(4)
                .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::COMPUTE),
        ];

        let layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings);
        self.desc_layout = self
            .device
            .create_descriptor_set_layout(&layout_info, None)?;

        let pool_sizes = [
            vk::DescriptorPoolSize {
                ty: vk::DescriptorType::COMBINED_IMAGE_SAMPLER,
                descriptor_count: 8,
            },
            vk::DescriptorPoolSize {
                ty: vk::DescriptorType::STORAGE_IMAGE,
                descriptor_count: 2,
            },
        ];

        let pool_info = vk::DescriptorPoolCreateInfo::default()
            .max_sets(2)
            .pool_sizes(&pool_sizes);

        self.descriptor_pool = self.device.create_descriptor_pool(&pool_info, None)?;

        let layouts = [self.desc_layout, self.desc_layout];
        let alloc_info = vk::DescriptorSetAllocateInfo::default()
            .descriptor_pool(self.descriptor_pool)
            .set_layouts(&layouts);

        let sets = self.device.allocate_descriptor_sets(&alloc_info)?;
        self.desc_sets = [sets[0], sets[1]];

        Ok(())
    }

    /// Create upscale compute pipeline
    unsafe fn create_pipeline(&mut self) -> Result<()> {
        let shader_code = include_bytes!(concat!(env!("OUT_DIR"), "/vsr_upscale.comp.spv"));

        let shader_module_info =
            vk::ShaderModuleCreateInfo::default().code(bytemuck::cast_slice(shader_code));
        let shader_module = self
            .device
            .create_shader_module(&shader_module_info, None)?;

        let push_constant_range = vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::COMPUTE)
            .offset(0)
            .size(std::mem::size_of::<TsrPushConstants>() as u32);

        let layout_info = vk::PipelineLayoutCreateInfo::default()
            .set_layouts(std::slice::from_ref(&self.desc_layout))
            .push_constant_ranges(std::slice::from_ref(&push_constant_range));

        self.upscale_layout = self.device.create_pipeline_layout(&layout_info, None)?;

        let stage_info = vk::PipelineShaderStageCreateInfo::default()
            .stage(vk::ShaderStageFlags::COMPUTE)
            .module(shader_module)
            .name(c"main");

        let pipeline_info = vk::ComputePipelineCreateInfo::default()
            .stage(stage_info)
            .layout(self.upscale_layout);

        let pipelines = self
            .device
            .create_compute_pipelines(vk::PipelineCache::null(), &[pipeline_info], None)
            .map_err(|(_, e)| e)?;

        self.upscale_pl = pipelines[0];
        self.device.destroy_shader_module(shader_module, None);

        Ok(())
    }

    pub fn get_jitter(&mut self) -> (f32, f32) {
        self.halton.next()
    }

    pub fn jitter_projection(&mut self, projection: glam::Mat4) -> glam::Mat4 {
        let (jx, jy) = self.get_jitter();

        let dx = jx / self.render_w as f32;
        let dy = jy / self.render_h as f32;

        let mat = glam::Mat4::from_translation(glam::Vec3::new(dx * 2.0, dy * 2.0, 0.0));
        mat * projection
    }

    pub fn motion_view(&self) -> vk::ImageView {
        self.motion_v
    }

    pub fn motion_image(&self) -> vk::Image {
        self.motion_img
    }

    pub fn render_size(&self) -> (u32, u32) {
        (self.render_w, self.render_h)
    }

    pub fn display_size(&self) -> (u32, u32) {
        (self.display_w, self.display_h)
    }

    /// Get current quality preset
    pub fn quality(&self) -> VsrQuality {
        self.quality
    }

    /// Perform temporal upscaling
    ///
    /// # Safety
    /// command_buffer must be in a recording state. Image views must be valid.
    pub unsafe fn upscale(
        &mut self,
        command_buffer: vk::CommandBuffer,
        input_view: vk::ImageView,
        depth_view: vk::ImageView,
        motion_view: vk::ImageView,
        jitter: [f32; 2],
    ) -> Result<()> {
        if !self.initialized || self.upscale_pl == vk::Pipeline::null() {
            return Ok(());
        }

        let prev = (self.frame_idx % 2) as usize;
        let curr = ((self.frame_idx + 1) % 2) as usize;

        self.update_descriptor_set(curr, input_view, depth_view, motion_view, prev)?;

        // Image barrier for history and output
        let image_barriers = [
            vk::ImageMemoryBarrier::default()
                .image(self.history_imgs[prev])
                .old_layout(vk::ImageLayout::GENERAL)
                .new_layout(vk::ImageLayout::GENERAL)
                .src_access_mask(vk::AccessFlags::SHADER_WRITE)
                .dst_access_mask(vk::AccessFlags::SHADER_READ)
                .subresource_range(vk::ImageSubresourceRange {
                    aspect_mask: vk::ImageAspectFlags::COLOR,
                    level_count: 1,
                    layer_count: 1,
                    ..Default::default()
                }),
            vk::ImageMemoryBarrier::default()
                .image(self.history_imgs[curr])
                .old_layout(vk::ImageLayout::UNDEFINED)
                .new_layout(vk::ImageLayout::GENERAL)
                .src_access_mask(vk::AccessFlags::empty())
                .dst_access_mask(vk::AccessFlags::SHADER_WRITE)
                .subresource_range(vk::ImageSubresourceRange {
                    aspect_mask: vk::ImageAspectFlags::COLOR,
                    level_count: 1,
                    layer_count: 1,
                    ..Default::default()
                }),
        ];

        self.device.cmd_pipeline_barrier(
            command_buffer,
            vk::PipelineStageFlags::COMPUTE_SHADER | vk::PipelineStageFlags::TRANSFER,
            vk::PipelineStageFlags::COMPUTE_SHADER,
            vk::DependencyFlags::empty(),
            &[],
            &[],
            &image_barriers,
        );

        // Bind pipeline and descriptor set
        self.device.cmd_bind_pipeline(
            command_buffer,
            vk::PipelineBindPoint::COMPUTE,
            self.upscale_pl,
        );

        self.device.cmd_bind_descriptor_sets(
            command_buffer,
            vk::PipelineBindPoint::COMPUTE,
            self.upscale_layout,
            0,
            &[self.desc_sets[curr]],
            &[],
        );

        // Push constants
        let push_constants = TsrPushConstants {
            jitter,
            render_size: [self.render_w as f32, self.render_h as f32],
            display_size: [self.display_w as f32, self.display_h as f32],
            history_weight: 0.95,
            frame_index: self.frame_idx,
        };

        self.device.cmd_push_constants(
            command_buffer,
            self.upscale_layout,
            vk::ShaderStageFlags::COMPUTE,
            0,
            bytemuck::bytes_of(&push_constants),
        );

        // Dispatch compute (8x8 groups)
        let gx = self.display_w.div_ceil(8);
        let gy = self.display_h.div_ceil(8);
        self.device.cmd_dispatch(command_buffer, gx, gy, 1);

        Ok(())
    }

    unsafe fn update_descriptor_set(
        &self,
        set_idx: usize,
        input: vk::ImageView,
        depth: vk::ImageView,
        motion: vk::ImageView,
        hist_idx: usize,
    ) -> Result<()> {
        let sampler_info = vk::DescriptorImageInfo::default()
            .sampler(self.sampler)
            .image_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL);

        let input_info = [sampler_info.image_view(input)];
        let motion_info = [sampler_info.image_view(motion)];
        let depth_info = [sampler_info.image_view(depth)];
        let history_info = [sampler_info.image_view(self.history_vs[hist_idx])];

        let output_info = [vk::DescriptorImageInfo::default()
            .image_view(self.history_vs[(hist_idx + 1) % 2])
            .image_layout(vk::ImageLayout::GENERAL)];

        let writes = [
            vk::WriteDescriptorSet::default()
                .dst_set(self.desc_sets[set_idx])
                .dst_binding(0)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .image_info(&input_info),
            vk::WriteDescriptorSet::default()
                .dst_set(self.desc_sets[set_idx])
                .dst_binding(1)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .image_info(&motion_info),
            vk::WriteDescriptorSet::default()
                .dst_set(self.desc_sets[set_idx])
                .dst_binding(2)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .image_info(&depth_info),
            vk::WriteDescriptorSet::default()
                .dst_set(self.desc_sets[set_idx])
                .dst_binding(3)
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .image_info(&history_info),
            vk::WriteDescriptorSet::default()
                .dst_set(self.desc_sets[set_idx])
                .dst_binding(4)
                .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
                .image_info(&output_info),
        ];

        self.device.update_descriptor_sets(&writes, &[]);
        Ok(())
    }

    /// Get current upscaled output view
    pub fn output_view(&self) -> vk::ImageView {
        self.history_vs[(self.frame_idx % 2) as usize]
    }

    /// Advance to next frame
    pub fn next_frame(&mut self) {
        self.frame_idx = self.frame_idx.wrapping_add(1);
    }

    pub fn next_jitter(&mut self) -> (f32, f32) {
        if !self.initialized {
            return (0.0, 0.0);
        }
        self.halton.next()
    }

    /// Get current frame index (for ping-pong buffer selection)
    pub fn frame_index(&self) -> u32 {
        self.frame_idx
    }

    /// Destroy GPU resources
    ///
    /// # Safety
    /// Resources must not be in use.
    pub unsafe fn destroy(&mut self, allocator: &vk_mem::Allocator) {
        if !self.initialized {
            return;
        }

        if self.motion_v != vk::ImageView::null() {
            self.device.destroy_image_view(self.motion_v, None);
        }
        if let Some(mut a) = self.motion_alloc.take() {
            allocator.destroy_image(self.motion_img, &mut a);
        }

        for i in 0..2 {
            if self.history_vs[i] != vk::ImageView::null() {
                self.device.destroy_image_view(self.history_vs[i], None);
            }
            if let Some(mut a) = self.history_allocs[i].take() {
                allocator.destroy_image(self.history_imgs[i], &mut a);
            }
        }

        if self.sampler != vk::Sampler::null() {
            self.device.destroy_sampler(self.sampler, None);
        }

        if self.upscale_pl != vk::Pipeline::null() {
            self.device.destroy_pipeline(self.upscale_pl, None);
            self.device
                .destroy_pipeline_layout(self.upscale_layout, None);
        }

        if self.descriptor_pool != vk::DescriptorPool::null() {
            self.device
                .destroy_descriptor_pool(self.descriptor_pool, None);
            self.device
                .destroy_descriptor_set_layout(self.desc_layout, None);
        }

        self.initialized = false;
    }
}
