use std::{
    borrow::Cow,
    collections::HashMap,
    sync::{Arc, Weak},
};

use crate::{
    define_device_owned_handle,
    device::{DeviceOwned, QueueFlags},
};

use super::Device;
use ash::{prelude::*, vk};
use itertools::Itertools;
use parking_lot::Mutex;
use vk_mem::Alloc;

#[derive(Clone)]
pub struct ImageDesc {
    pub flags: vk::ImageCreateFlags,
    pub name: Option<Cow<'static, str>>,
    pub format: vk::Format,
    pub kind: vk::ImageType,
    pub mip_levels: u32,
    pub array_layers: u32,
    pub samples: vk::SampleCountFlags,
    pub extent: vk::Extent3D,
    pub tiling: vk::ImageTiling,
    pub usage: vk::ImageUsageFlags,
    pub queue_families: QueueFlags,
    pub layout: vk::ImageLayout,

    pub mem_usage: vk_mem::MemoryUsage,
    pub alloc_flags: vk_mem::AllocationCreateFlags,
}

impl std::hash::Hash for ImageDesc {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.flags.hash(state);
        self.format.hash(state);
        self.kind.hash(state);
        self.mip_levels.hash(state);
        self.array_layers.hash(state);
        self.samples.hash(state);
        self.extent.hash(state);
        self.tiling.hash(state);
        self.usage.hash(state);
        self.queue_families.hash(state);
        self.layout.hash(state);
        self.mem_usage.hash(state);
        self.alloc_flags.bits().hash(state);
    }
}

impl Eq for ImageDesc {}
impl PartialEq for ImageDesc {
    fn eq(&self, other: &Self) -> bool {
        self.flags == other.flags
            && self.name == other.name
            && self.format == other.format
            && self.kind == other.kind
            && self.mip_levels == other.mip_levels
            && self.array_layers == other.array_layers
            && self.samples == other.samples
            && self.extent == other.extent
            && self.tiling == other.tiling
            && self.usage == other.usage
            && self.queue_families == other.queue_families
            && self.layout == other.layout
            && self.mem_usage == other.mem_usage
            && self.alloc_flags.bits() == other.alloc_flags.bits()
    }
}

impl<'a> std::fmt::Debug for ImageDesc {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ImageDesc")
            .field("flags", &self.flags)
            .field("name", &self.name)
            .field("format", &self.format)
            .field("kind", &self.kind)
            .field("mip_levels", &self.mip_levels)
            .field("array_layers", &self.array_layers)
            .field("samples", &self.samples)
            .field("extent", &self.extent)
            .field("tiling", &self.tiling)
            .field("usage", &self.usage)
            .field("queue_families", &self.queue_families)
            .field("layout", &self.layout)
            .field("mem_usage", &self.mem_usage)
            .field_with("alloc_flags", |f| {
                write!(
                    f,
                    "{}",
                    self.alloc_flags
                        .iter_names()
                        .map(|(name, _)| name)
                        .format(" | ")
                )
            })
            .finish()
    }
}

impl Default for ImageDesc {
    fn default() -> Self {
        Self {
            flags: Default::default(),
            name: Default::default(),
            format: Default::default(),
            kind: vk::ImageType::TYPE_2D,
            samples: vk::SampleCountFlags::TYPE_1,
            mip_levels: 1,
            array_layers: 1,
            extent: Default::default(),
            tiling: vk::ImageTiling::OPTIMAL,
            usage: Default::default(),
            queue_families: QueueFlags::empty(),
            layout: vk::ImageLayout::UNDEFINED,
            alloc_flags: vk_mem::AllocationCreateFlags::empty(),
            mem_usage: vk_mem::MemoryUsage::Auto,
        }
    }
}

