use std::{borrow::Cow, collections::BTreeMap, ops::Deref, sync::Arc};

use ash::{
    khr,
    prelude::VkResult,
    vk::{self, Handle},
};
use parking_lot::Mutex;
use tinyvec::{array_vec, ArrayVec};

use crate::{sync, Instance, PhysicalDevice, Queue};

#[derive(Debug, Default)]
pub struct DeviceQueueFamilies {
    pub(crate) families: Vec<(u32, u32)>,
    pub(crate) graphics: (u32, u32),
    pub(crate) present: (u32, u32),
    pub(crate) async_compute: (u32, u32),
    pub(crate) transfer: (u32, u32),
}

impl DeviceQueueFamilies {
    pub fn swapchain_family_indices(&self) -> ArrayVec<[u32; 2]> {
        let mut indices = array_vec!([u32; 2] => self.graphics.0);

        if self.present.0 != self.graphics.0 {
            indices.push(self.present.0);
        }

        indices
    }
    pub fn graphics_familty(&self) -> u32 {
        self.graphics.0
    }
    pub fn present_familty(&self) -> u32 {
        self.present.0
    }
    pub fn async_compute_familty(&self) -> u32 {
        self.async_compute.0
    }
    pub fn transfer_familty(&self) -> u32 {
        self.transfer.0
    }
}

#[repr(transparent)]
struct DeviceWrapper(ash::Device);

impl Deref for DeviceWrapper {
    type Target = ash::Device;

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

impl Drop for DeviceWrapper {
    fn drop(&mut self) {
        unsafe {
            _ = self.0.device_wait_idle();
            self.0.destroy_device(None);
        }
    }
}

pub struct DeviceInner {
    alloc: vk_mem::Allocator,
    device: DeviceWrapper,
    physical: PhysicalDevice,
    instance: Arc<Instance>,
    swapchain: khr::swapchain::Device,
    debug_utils: ash::ext::debug_utils::Device,
    allocated_queues: BTreeMap<(u32, u32), Queue>,
    // these are resident in allocated_queues, and may in fact be clones of each
    // other, for ease of access
    main_queue: Queue,
    compute_queue: Queue,
    transfer_queue: Queue,
    present_queue: Queue,
    sync_threadpool: sync::SyncThreadpool,
}

impl core::fmt::Debug for DeviceInner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DeviceInner")
            .field("device", &self.device.handle())
            .finish()
    }
}

#[derive(Clone, Debug)]
pub struct Device(Arc<DeviceInner>);
pub type WeakDevice = std::sync::Weak<DeviceInner>;

