use std::{borrow::Cow, sync::Arc};

use crate::{
    device::{Allocation, AllocationStrategy, DeviceHandle, DeviceObject, QueueFlags},
    swapchain::Swapchain,
    util::weak_vec::WeakVec,
};

use super::Device;
use ash::vk;
use gpu_allocator::vulkan::{AllocationCreateDesc, AllocationScheme};
use parking_lot::Mutex;

#[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_location: gpu_allocator::MemoryLocation,
    pub alloc_scheme: AllocationStrategy,
}

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_location.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_location == other.mem_location
    }
}

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_location", &self.mem_location)
            .field("alloc_scheme", &self.alloc_scheme)
            .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,
            mem_location: gpu_allocator::MemoryLocation::Unknown,
            alloc_scheme: AllocationStrategy::AllocatorManaged,
        }
    }
}

#[derive(Debug)]
enum ImageInner {
    Swapchain(vk::Image, Device),
    Allocated(DeviceObject<vk::Image>, Allocation),
}

impl DeviceHandle for vk::Image {
    unsafe fn destroy(&mut self, device: &Device) {
        unsafe {
            device.dev().destroy_image(*self, None);
        }
    }
}

impl ImageInner {
    fn image(&self) -> vk::Image {
        match self {
            Self::Swapchain(image, _) => *image,
            Self::Allocated(image, _) => **image,
        }
    }
    fn device(&self) -> &Device {
        match self {
            Self::Swapchain(_, device) => device,
            Self::Allocated(image, _) => image.device(),
        }
    }

    fn allocation(&self) -> Option<&Allocation> {
        match self {
            Self::Swapchain(_, _) => None,
            Self::Allocated(_, alloc) => Some(alloc),
        }
    }
    fn allocation_mut(&mut self) -> Option<&mut Allocation> {
        match self {
            Self::Swapchain(_, _) => None,
            Self::Allocated(_, alloc) => Some(alloc),
        }
    }
}

#[derive(Debug)]
pub struct Image {
    image: ImageInner,
    desc: ImageDesc,
    views: Mutex<WeakVec<ImageView>>,
}

impl Image {
    pub fn new(device: Device, desc: ImageDesc) -> crate::Result<Self> {
        let (image, requirements) = Self::new_raw(device.clone(), &desc)?;
        let alloc = device.alloc2.lock().allocate(&AllocationCreateDesc {
            name: desc.name.as_deref().unwrap_or(""),
            requirements,
            location: desc.mem_location,
            linear: desc.tiling == vk::ImageTiling::LINEAR,
            allocation_scheme: match desc.alloc_scheme {
                AllocationStrategy::AllocatorManaged => AllocationScheme::GpuAllocatorManaged,
                AllocationStrategy::Dedicated => AllocationScheme::DedicatedImage(image),
            },
        })?;

        Ok(Self {
            image: ImageInner::Allocated(
                DeviceObject::new(image, device.clone(), desc.name.clone()),
                Allocation::Owned(DeviceObject::new_without_name(alloc, device)),
            ),
            desc,
            views: Default::default(),
        })
    }

    pub fn new_with_allocation(
        device: Device,
        allocation: Allocation,
        desc: ImageDesc,
    ) -> crate::Result<Self> {
        let (image, requirements) = Self::new_raw(device.clone(), &desc)?;

        // validate allocation
        let alloc_size = allocation
            .allocation()
            .map(|alloc| alloc.size())
            .unwrap_or(0);
        if alloc_size < requirements.size {
            tracing::error!(
                "allocation size {} is smaller than image memory requirements {}",
                alloc_size,
                requirements.size
            );
            return Err(crate::Error::Unspecified);
        }

        if allocation
            .allocation()
            .map(|alloc| 1 << alloc.memory_type_index())
            .unwrap_or(0)
            & requirements.memory_type_bits
            == 0
        {
            return Err(crate::Error::Unspecified);
        }

        Ok(Self {
            image: ImageInner::Allocated(
                DeviceObject::new(image, device.clone(), desc.name.clone()),
                allocation,
            ),
            desc,
            views: Default::default(),
        })
    }

    pub fn from_swapchain_image(image: vk::Image, swapchain: &Swapchain) -> Self {
        Self {
            image: ImageInner::Swapchain(image, swapchain.swapchain.device().clone()),
            desc: ImageDesc {
                format: swapchain.config.format,
                kind: vk::ImageType::TYPE_2D,
                mip_levels: 1,
                array_layers: 1,
                samples: vk::SampleCountFlags::TYPE_1,
                extent: vk::Extent3D {
                    width: swapchain.config.extent.width,
                    height: swapchain.config.extent.height,
                    depth: 1,
                },
                tiling: vk::ImageTiling::OPTIMAL,
                usage: swapchain.config.usage,
                queue_families: QueueFlags::PRESENT,
                layout: vk::ImageLayout::UNDEFINED,
                mem_location: gpu_allocator::MemoryLocation::GpuOnly,
                ..Default::default()
            },
            views: Default::default(),
        }
    }