define_device_owned_handle! {
    #[derive(Debug)]
    pub Image(vk::Image) {
        alloc: Option<vk_mem::Allocation>,
        size: vk::Extent3D,
        format: vk::Format,
        views: Mutex<HashMap<ImageViewDesc, vk::ImageView>>,
        aliases: Mutex<HashMap<ImageDesc, Arc<Image>>>,
        parent: Option<Weak<Image>>,
        is_swapchain_image: bool,
    } => |this| if !this.is_swapchain_image {
        unsafe {
            for &view in this.views.lock().values() {
                this.inner.dev().dev().destroy_image_view(view, None);
            }

            let handle = this.handle();
            let dev = this.device().clone();
            if let Some(alloc) = this.alloc.as_mut() {
                // destroy image handle and allocation
                dev.alloc().destroy_image(handle, alloc);
            } else {
                // destroy image handle
                dev.dev().destroy_image(handle, None);
            }
        }
    }
}

impl Eq for Image {}
impl PartialEq for Image {
    fn eq(&self, other: &Self) -> bool {
        self.inner == other.inner
    }
}

impl Image {
    pub fn new(device: Device, desc: ImageDesc) -> VkResult<Self> {
        tracing::trace!("allocate new image with desc={desc:?}");
        let ImageDesc {
            flags,
            name,
            format,
            kind,
            mip_levels,
            array_layers,
            samples,
            extent,
            tiling,
            usage,
            queue_families,
            layout,
            mem_usage,
            alloc_flags,
        } = desc;

        let queue_families = device.queue_families().family_indices(queue_families);

        let sharing_mode = if queue_families.len() > 1 {
            vk::SharingMode::CONCURRENT
        } else {
            vk::SharingMode::EXCLUSIVE
        };

        let info = &vk::ImageCreateInfo::default()
            .flags(flags)
            .image_type(kind)
            .format(format)
            .extent(extent)
            .samples(samples)
            .initial_layout(layout)
            .tiling(tiling)
            .usage(usage)
            .sharing_mode(sharing_mode)
            .queue_family_indices(&queue_families)
            .array_layers(array_layers)
            .mip_levels(mip_levels);

        let alloc_info = &vk_mem::AllocationCreateInfo {
            usage: mem_usage,
            flags: alloc_flags,
            ..Default::default()
        };

        let (handle, alloc) = unsafe { device.alloc().create_image(info, alloc_info)? };

        Self::construct(
            device,
            handle,
            name,
            Some(alloc),
            extent,
            format,
            Mutex::new(HashMap::new()),
            Mutex::new(HashMap::new()),
            None, // aliased
            false,
        )
    }

    pub unsafe fn from_swapchain_image(
        device: Device,
        image: vk::Image,
        name: Option<Cow<'static, str>>,
        extent: vk::Extent3D,
        format: vk::Format,
    ) -> Result<Image, vk::Result> {
        Self::construct(
            device,
            image,
            name,
            None,
            extent,
            format,
            Mutex::new(HashMap::new()),
            Mutex::new(HashMap::new()),
            None,
            true,
        )
    }

    pub fn format(&self) -> vk::Format {
        self.format
    }
    pub fn image(&self) -> vk::Image {
        self.handle()
    }
    pub fn size(&self) -> vk::Extent3D {
        self.size
    }
    pub fn extent_2d(&self) -> vk::Extent2D {
        vk::Extent2D {
            width: self.size.width,
            height: self.size.height,
        }
    }
    pub fn width(&self) -> u32 {
        self.size.width
    }
    pub fn height(&self) -> u32 {
        self.size.height
    }
    pub fn depth(&self) -> u32 {
        self.size.depth
    }

    fn get_parent_or_self(self: &Arc<Self>) -> Arc<Image> {
        self.parent
            .as_ref()
            .map(|weak| weak.upgrade().unwrap())
            .unwrap_or_else(|| self.clone())
    }

    pub unsafe fn get_alias(self: &Arc<Self>, desc: ImageDesc) -> VkResult<Arc<Self>> {
        self.get_parent_or_self().get_alias_inner(desc)
    }