impl Device {
    pub fn new(
        instance: Arc<Instance>,
        physical: PhysicalDevice,
        mut features: crate::PhysicalDeviceFeatures,
    ) -> VkResult<Self> {
        // we have 4 queues at most: graphics, compute, transfer, present
        let priorities = [1.0f32; 4];

        let queue_infos = physical
            .queue_families
            .families
            .iter()
            .map(|&(family, queues)| {
                vk::DeviceQueueCreateInfo::default()
                    .queue_family_index(family)
                    .queue_priorities(&priorities[..queues as usize])
            })
            .collect::<Vec<_>>();

        let extensions = features
            .device_extensions
            .iter()
            .map(|ext| ext.extension_name.as_ptr())
            .collect::<Vec<_>>();

        let mut features2 = features.features2();
        let device_info = vk::DeviceCreateInfo::default()
            .queue_create_infos(&queue_infos)
            .enabled_extension_names(&extensions)
            .push_next(&mut features2);

        let device = unsafe {
            let device = instance
                .instance
                .create_device(physical.pdev, &device_info, None)?;

            let allocated_queues = queue_infos
                .iter()
                .flat_map(|info| {
                    (0..info.queue_count).map(|i| {
                        (
                            (info.queue_family_index, i),
                            Queue::new(&device, info.queue_family_index, i),
                        )
                    })
                })
                .collect::<BTreeMap<_, _>>();

            let get_queue =
                |(family, index)| allocated_queues.get(&(family, index)).cloned().unwrap();

            let main_queue = get_queue(physical.queue_families.graphics);
            let present_queue = get_queue(physical.queue_families.present);
            let compute_queue = get_queue(physical.queue_families.async_compute);
            let transfer_queue = get_queue(physical.queue_families.transfer);

            let alloc_info =
                vk_mem::AllocatorCreateInfo::new(&instance.instance, &device, physical.pdev);

            let alloc = unsafe { vk_mem::Allocator::new(alloc_info)? };

            DeviceInner {
                device: DeviceWrapper(device.clone()),
                physical,
                swapchain: khr::swapchain::Device::new(&instance.instance, &device),
                debug_utils: ash::ext::debug_utils::Device::new(&instance.instance, &device),
                instance,
                alloc,
                allocated_queues,
                main_queue,
                present_queue,
                compute_queue,
                transfer_queue,
                sync_threadpool: sync::SyncThreadpool::new(),
            }
        };

        Ok(Self(Arc::new(device)))
    }
    pub fn sync_threadpool(&self) -> &sync::SyncThreadpool {
        &self.0.sync_threadpool
    }
    pub fn weak(&self) -> WeakDevice {
        Arc::downgrade(&self.0)
    }
    pub fn alloc(&self) -> &vk_mem::Allocator {
        &self.0.alloc
    }
    pub fn dev(&self) -> &ash::Device {
        &self.0.device
    }
    pub fn swapchain(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
    pub fn debug_utils(&self) -> &ash::ext::debug_utils::Device {
        &self.0.debug_utils
    }
    pub fn queue_families(&self) -> &DeviceQueueFamilies {
        &self.0.physical.queue_families
    }
    pub fn phy(&self) -> vk::PhysicalDevice {
        self.0.physical.pdev
    }
    pub fn graphics_queue(&self) -> &Queue {
        &self.0.main_queue
    }
    pub fn present_queue(&self) -> &Queue {
        &self.0.present_queue
    }

    pub unsafe fn lock_queues(&self) {
        // this is obviously awful, allocating for this
        self.0
            .allocated_queues
            .values()
            .for_each(|q| core::mem::forget(q.lock()));
    }

    pub unsafe fn unlock_queues(&self) {
        self.0
            .allocated_queues
            .values()
            .for_each(|q| unsafe { q.0.force_unlock() });
    }

    pub fn wait_queue_idle(&self, queue: &Queue) -> VkResult<()> {
        tracing::warn!("locking queue {queue:?} and waiting for idle");

        queue.with_locked(|q| unsafe { self.dev().queue_wait_idle(q) })?;

        tracing::warn!("finished waiting: unlocking queue {queue:?}.");
        Ok(())
    }

    pub fn wait_idle(&self) -> VkResult<()> {
        tracing::warn!("locking all queues and waiting for device to idle");
        unsafe {
            self.lock_queues();
            self.dev().device_wait_idle()?;
            self.unlock_queues();
        }
        tracing::warn!("finished waiting: unlocking all queues.");
        Ok(())
    }
}

impl AsRef<khr::swapchain::Device> for Device {
    fn as_ref(&self) -> &khr::swapchain::Device {
        &self.0.swapchain
    }
}

impl AsRef<ash::Device> for Device {
    fn as_ref(&self) -> &ash::Device {
        &self.0.device
    }
}

pub struct DeviceAndQueues {
    pub(crate) device: Device,
    pub(crate) main_queue: Queue,
    pub(crate) compute_queue: Queue,
    pub(crate) transfer_queue: Queue,
    pub(crate) present_queue: Queue,
}

impl AsRef<ash::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::Device {
        &self.device.as_ref()
    }
}
impl AsRef<ash::khr::swapchain::Device> for DeviceAndQueues {
    fn as_ref(&self) -> &ash::khr::swapchain::Device {
        &self.device.as_ref()
    }
}

#[derive(Clone)]
pub struct DeviceOwnedDebugObject<T> {
    device: Device,
    object: T,
    name: Option<Cow<'static, str>>,
}

impl<T: std::fmt::Debug + vk::Handle + Copy> std::fmt::Debug for DeviceOwnedDebugObject<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct(core::any::type_name::<T>())
            .field_with("device", |f| {
                write!(f, "0x{:x}", self.device.0.device.handle().as_raw())
            })
            .field_with("handle", |f| write!(f, "0x{:x}", &self.object.as_raw()))
            .field("name", &self.name)
            .finish()
    }
}

impl<T> DeviceOwnedDebugObject<T> {
    pub fn new<S: Into<Cow<'static, str>>>(
        device: crate::Device,
        object: T,
        name: Option<S>,
    ) -> ash::prelude::VkResult<Self>
    where
        T: vk::Handle + Copy,
    {
        let name = name.map(Into::<Cow<_>>::into);
        if let Some(name) = name.as_ref() {
            let name =
                std::ffi::CString::new(name.as_bytes()).unwrap_or(c"invalid name".to_owned());
            unsafe {
                device.debug_utils().set_debug_utils_object_name(
                    &vk::DebugUtilsObjectNameInfoEXT::default()
                        .object_handle(object)
                        .object_name(&name),
                )?;
            }
        }

        Ok(Self {
            device,
            object,
            name,
        })
    }

    pub fn dev(&self) -> &crate::Device {
        &self.device
    }
    pub fn handle(&self) -> T
    where
        T: Copy,
    {
        self.object
    }
}