use bitflags::bitflags;
use parking_lot::Mutex;
use raw_window_handle::RawDisplayHandle;
use std::{collections::HashMap, ops::Deref, sync::Arc};
use tinyvec::{ArrayVec, array_vec};

use ash::vk;

use crate::{Instance, PhysicalDeviceInfo, Result};

#[derive(Debug, Clone)]
pub struct Queue {
    inner: Arc<QueueInner>,
}

impl Deref for Queue {
    type Target = QueueInner;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl Queue {
    pub fn from_vk_queue(raw: vk::Queue, family: QueueFamily) -> Self {
        Self {
            inner: Arc::new(QueueInner {
                raw,
                family,
                lock: Mutex::new(()),
            }),
        }
    }
    pub fn with_locked<F, R>(&self, f: F) -> R
    where
        F: FnOnce(&Queue) -> R,
    {
        let _lock = self.inner.lock.lock();
        f(&self)
    }
}

#[derive(Debug)]
pub struct QueueInner {
    pub(crate) raw: vk::Queue,
    pub(crate) family: QueueFamily,
    pub(crate) lock: Mutex<()>,
}

impl QueueInner {
    pub fn raw(&self) -> vk::Queue {
        self.raw
    }
    pub fn family_index(&self) -> u32 {
        self.family.index
    }
}

#[derive(Debug, Clone, Copy)]
pub struct QueueFamily {
    pub index: u32,
    pub count: u32,
    pub flags: QueueFlags,
    pub granularity: TransferGranuality,
    pub timestamp_valid_bits: u32,
}

bitflags! {
#[derive(Debug, Clone, Copy)]
    pub struct QueueFlags: u16 {
        const GRAPHICS = 0b1;
        const COMPUTE = 0b10;
        const TRANSFER = 0b100;
        const PRESENT = 0b1000;
        const ENCODE = 0b1_0000;
        const DECODE = 0b10_0000;
    }
}

impl QueueFlags {
    pub fn from_vk_flags(vk_flags: vk::QueueFlags) -> Self {
        let mut flags = QueueFlags::empty();
        if vk_flags.contains(vk::QueueFlags::GRAPHICS) {
            flags |= QueueFlags::GRAPHICS;
        }
        if vk_flags.contains(vk::QueueFlags::COMPUTE) {
            flags |= QueueFlags::COMPUTE;
        }
        if vk_flags.contains(vk::QueueFlags::TRANSFER) {
            flags |= QueueFlags::TRANSFER;
        }
        if vk_flags.contains(vk::QueueFlags::VIDEO_ENCODE_KHR) {
            flags |= QueueFlags::ENCODE;
        }
        if vk_flags.contains(vk::QueueFlags::VIDEO_DECODE_KHR) {
            flags |= QueueFlags::DECODE;
        }
        flags
    }
}

#[derive(Debug, Clone, Copy)]
pub enum TransferGranuality {
    FullImage,
    Unrestricted,
    Tiled(u32, u32, u32),
}

impl TransferGranuality {
    pub fn from_vk_granularity(vk_granularity: vk::Extent3D) -> Self {
        if vk_granularity.width == 0 || vk_granularity.height == 0 || vk_granularity.depth == 0 {
            TransferGranuality::FullImage
        } else if vk_granularity.width == 1
            && vk_granularity.height == 1
            && vk_granularity.depth == 1
        {
            TransferGranuality::Unrestricted
        } else {
            TransferGranuality::Tiled(
                vk_granularity.width,
                vk_granularity.height,
                vk_granularity.depth,
            )
        }
    }
}

#[derive(Debug)]
pub struct DeviceQueues {
    graphics: Queue,
    compute: Queue,
    transfer: Queue,
}

impl DeviceQueues {
    pub fn graphics(&self) -> &Queue {
        &self.graphics
    }
    pub fn compute(&self) -> &Queue {
        &self.compute
    }
    pub fn transfer(&self) -> &Queue {
        &self.transfer
    }

    pub fn family_indices(&self, flags: crate::device::QueueFlags) -> ArrayVec<[u32; 4]> {
        let mut indices = array_vec!([u32; 4]);
        use crate::device::QueueFlags as QF;
        if flags.intersects(QF::GRAPHICS | QF::PRESENT) {
            indices.push(self.graphics.family.index);
        }
        if flags.contains(QF::ASYNC_COMPUTE) {
            indices.push(self.compute.family.index);
        }
        if flags.contains(QF::TRANSFER) {
            indices.push(self.transfer.family.index);
        }
        let len = indices.partition_dedup().0.len();
        _ = indices.drain(len..);
        indices
    }

    pub fn swapchain_family_indices(&self) -> &[u32] {
        core::slice::from_ref(&self.graphics.family.index)
    }

    /// # Safety
    ///
    /// The caller must ensure that the queues aren't already locked when calling this function.
    pub unsafe fn lock(&self) {
        core::mem::forget((
            self.graphics.inner.lock.lock(),
            self.compute.inner.lock.lock(),
            self.transfer.inner.lock.lock(),
        ));
    }