    /// must only be called on the primogenitor of an image.
    /// get the primogenitor with [`Self::get_parent_or_self()`]
    unsafe fn get_alias_inner(self: Arc<Self>, desc: ImageDesc) -> VkResult<Arc<Image>> {
        use std::collections::hash_map::Entry::*;
        match self.aliases.lock().entry(desc.clone()) {
            Occupied(occupied) => Ok(occupied.get().clone()),
            Vacant(vacant) => {
                let ImageDesc {
                    flags,
                    name,
                    format,
                    kind,
                    mip_levels,
                    array_layers,
                    samples,
                    extent,
                    tiling,
                    usage,
                    queue_families,
                    layout,
                    ..
                } = desc;

                let queue_families = self
                    .device()
                    .queue_families()
                    .family_indices(queue_families);

                let sharing_mode = if queue_families.len() > 1 {
                    vk::SharingMode::CONCURRENT
                } else {
                    vk::SharingMode::EXCLUSIVE
                };

                let info = &vk::ImageCreateInfo::default()
                    .flags(flags)
                    .image_type(kind)
                    .format(format)
                    .extent(extent)
                    .samples(samples)
                    .initial_layout(layout)
                    .tiling(tiling)
                    .usage(usage)
                    .sharing_mode(sharing_mode)
                    .queue_family_indices(&queue_families)
                    .array_layers(array_layers)
                    .mip_levels(mip_levels);

                let alloc = self
                    .alloc
                    .as_ref()
                    .expect("no alloc associated with image. is this the framebuffer?");

                let image = unsafe {
                    let image = self.device().dev().create_image(info, None)?;

                    let req = self.device().dev().get_image_memory_requirements(image);
                    if self.device().alloc().get_allocation_info(alloc).size < req.size {
                        return Err(vk::Result::ERROR_MEMORY_MAP_FAILED);
                    }

                    self.device().alloc().bind_image_memory(alloc, image)?;
                    image
                };

                let alias = Self::construct(
                    self.device().clone(),
                    image,
                    name,
                    None,
                    extent,
                    format,
                    Mutex::new(HashMap::new()),
                    Mutex::new(HashMap::new()),
                    Some(Arc::downgrade(&self)),
                    self.is_swapchain_image,
                )?;
                Ok(vacant.insert(Arc::new(alias)).clone())
            }
        }
    }

    /// technically, this ImageView belongs to the image and is managed by it.
    pub fn get_view(&self, desc: ImageViewDesc) -> VkResult<vk::ImageView> {
        use std::collections::hash_map::Entry::*;
        match self.views.lock().entry(desc.hash_eq_copy()) {
            Occupied(occupied) => Ok(*occupied.get()),
            Vacant(vacant) => {
                let view = unsafe {
                    let create_info = vk::ImageViewCreateInfo::default()
                        .flags(desc.flags)
                        .image(self.image())
                        .view_type(vk::ImageViewType::TYPE_2D)
                        .format(desc.format)
                        .components(desc.components)
                        .subresource_range(
                            vk::ImageSubresourceRange::default()
                                .aspect_mask(desc.aspect)
                                .base_mip_level(desc.mip_range.0)
                                .level_count(desc.mip_range.count())
                                .base_array_layer(desc.layer_range.0)
                                .layer_count(desc.layer_range.count()),
                        );
                    self.device().dev().create_image_view(&create_info, None)?
                };
                Ok(*vacant.insert(view))
            }
        }
    }

    pub fn create_view(&self, desc: ImageViewDesc) -> VkResult<ImageView> {
        let create_info = vk::ImageViewCreateInfo::default()
            .flags(desc.flags)
            .image(self.image())
            .view_type(vk::ImageViewType::TYPE_2D)
            .format(desc.format)
            .components(desc.components)
            .subresource_range(
                vk::ImageSubresourceRange::default()
                    .aspect_mask(desc.aspect)
                    .base_mip_level(desc.mip_range.0)
                    .level_count(desc.mip_range.count())
                    .base_array_layer(desc.layer_range.0)
                    .layer_count(desc.layer_range.count()),
            );

        let view = unsafe { self.device().dev().create_image_view(&create_info, None)? };

        ImageView::construct(self.device().clone(), view, desc.name)
    }
}

