use std::borrow::Cow;

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

use super::Device;
use ash::{prelude::*, vk};
use itertools::Itertools;
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<'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: vk_mem::Allocation,
        size: vk::Extent3D,
        format: vk::Format,
    } => |this|  unsafe {
        this.inner.dev().alloc().destroy_image(this.handle(), &mut this.alloc);
    }
}

impl Image {
    pub fn new(device: Device, desc: ImageDesc) -> VkResult<Self> {
        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, alloc, extent, format)
    }

    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
    }

    pub fn 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,
}

#[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_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,
};