use std::sync::{atomic::AtomicU8, Arc};

use crate::{
    define_device_owned_handle,
    device::DeviceOwned,
    images::{Image, QueueOwnership},
    pipeline::{Pipeline, PipelineLayout},
    sync::{self, FenceFuture},
    util::{self, FormatExt, MutexExt},
};

use super::{Device, Queue};
use ash::{prelude::*, vk};
use parking_lot::Mutex;

#[derive(Debug)]
pub struct SingleUseCommandPool {
    device: Device,
    pool: Mutex<vk::CommandPool>,
    queue: Queue,
}

impl Drop for SingleUseCommandPool {
    fn drop(&mut self) {
        unsafe {
            self.pool.with_locked(|&pool| {
                self.device.dev().destroy_command_pool(pool, None);
            });
        }
    }
}

impl SingleUseCommandPool {
    pub fn new(device: Device, queue: Queue) -> VkResult<Arc<Self>> {
        let pool_info = vk::CommandPoolCreateInfo::default()
            .queue_family_index(queue.family())
            .flags(vk::CommandPoolCreateFlags::TRANSIENT);

        let pool = unsafe { device.dev().create_command_pool(&pool_info, None)? };

        Ok(Arc::new(Self {
            device,
            pool: Mutex::new(pool),
            queue,
        }))
    }

    #[allow(dead_code)]
    pub fn reset(&self) -> VkResult<()> {
        unsafe {
            self.pool.with_locked(|pool| {
                self.device
                    .dev()
                    .reset_command_pool(*pool, vk::CommandPoolResetFlags::empty())
            })
        }
    }
    pub fn alloc(self: &Arc<Self>) -> VkResult<SingleUseCommand> {
        SingleUseCommand::new(self.device.clone(), self.clone())
    }

    pub fn queue(&self) -> &Queue {
        &self.queue
    }

    /// get the underlying pool, bypassing the mutex
    #[allow(dead_code)]
    pub unsafe fn pool(&self) -> vk::CommandPool {
        self.pool.data_ptr().read()
    }
}

pub trait CommandBuffer: DeviceOwned<vk::CommandBuffer> {
    fn queue(&self) -> &Queue;
}

impl CommandBuffer for SingleUseCommand {
    fn queue(&self) -> &Queue {
        &self.pool.queue
    }
}

pub struct CommandList<T: CommandBuffer>(pub Vec<T>);

impl<T: CommandBuffer> CommandList<T> {
    /// all commands in list must be allocated from the same queue.
    pub fn submit<'a>(
        &'a self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
        fence: Arc<sync::Fence>,
    ) -> VkResult<FenceFuture<'a>> {
        if self.0.is_empty() {
            //exit
        }
        let buffers = self.0.iter().map(|cmd| cmd.handle()).collect::<Vec<_>>();

        let mut info = vk::SubmitInfo::default().command_buffers(&buffers);
        if let Some((sema, stage)) = wait.as_ref() {
            info = info
                .wait_dst_stage_mask(core::slice::from_ref(stage))
                .wait_semaphores(core::slice::from_ref(sema));
        }

        if let Some(signal) = signal.as_ref() {
            info = info.signal_semaphores(core::slice::from_ref(signal));
        }

        self.0[0].queue().with_locked(|queue| unsafe {
            self.0[0]
                .device()
                .dev()
                .queue_submit(queue, &[info], fence.fence())
        })?;

        Ok(FenceFuture::<'a>::new(fence))
    }
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
#[repr(u8)]
pub enum CommandBufferState {
    Initial = 0,
    Recording,
    Executable,
    Pending,
    Invalid,
}

#[derive(Debug)]
#[repr(transparent)]
struct CommandBufferState2(AtomicU8);

#[allow(dead_code)]
impl CommandBufferState2 {
    fn initial() -> Self {
        Self(AtomicU8::new(CommandBufferState::Initial as u8))
    }
    fn recording() -> Self {
        Self(AtomicU8::new(CommandBufferState::Recording as u8))
    }

