//! This module defines the commands that can be recorded in the render
//! graph.
//! Commands can be allocated, and then recorded into a command
//! buffer.
//! the lifecycle of a command is thus as follows:
//! 1. A command is constructed, with its parameters and resource dependencies.
//! 2. The command is inserted into a command list. At this stage, the command's side effects are recorded and barriers are automatically inserted as required.
//! 3. When the command list is recorded into a command buffer, the command's `apply` method is called, which translates the command into actual GPU commands using the provided `CommandRecorder`.
//!
//! Each command declares its side effects, such as what resources it
//! reads from or write to, and which pipeline stages it affects. Commands are
//! recorded sequentially, in the order they were inserted into a command
//! buffer. During command insertion

use crate::{
    device::DeviceOwned,
    render_graph::recorder::{CommandRecorder, SideEffectMap},
};

use super::resources::*;

pub struct ImportResource<T: Resource> {
    pub resource: T,
    pub access: vk::AccessFlags2,
    pub stage: vk::PipelineStageFlags2,
    pub layout: Option<vk::ImageLayout>,
}

impl<T: Resource> Command for ImportResource<T> {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.resource
            .side_effect(map.reborrow().into_side_effect_map2(
                self.stage,
                self.access,
                self.layout,
            ));
    }

    fn apply(self, _recorder: &mut CommandRecorder) {
        // No actual GPU commands needed for import, as it's just a logical operation
    }
}

pub struct CopyBuffers {
    pub src: BufferSlice,
    pub dst: BufferSlice,
}

pub struct CopyTextures {
    pub src: TextureRegion,
    pub dst: TextureRegion,
}

pub struct CopyBufferToTexture {
    pub src: BufferSlice,
    pub rows: Option<BufferRows>,
    pub dst: TextureRegion,
}

pub struct CopyTextureToBuffer {
    pub src: TextureRegion,
    pub rows: Option<BufferRows>,
    pub dst: BufferSlice,
}

pub struct CopyBuffersToTextures {
    pub src: BufferSlice,
    pub src_rows: Vec<BufferRows>,
    pub dst: Texture,
    pub src_ranges: Vec<(TextureRegion, vk::Offset3D, vk::Extent3D)>,
}

pub struct Copy<T: Resource, U: Resource> {
    pub src: T,
    pub dst: U,
}

fn copy_side_effects_inner<T, U>(src: &T, dst: &U, mut map: SideEffectMap)
where
    T: Resource,
    U: Resource,
{
    src.side_effect(map.reborrow().into_side_effect_map2(
        vk::PipelineStageFlags2::COPY,
        vk::AccessFlags2::TRANSFER_READ,
        Some(vk::ImageLayout::TRANSFER_SRC_OPTIMAL),
    ));
    dst.side_effect(map.reborrow().into_side_effect_map2(
        vk::PipelineStageFlags2::COPY,
        vk::AccessFlags2::TRANSFER_WRITE,
        Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL),
    ));
}

impl Command for CopyBufferToTexture {
    fn side_effects(&self, map: SideEffectMap) {
        copy_side_effects_inner(&self.src, &self.dst, map)
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        // it doesn't really make sense to copy the same data into multiple mip
        // levels, considering the mip extents are different for each level.
        debug_assert_eq!(self.dst.range.mip_levels.len(), 1);

        let rows = self.rows.unwrap_or_default();
        let regions = &[vk::BufferImageCopy2::default()
            .buffer_offset(self.src.range.offset + rows.offset)
            .buffer_row_length(rows.row_size)
            .buffer_image_height(rows.row_count)
            .image_offset(self.dst.offset)
            .image_extent(self.dst.extent)
            .image_subresource(
                vk::ImageSubresourceLayers::default()
                    .aspect_mask(self.dst.range.aspect)
                    .mip_level(self.dst.range.mip_levels.start())
                    .base_array_layer(self.dst.range.array_layers.start())
                    .layer_count(self.dst.range.array_layers.len()),
            )];

        let info = vk::CopyBufferToImageInfo2::default()
            .src_buffer(recorder.get_buffer_handle(self.src.id()))
            .dst_image(recorder.get_image_handle(self.dst.id()))
            .dst_image_layout(recorder.get_image_layout(self.dst.id()))
            .regions(regions);

        unsafe {
            dev.cmd_copy_buffer_to_image2(cmd.raw(), &info);
        }
    }
}