    fn new_raw(
        device: Device,
        desc: &ImageDesc,
    ) -> crate::Result<(vk::Image, vk::MemoryRequirements)> {
        tracing::trace!("allocate new image with desc={desc:?}");
        let ImageDesc {
            flags,
            format,
            kind,
            mip_levels,
            array_layers,
            samples,
            extent,
            tiling,
            usage,
            queue_families,
            layout,
            ..
        } = desc;

        let queue_families = device.queues.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);

        // validate
        let limits = &device.adapter.properties.core.limits;
        let max_dim = match *kind {
            vk::ImageType::TYPE_1D => limits.max_image_dimension1_d,
            vk::ImageType::TYPE_2D => limits.max_image_dimension2_d,
            vk::ImageType::TYPE_3D => limits.max_image_dimension3_d,
            _ => unreachable!(),
        };

        if extent.width > max_dim || extent.height > max_dim || extent.depth > max_dim {
            tracing::error!(
                "image extent {extent:?} exceeds device limits (max dimension: {max_dim})"
            );

            return Err(crate::Error::ImageTooLarge {
                width: extent.width,
                height: extent.height,
                max_size: max_dim,
            });
        }

        let image = unsafe { device.raw.create_image(&info, None)? };

        let requirements = unsafe { device.raw.get_image_memory_requirements(image) };

        Ok((image, requirements))
    }

    pub fn raw(&self) -> vk::Image {
        self.image.image()
    }
}

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

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

    pub fn allocation(&self) -> Option<&Allocation> {
        self.image.allocation()
    }

    pub fn allocation_mut(&mut self) -> Option<&mut Allocation> {
        self.image.allocation_mut()
    }

    // /// 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: &Arc<Self>, mut desc: ImageViewDesc) -> crate::Result<ImageView> {
        // validate
        if !view_kind_compatible(self.desc.kind, desc.kind) {
            tracing::error!(
                "image view kind {:?} is not compatible with image kind {:?}",
                desc.kind,
                self.desc.kind
            );
            return Err(crate::Error::IncompatibleImageViewKind {
                image_kind: self.desc.kind,
                view_kind: desc.kind,
            });
        }

        if desc.mip_range.0 > self.desc.mip_levels || desc.mip_range.1 > self.desc.mip_levels {
            tracing::error!(
                "image view mip range {:?} exceeds image mip levels {}",
                desc.mip_range,
                self.desc.mip_levels
            );
            return Err(crate::Error::Unspecified);
        }

        if desc.format == vk::Format::UNDEFINED {
            desc.format = self.desc.format;
        }

        let create_info = vk::ImageViewCreateInfo::default()
            .flags(desc.flags)
            .image(self.image())
            .view_type(desc.kind)
            .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 device = self.image.device();
        let view = unsafe { device.raw.create_image_view(&create_info, None)? };

        Ok(ImageView {
            view: DeviceObject::new(view, device.clone(), desc.name.clone()),
            desc,
            image: self.clone(),
        })
    }
}

fn validate_image_view_format(image: &ImageDesc, view_format: vk::Format) -> bool {}

fn view_kind_compatible(image_kind: vk::ImageType, view_kind: vk::ImageViewType) -> bool {
    use vk::ImageType as IT;
    use vk::ImageViewType as VT;
    match (image_kind, view_kind) {
        (IT::TYPE_1D, VT::TYPE_1D | VT::TYPE_1D_ARRAY) => true,
        (IT::TYPE_2D, VT::TYPE_2D | VT::TYPE_2D_ARRAY | VT::CUBE | VT::CUBE_ARRAY) => true,
        (IT::TYPE_3D, VT::TYPE_2D | VT::TYPE_2D_ARRAY | VT::TYPE_3D) => true,
        _ => false,
    }
}

#[derive(Debug, Default, Clone)]
pub struct ImageViewDesc {
    pub name: Option<Cow<'static, str>>,
    pub flags: vk::ImageViewCreateFlags,
    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 color_2d() -> Self {
        Self {
            kind: vk::ImageViewType::TYPE_2D,
            aspect: vk::ImageAspectFlags::COLOR,
            ..Default::default()
        }
    }

    pub fn with_format(self, format: vk::Format) -> Self {
        Self { format, ..self }
    }

    pub(crate) 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
    }
}

#[derive(Debug)]
pub struct ImageView {
    view: DeviceObject<vk::ImageView>,
    desc: ImageViewDesc,
    image: Arc<Image>,
}

impl ImageView {
    pub fn raw(&self) -> vk::ImageView {
        *self.view
    }
    pub fn image(&self) -> &Arc<Image> {
        &self.image
    }
}

impl DeviceHandle for vk::ImageView {
    unsafe fn destroy(&mut self, device: &Device) {
        unsafe {
            device.dev().destroy_image_view(*self, 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,
};