    fn state(&self) -> CommandBufferState {
        let value = self.0.load(std::sync::atomic::Ordering::Relaxed);
        unsafe { *<*const _>::from(&value).cast::<CommandBufferState>() }
    }

    fn is_initial(&self) -> bool {
        self.0.load(std::sync::atomic::Ordering::Relaxed) == 0
    }
    fn set_executable(&self) {
        self.0.store(
            CommandBufferState::Executable as u8,
            std::sync::atomic::Ordering::Relaxed,
        );
    }
    fn set_recording(&self) {
        self.0.store(
            CommandBufferState::Recording as u8,
            std::sync::atomic::Ordering::Relaxed,
        );
    }
    fn set_pending(&self) {
        self.0.store(
            CommandBufferState::Pending as u8,
            std::sync::atomic::Ordering::Relaxed,
        );
    }
    fn set_invalid(&self) {
        self.0.store(
            CommandBufferState::Invalid as u8,
            std::sync::atomic::Ordering::Relaxed,
        );
    }
}

define_device_owned_handle! {
    #[derive(Debug)]
    pub SingleUseCommand(vk::CommandBuffer) {
        pool: Arc<SingleUseCommandPool>,
        state: CommandBufferState2,
    } => |this|  unsafe {
            this.pool
                .pool
                .with_locked(|&pool| this.device().dev().free_command_buffers(pool, &[this.handle()]))
    }
}

impl !Sync for SingleUseCommand {}

impl SingleUseCommand {
    pub fn new(device: Device, pool: Arc<SingleUseCommandPool>) -> VkResult<Self> {
        let buffer = unsafe {
            let buffer = pool.pool.with_locked(|pool| {
                let alloc_info = vk::CommandBufferAllocateInfo::default()
                    .command_buffer_count(1)
                    .command_pool(*pool)
                    .level(vk::CommandBufferLevel::PRIMARY);
                Ok(device.dev().allocate_command_buffers(&alloc_info)?[0])
            })?;

            device.dev().begin_command_buffer(
                buffer,
                &vk::CommandBufferBeginInfo::default()
                    .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
            )?;

            buffer
        };

        Self::construct(device, buffer, None, pool, CommandBufferState2::recording())
    }

    pub fn state(&self) -> CommandBufferState {
        self.state.state()
    }