impl Command for Copy<BufferRowSlices, TextureRegions> {
    fn side_effects(&self, map: SideEffectMap) {
        self.side_effects_inner(map)
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        let regions = &self
            .src
            .rows
            .iter()
            .zip(self.dst.ranges.iter())
            .map(|(src_row, dst_range)| {
                vk::BufferImageCopy2::default()
                    .buffer_offset(src_row.offset)
                    .buffer_row_length(src_row.row_size)
                    .buffer_image_height(src_row.row_count)
                    .image_offset(dst_range.offset())
                    .image_extent(dst_range.extent())
                    .image_subresource(
                        vk::ImageSubresourceLayers::default()
                            .aspect_mask(dst_range.aspect)
                            .mip_level(dst_range.mip_levels.start())
                            .base_array_layer(dst_range.array_layers.start())
                            .layer_count(dst_range.array_layers.len()),
                    )
            })
            .collect::<Vec<_>>();

        let info = vk::CopyBufferToImageInfo2::default()
            .src_buffer(recorder.get_buffer_handle(self.src.id()))
            .dst_image(recorder.get_image_handle(self.dst.id()))
            .dst_image_layout(recorder.get_image_layout(self.dst.id()))
            .regions(regions);

        unsafe {
            dev.cmd_copy_buffer_to_image2(cmd.raw(), &info);
        }
    }
}

impl Command for Copy<TextureRegion, BufferRowSlice> {
    fn side_effects(&self, map: SideEffectMap) {
        self.side_effects_inner(map)
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        let regions = &[vk::BufferImageCopy2::default()
            .buffer_offset(self.dst.row.offset)
            .buffer_row_length(self.dst.row.row_size)
            .buffer_image_height(self.dst.row.row_count)
            .image_offset(self.src.range.offset())
            .image_extent(self.src.range.extent())
            .image_subresource(
                vk::ImageSubresourceLayers::default()
                    .aspect_mask(self.src.range.aspect)
                    .mip_level(self.src.range.mip_levels.start())
                    .base_array_layer(self.src.range.array_layers.start())
                    .layer_count(self.src.range.array_layers.len()),
            )];

        let info = vk::CopyImageToBufferInfo2::default()
            .dst_buffer(recorder.get_buffer_handle(self.dst.id()))
            .src_image(recorder.get_image_handle(self.src.id()))
            .src_image_layout(recorder.get_image_layout(self.src.id()))
            .regions(regions);

        unsafe {
            dev.cmd_copy_image_to_buffer2(cmd.raw(), &info);
        }
    }
}
impl Command for Copy<TextureRegion, TextureRegion> {
    fn side_effects(&self, map: SideEffectMap) {
        self.side_effects_inner(map)
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        let regions = &[vk::ImageCopy2::default()
            .extent(self.src.range.extent())
            .dst_offset(self.dst.range.offset())
            .src_offset(self.src.range.offset())
            .dst_subresource(
                vk::ImageSubresourceLayers::default()
                    .aspect_mask(self.dst.range.aspect)
                    .mip_level(self.dst.range.mip_levels.start())
                    .base_array_layer(self.dst.range.array_layers.start())
                    .layer_count(self.dst.range.array_layers.len()),
            )
            .src_subresource(
                vk::ImageSubresourceLayers::default()
                    .aspect_mask(self.src.range.aspect)
                    .mip_level(self.src.range.mip_levels.start())
                    .base_array_layer(self.src.range.array_layers.start())
                    .layer_count(self.src.range.array_layers.len()),
            )];

        let info = vk::CopyImageInfo2::default()
            .dst_image(recorder.get_image_handle(self.dst.id()))
            .dst_image_layout(recorder.get_image_layout(self.dst.id()))
            .src_image(recorder.get_image_handle(self.src.id()))
            .src_image_layout(recorder.get_image_layout(self.src.id()))
            .regions(regions);

        unsafe {
            dev.cmd_copy_image2(cmd.raw(), &info);
        }
    }
}
impl Command for Copy<BufferSlice, BufferSlice> {
    fn side_effects(&self, map: SideEffectMap) {
        self.side_effects_inner(map)
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        debug_assert_eq!(self.src.range.size, self.dst.range.size);

        let regions = &[vk::BufferCopy2::default()
            .dst_offset(self.dst.range.offset)
            .src_offset(self.src.range.offset)
            .size(self.src.range.size)];

        let info = vk::CopyBufferInfo2::default()
            .dst_buffer(recorder.get_buffer_handle(self.dst.id()))
            .src_buffer(recorder.get_buffer_handle(self.src.id()))
            .regions(regions);

        unsafe {
            dev.cmd_copy_buffer2(cmd.raw(), &info);
        }
    }
}

pub struct ClearTexture {
    pub dst: Write<TextureRegion>,
    pub clear_value: [f32; 4],
}