#[derive(Debug, Default, Clone)]
pub struct ImageViewDesc {
    pub flags: vk::ImageViewCreateFlags,
    pub name: Option<Cow<'static, str>>,
    pub kind: vk::ImageViewType,
    pub format: vk::Format,
    pub components: vk::ComponentMapping,
    pub aspect: vk::ImageAspectFlags,
    pub mip_range: MipRange,
    pub layer_range: MipRange,
}

impl ImageViewDesc {
    pub fn hash_eq_copy(&self) -> Self {
        let &Self {
            flags,
            kind,
            format,
            components,
            aspect,
            mip_range,
            layer_range,
            ..
        } = self;

        Self {
            flags,
            name: None,
            kind,
            format,
            components,
            aspect,
            mip_range,
            layer_range,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct MipRange(u32, u32);

impl Default for MipRange {
    fn default() -> Self {
        Self(0, vk::REMAINING_ARRAY_LAYERS)
    }
}

impl MipRange {
    fn count(&self) -> u32 {
        self.1
    }
}

impl<R: core::ops::RangeBounds<u32>> From<R> for MipRange {
    fn from(value: R) -> Self {
        let start = match value.start_bound() {
            std::ops::Bound::Included(v) => *v,
            std::ops::Bound::Excluded(v) => *v + 1,
            std::ops::Bound::Unbounded => 0,
        };
        let count = match value.end_bound() {
            std::ops::Bound::Included(v) => *v + 1 - start,
            std::ops::Bound::Excluded(v) => *v - start,
            std::ops::Bound::Unbounded => vk::REMAINING_MIP_LEVELS,
        };

        Self(start, count)
    }
}

impl std::hash::Hash for ImageViewDesc {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.flags.hash(state);
        self.kind.hash(state);
        self.format.hash(state);
        (
            self.components.r,
            self.components.g,
            self.components.b,
            self.components.a,
        )
            .hash(state);
        self.aspect.hash(state);
        self.layer_range.hash(state);
        self.mip_range.hash(state);
    }
}

impl Eq for ImageViewDesc {}
impl PartialEq for ImageViewDesc {
    fn eq(&self, other: &Self) -> bool {
        self.flags == other.flags
            && self.kind == other.kind
            && self.format == other.format
            && (
                self.components.r,
                self.components.g,
                self.components.b,
                self.components.a,
            ) == (
                other.components.r,
                other.components.g,
                other.components.b,
                other.components.a,
            )
            && self.aspect == other.aspect
            && self.mip_range == other.mip_range
            && self.layer_range == other.layer_range
    }
}

define_device_owned_handle! {
    #[derive(Debug)]
    pub ImageView(vk::ImageView) {} => |this|  unsafe {
        this.device().dev().destroy_image_view(this.handle(), None);
    }
}

pub struct QueueOwnership {
    pub src: u32,
    pub dst: u32,
}

pub const SUBRESOURCERANGE_ALL: vk::ImageSubresourceRange = vk::ImageSubresourceRange {
    aspect_mask: vk::ImageAspectFlags::empty(),
    base_mip_level: 0,
    level_count: vk::REMAINING_MIP_LEVELS,
    base_array_layer: 0,
    layer_count: vk::REMAINING_ARRAY_LAYERS,
};

pub const SUBRESOURCERANGE_COLOR_ALL: vk::ImageSubresourceRange = vk::ImageSubresourceRange {
    aspect_mask: vk::ImageAspectFlags::COLOR,
    base_mip_level: 0,
    level_count: vk::REMAINING_MIP_LEVELS,
    base_array_layer: 0,
    layer_count: vk::REMAINING_ARRAY_LAYERS,
};