    pub fn end(&self) -> VkResult<()> {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.inner.dev().dev().end_command_buffer(self.handle())?;
        }
        self.state.set_executable();
        Ok(())
    }

    /// Safety: commandbuffer must not be accessed from multiple threads at the same time
    pub unsafe fn buffer(&self) -> vk::CommandBuffer {
        self.handle()
    }

    //pub fn copy_buffer_to_image(&self, image: &Image2D, buffer: &Buffer, )

    pub fn image_barrier(
        &self,
        image: vk::Image,
        aspects: vk::ImageAspectFlags,
        src_stage: vk::PipelineStageFlags2,
        src_access: vk::AccessFlags2,
        dst_stage: vk::PipelineStageFlags2,
        dst_access: vk::AccessFlags2,
        old_layout: vk::ImageLayout,
        new_layout: vk::ImageLayout,
        queue_ownership_op: Option<QueueOwnership>,
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        let (src_family, dst_family) = queue_ownership_op
            .map(|t| (t.src, t.dst))
            .unwrap_or((vk::QUEUE_FAMILY_IGNORED, vk::QUEUE_FAMILY_IGNORED));

        let barrier = vk::ImageMemoryBarrier2::default()
            .image(image)
            .subresource_range(
                vk::ImageSubresourceRange::default()
                    .aspect_mask(aspects)
                    .base_mip_level(0)
                    .base_array_layer(0)
                    .level_count(vk::REMAINING_MIP_LEVELS)
                    .layer_count(vk::REMAINING_ARRAY_LAYERS),
            )
            .src_stage_mask(src_stage)
            .src_access_mask(src_access)
            .dst_stage_mask(dst_stage)
            .dst_access_mask(dst_access)
            .dst_queue_family_index(dst_family)
            .src_queue_family_index(src_family)
            .old_layout(old_layout)
            .new_layout(new_layout);

        unsafe {
            self.device().dev().cmd_pipeline_barrier2(
                self.handle(),
                &vk::DependencyInfo::default()
                    .dependency_flags(vk::DependencyFlags::BY_REGION)
                    .image_memory_barriers(core::slice::from_ref(&barrier)),
            );
        }
    }

    pub fn blit_images(
        &self,
        src: &Image,
        src_region: util::Rect2D,
        dst: &Image,
        dst_region: util::Rect2D,
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_blit_image(
                self.buffer(),
                src.image(),
                vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
                dst.image(),
                vk::ImageLayout::TRANSFER_DST_OPTIMAL,
                &[vk::ImageBlit::default()
                    .src_subresource(
                        vk::ImageSubresourceLayers::default()
                            .aspect_mask(vk::ImageAspectFlags::COLOR)
                            .layer_count(1),
                    )
                    .dst_subresource(
                        vk::ImageSubresourceLayers::default()
                            .aspect_mask(vk::ImageAspectFlags::COLOR)
                            .layer_count(1),
                    )
                    .src_offsets(src_region.into_offsets_3d())
                    .dst_offsets(dst_region.into_offsets_3d())],
                vk::Filter::LINEAR,
            );
        }
    }

    pub fn copy_buffer_to_image(
        &self,
        buffer: vk::Buffer,
        image: vk::Image,
        layout: vk::ImageLayout,
        regions: &[vk::BufferImageCopy],
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_copy_buffer_to_image(
                self.handle(),
                buffer,
                image,
                layout,
                regions,
            );
        }
    }

    pub fn copy_buffers(&self, src: vk::Buffer, dst: vk::Buffer, regions: &[vk::BufferCopy]) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_copy_buffer(self.handle(), src, dst, regions);
        }
    }
    #[allow(dead_code)]
    pub fn copy_images(
        &self,
        src: vk::Image,
        src_layout: vk::ImageLayout,
        dst: vk::Image,
        dst_layout: vk::ImageLayout,
        regions: &[vk::ImageCopy],
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_copy_image(
                self.handle(),
                src,
                src_layout,
                dst,
                dst_layout,
                regions,
            );
        }
    }

    pub fn clear_color_image(
        &self,
        image: vk::Image,
        format: vk::Format,
        layout: vk::ImageLayout,
        color: crate::Rgba,
        subresources: &[vk::ImageSubresourceRange],
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        let clear_colors = match format.get_component_kind() {
            crate::util::FormatComponentKind::Float => vk::ClearColorValue {
                float32: color.into_f32(),
            },
            crate::util::FormatComponentKind::UInt => vk::ClearColorValue {
                uint32: color.into_u32(),
            },
            crate::util::FormatComponentKind::SInt => vk::ClearColorValue {
                int32: color.into_i32(),
            },
        };

        unsafe {
            self.device().dev().cmd_clear_color_image(
                self.handle(),
                image,
                layout,
                &clear_colors,
                subresources,
            );
        }
    }

    pub fn begin_rendering(&self, rendering_info: vk::RenderingInfo<'_>) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_begin_rendering(self.buffer(), &rendering_info);
        }
    }

    pub fn set_viewport(&self, viewports: &[vk::Viewport]) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_set_viewport(self.buffer(), 0, viewports);
        }
    }
    pub fn set_scissors(&self, scissors: &[vk::Rect2D]) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_set_scissor(self.buffer(), 0, scissors);
        }
    }
    pub fn push_constants(
        &self,
        layout: &PipelineLayout,
        stage: vk::ShaderStageFlags,
        offset: u32,
        bytes: &[u8],
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_push_constants(
                self.handle(),
                layout.handle(),
                stage,
                offset,
                bytes,
            );
        }
    }
    pub fn bind_pipeline(&self, pipeline: &Pipeline) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_bind_pipeline(
                self.buffer(),
                pipeline.bind_point(),
                pipeline.handle(),
            );
        }
    }
    pub fn bind_vertex_buffers(&self, buffers: &[vk::Buffer], offsets: &[u64]) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_bind_vertex_buffers(self.buffer(), 0, buffers, offsets);
        }
    }
    pub fn bind_indices(&self, buffer: vk::Buffer, offset: u64, kind: vk::IndexType) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device()
                .dev()
                .cmd_bind_index_buffer(self.buffer(), buffer, offset, kind);
        }
    }

    pub fn bind_descriptor_sets(
        &self,
        layout: &PipelineLayout,
        bind_point: vk::PipelineBindPoint,
        descriptor_sets: &[vk::DescriptorSet],
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        use crate::device::DeviceOwned;
        unsafe {
            self.device().dev().cmd_bind_descriptor_sets(
                self.buffer(),
                bind_point,
                layout.handle(),
                0,
                descriptor_sets,
                &[],
            );
        }
    }

    #[allow(dead_code)]
    pub fn draw_indexed(
        &self,
        indices: u32,
        instances: u32,
        index_offset: u32,
        vertex_offset: i32,
        instance_offset: u32,
    ) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_draw_indexed(
                self.buffer(),
                indices,
                instances,
                index_offset,
                vertex_offset,
                instance_offset,
            );
        }
    }

    pub fn draw_indexed_indirect(&self, buffer: vk::Buffer, offset: u64, count: u32, stride: u32) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_draw_indexed_indirect(
                self.buffer(),
                buffer,
                offset,
                count,
                stride,
            );
        }
    }

    pub fn end_rendering(&self) {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe {
            self.device().dev().cmd_end_rendering(self.buffer());
        }
    }

    pub fn submit_fence(
        &self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
        fence: Option<vk::Fence>,
    ) -> VkResult<()> {
        assert_eq!(self.state(), CommandBufferState::Recording);
        unsafe { self.device().dev().end_command_buffer(self.handle())? };

        let buffers = [self.handle()];
        let mut submit_info = vk::SubmitInfo::default().command_buffers(&buffers);

        if let Some(semaphore) = signal.as_ref() {
            // SAFETY: T and [T;1] have the same layout
            submit_info = submit_info.signal_semaphores(unsafe {
                core::mem::transmute::<&vk::Semaphore, &[vk::Semaphore; 1]>(semaphore)
            });
        }
        if let Some((semaphore, stage)) = wait.as_ref() {
            submit_info = submit_info
                .wait_semaphores(core::slice::from_ref(semaphore))
                .wait_dst_stage_mask(core::slice::from_ref(stage));
        }

        let fence = fence.unwrap_or(vk::Fence::null());
        self.pool.queue().with_locked(|queue| unsafe {
            self.device()
                .dev()
                .queue_submit(queue, &[submit_info], fence)
        })?;
        tracing::trace!(
            "submitted queue {:?} and fence {:?}",
            self.pool.queue(),
            fence
        );

        Ok(())
    }

    pub fn submit_async<'a>(
        &'a self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
        fence: Arc<sync::Fence>,
    ) -> VkResult<FenceFuture<'a>> {
        self.submit_fence(wait, signal, Some(fence.fence()))?;

        Ok(FenceFuture::new(fence))
    }

    #[allow(dead_code)]
    pub fn submit_blocking(
        self,
        wait: Option<(vk::Semaphore, vk::PipelineStageFlags)>,
        signal: Option<vk::Semaphore>,
    ) -> VkResult<()> {
        let fence = Arc::new(sync::Fence::create(self.device().clone())?);
        let future = self.submit_async(wait, signal, fence)?;
        future.block()?;
        Ok(())
    }
}