impl Command for ClearTexture {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.dst.side_effect(map.reborrow().into_side_effect_map2(
            vk::PipelineStageFlags2::TRANSFER,
            vk::AccessFlags2::TRANSFER_WRITE,
            Some(vk::ImageLayout::TRANSFER_DST_OPTIMAL),
        ));
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct UpdateBuffer {
    pub dst: Write<BufferSlice>,
    pub data: Vec<u8>,
}

impl Command for UpdateBuffer {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.dst.side_effect(map.reborrow().into_side_effect_map2(
            vk::PipelineStageFlags2::TRANSFER,
            vk::AccessFlags2::TRANSFER_WRITE,
            None,
        ));
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct BeginRendering {
    pub color_attachments: Vec<TextureRegion>,
    pub depth_attachment: Option<TextureRegion>,
    pub stencil_attachment: Option<TextureRegion>,
    pub area: (u32, u32),
    pub layers: u32,
}

impl Command for BeginRendering {
    fn side_effects(&self, mut map: SideEffectMap) {
        for attachment in &self.color_attachments {
            // The dependencies declared here need to be barrier'd before the
            // render pass: the spec says pipelineBarrier inside of a render
            // pass instance referencing framebuffer-local stages must contain
            // only framebuffer-local stages.

            // TODO: consider loadop and storeop?
            attachment.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT,
                vk::AccessFlags2::COLOR_ATTACHMENT_READ,
                Some(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL),
            ));
        }
        if let Some(depth) = &self.depth_attachment {
            depth.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::EARLY_FRAGMENT_TESTS,
                vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_READ,
                Some(vk::ImageLayout::DEPTH_ATTACHMENT_OPTIMAL),
            ));
        }
        if let Some(stencil) = &self.stencil_attachment {
            stencil.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::EARLY_FRAGMENT_TESTS,
                vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_READ,
                Some(vk::ImageLayout::STENCIL_ATTACHMENT_OPTIMAL),
            ));
        }
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct EndRendering;

impl Command for EndRendering {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but ends the render pass
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct BindPipeline(pub vk::Pipeline, pub vk::PipelineBindPoint);

impl Command for BindPipeline {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but affects the pipeline state
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        unsafe {
            dev.cmd_bind_pipeline(cmd.raw(), self.1, self.0);
        }
    }
}

pub struct BindVertexBuffers {
    pub buffers: Vec<BufferSlice>,
    pub first_binding: u32,
}

impl Command for BindVertexBuffers {
    fn side_effects(&self, mut map: SideEffectMap) {
        for buffer in &self.buffers {
            buffer.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::VERTEX_ATTRIBUTE_INPUT,
                vk::AccessFlags2::VERTEX_ATTRIBUTE_READ,
                None,
            ));
        }
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        let buffers = self
            .buffers
            .iter()
            .map(|buffer| recorder.get_buffer_handle(buffer.id()))
            .collect::<Vec<_>>();

        let offsets = self
            .buffers
            .iter()
            .map(|buffer| buffer.range.offset)
            .collect::<Vec<_>>();

        unsafe {
            dev.cmd_bind_vertex_buffers(cmd.raw(), self.first_binding, &buffers, &offsets);
        }
    }
}

pub struct BindIndexBuffer(pub BufferSlice, pub IndexFormat);

impl Command for BindIndexBuffer {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.0.side_effect(map.reborrow().into_side_effect_map2(
            vk::PipelineStageFlags2::INDEX_INPUT,
            vk::AccessFlags2::INDEX_READ,
            None,
        ));
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        let buffer = recorder.get_buffer_handle(self.0.id());
        let index_type = match self.1 {
            IndexFormat::Uint8 => vk::IndexType::UINT8_KHR,
            IndexFormat::Uint16 => vk::IndexType::UINT16,
            IndexFormat::Uint32 => vk::IndexType::UINT32,
        };

        unsafe {
            dev.cmd_bind_index_buffer(cmd.raw(), buffer, self.0.range.offset, index_type);
        }
    }
}

pub struct BindDescriptorSets {
    pub first_set: u32,
    pub pipeline_layout: vk::PipelineLayout,
    pub bind_point: vk::PipelineBindPoint,
    pub sets: Vec<vk::DescriptorSet>,
}

impl Command for BindDescriptorSets {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but affects the descriptor set state
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        unsafe {
            dev.cmd_bind_descriptor_sets(
                cmd.raw(),
                self.bind_point,
                self.pipeline_layout,
                self.first_set,
                &self.sets,
                &[],
            );
        }
    }
}

pub struct SetViewport {
    pub x: f32,
    pub y: f32,
    pub width: f32,
    pub height: f32,
    pub min_depth: f32,
    pub max_depth: f32,
}