    /// # Safety
    ///
    /// The caller must have acquired and have logical ownership of the lock on the queues.
    pub unsafe fn unlock(&self) {
        unsafe {
            self.graphics.inner.lock.force_unlock();
            self.compute.inner.lock.force_unlock();
            self.transfer.inner.lock.force_unlock();
        }
    }
}

#[derive(Debug)]
pub struct DeviceQueueInfos {
    graphics: QueueFamily,
    compute: Option<QueueFamily>,
    transfer: Option<QueueFamily>,
}

impl DeviceQueueInfos {
    const PRIORITIES: [f32; 4] = [1.0, 1.0, 1.0, 1.0];
    pub fn into_create_infos(&self) -> Vec<vk::DeviceQueueCreateInfo<'_>> {
        let families = self.queue_family_indices();

        families
            .into_iter()
            .map(|(index, count)| {
                vk::DeviceQueueCreateInfo::default()
                    .queue_family_index(index)
                    .queue_priorities(&Self::PRIORITIES[..count as usize])
            })
            .collect()
    }

    fn queue_family_indices(&self) -> HashMap<u32, u32> {
        let mut families = HashMap::new();

        families.insert(self.graphics.index, 1);
        if let Some(compute) = self.compute {
            *families.entry(compute.index).or_insert(0) += 1;
        }
        if let Some(transfer) = self.transfer {
            *families.entry(transfer.index).or_insert(0) += 1;
        }

        families
    }

    pub fn retrieve_queues(self, dev: &ash::Device) -> DeviceQueues {
        let families = self.queue_family_indices();

        let mut queues = families
            .into_iter()
            .map(|(queue_family_index, count)| {
                let queues = (0..count)
                    .map(|queue_index| unsafe {
                        dev.get_device_queue(queue_family_index, queue_index)
                    })
                    .collect();
                (queue_family_index, queues)
            })
            .collect::<HashMap<u32, Vec<vk::Queue>>>();

        let graphics = Queue::from_vk_queue(
            queues.get_mut(&self.graphics.index).unwrap().pop().unwrap(),
            self.graphics,
        );
        let compute = self
            .compute
            .map(|compute| {
                Queue::from_vk_queue(
                    queues.get_mut(&compute.index).unwrap().pop().unwrap(),
                    compute,
                )
            })
            .unwrap_or_else(|| graphics.clone());
        let transfer = self
            .transfer
            .map(|transfer| {
                Queue::from_vk_queue(
                    queues.get_mut(&transfer.index).unwrap().pop().unwrap(),
                    transfer,
                )
            })
            .unwrap_or_else(|| graphics.clone());

        DeviceQueues {
            graphics,
            compute,
            transfer,
        }
    }

    pub fn select_queue_families(
        instance: &Instance,
        pdev: &PhysicalDeviceInfo,
        display_handle: Option<RawDisplayHandle>,
    ) -> Result<Self> {
        let queue_families = unsafe {
            instance
                .inner
                .raw
                .get_physical_device_queue_family_properties(pdev.pdev)
        };

        let queue_families = queue_families
            .into_iter()
            .enumerate()
            .map(|(index, props)| {
                let mut flags = QueueFlags::from_vk_flags(props.queue_flags);
                if let Some(display_handle) = display_handle {
                    if instance.query_presentation_support(pdev.pdev, index as u32, display_handle)
                    {
                        flags |= QueueFlags::PRESENT;
                    }
                }

                QueueFamily {
                    index: index as u32,
                    count: props.queue_count,
                    flags,
                    granularity: TransferGranuality::from_vk_granularity(
                        props.min_image_transfer_granularity,
                    ),
                    timestamp_valid_bits: props.timestamp_valid_bits,
                }
            })
            .collect::<Vec<_>>();

        let mut queue_counts = queue_families
            .iter()
            .map(|family| family.count)
            .collect::<Vec<_>>();

        let graphics = queue_families
            .iter()
            .enumerate()
            .find(|(_, family)| {
                family
                    .flags
                    .contains(QueueFlags::GRAPHICS | QueueFlags::COMPUTE)
            })
            .map(|(i, family)| {
                queue_counts[i] -= 1;
                *family
            })
            .expect("No graphics queue family found");

        assert!(
            graphics.flags.contains(QueueFlags::PRESENT),
            "Selected graphics queue family does not support presentation"
        );

        let compute = queue_families
            .iter()
            .enumerate()
            .find(|(i, family)| queue_counts[*i] > 0 && family.flags.contains(QueueFlags::COMPUTE))
            .map(|(i, family)| {
                queue_counts[i] -= 1;
                *family
            });

        let transfer = queue_families
            .iter()
            .enumerate()
            .find(|(i, family)| queue_counts[*i] > 0 && family.flags.contains(QueueFlags::TRANSFER))
            .map(|(i, family)| {
                queue_counts[i] -= 1;
                *family
            });

        Ok(Self {
            graphics,
            compute,
            transfer,
        })
    }
}