impl Command for SetViewport {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but affects the viewport state
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct SetScissor {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

impl Command for SetScissor {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but affects the scissor state
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct PushConstants {
    pub pipeline_layout: vk::PipelineLayout,
    pub stage_flags: vk::ShaderStageFlags,
    pub offset: u32,
    pub data: Vec<u8>,
}

impl Command for PushConstants {
    fn side_effects(&self, _map: SideEffectMap) {
        // No resource access, but affects the push constant state
    }

    fn apply(self, recorder: &mut CommandRecorder) {
        let cmd = recorder.cmd_buffer();
        let dev = &cmd.device().raw;

        unsafe {
            dev.cmd_push_constants(
                cmd.raw(),
                self.pipeline_layout,
                self.stage_flags,
                self.offset,
                &self.data,
            );
        }
    }
}

pub struct DrawData {
    pub vertex_count: u32,
    pub instance_count: u32,
    pub first_vertex: u32,
    pub first_instance: u32,
}

pub struct DrawIndexedData {
    pub index_count: u32,
    pub instance_count: u32,
    pub first_index: u32,
    pub vertex_offset: i32,
    pub first_instance: u32,
}

pub struct DrawIndirectData {
    pub indirect_buffer: BufferSlice,
    pub count: u32,
    pub stride: u32,
}

pub struct DrawIndexedIndirectData {
    pub indirect_buffer: BufferSlice,
    pub count: u32,
    pub stride: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum IndexFormat {
    Uint8,
    Uint16,
    Uint32,
}

pub struct Draw<T: IsDrawData> {
    pub data: T,
    pub vertex_buffers: Vec<BufferSlice>,
    pub index_buffer: Option<(BufferSlice, IndexFormat)>,
}

impl<T: IsDrawData> Command for Draw<T> {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.data.side_effects(map.reborrow());
        for vertex_buffer in &self.vertex_buffers {
            vertex_buffer.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::VERTEX_ATTRIBUTE_INPUT,
                vk::AccessFlags2::VERTEX_ATTRIBUTE_READ,
                None,
            ));
        }
        if let Some((index_buffer, _)) = &self.index_buffer {
            index_buffer.side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::INDEX_INPUT,
                vk::AccessFlags2::INDEX_READ,
                None,
            ));
        }
    }

    fn apply(self, _recorder: &mut CommandRecorder) {}
}

pub struct PipelineBarrier;

impl Command for PipelineBarrier {
    fn side_effects(&self, _map: SideEffectMap) {}

    fn apply(self, _recorder: &mut CommandRecorder) {
        todo!()
    }
}

pub trait Command: Sized {
    fn side_effects(&self, map: SideEffectMap);
    fn apply(self, recorder: &mut CommandRecorder);
}

mod sealed {
    pub trait IsDrawData {
        fn side_effects(&self, _map: super::SideEffectMap) {}
    }
    pub trait InsideRenderPass {}
    pub trait OutsideRenderPass {}
}
use ash::vk;
pub(super) use sealed::*;

impl IsDrawData for DrawData {}
impl IsDrawData for DrawIndexedData {}
impl IsDrawData for DrawIndirectData {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.indirect_buffer
            .side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::DRAW_INDIRECT,
                vk::AccessFlags2::INDIRECT_COMMAND_READ,
                None,
            ));
    }
}
impl IsDrawData for DrawIndexedIndirectData {
    fn side_effects(&self, mut map: SideEffectMap) {
        self.indirect_buffer
            .side_effect(map.reborrow().into_side_effect_map2(
                vk::PipelineStageFlags2::DRAW_INDIRECT,
                vk::AccessFlags2::INDIRECT_COMMAND_READ,
                None,
            ));
    }
}

impl OutsideRenderPass for BeginRendering {}
impl InsideRenderPass for EndRendering {}
impl<T: Resource, U: Resource> OutsideRenderPass for Copy<T, U> {}
impl OutsideRenderPass for ClearTexture {}
impl OutsideRenderPass for UpdateBuffer {}

impl InsideRenderPass for BindPipeline {}
impl OutsideRenderPass for BindPipeline {}
impl InsideRenderPass for BindVertexBuffers {}
impl OutsideRenderPass for BindVertexBuffers {}
impl InsideRenderPass for BindIndexBuffer {}
impl OutsideRenderPass for BindIndexBuffer {}
impl InsideRenderPass for BindDescriptorSets {}
impl OutsideRenderPass for BindDescriptorSets {}
impl OutsideRenderPass for PushConstants {}

impl InsideRenderPass for SetViewport {}
impl InsideRenderPass for SetScissor {}
impl InsideRenderPass for PushConstants {}
impl<T: IsDrawData> InsideRenderPass for Draw<T> {}
impl<T: IsDrawData> InsideRenderPass for T {}