mev 0.1.0

use core::fmt;
use std::{
    hash::{Hash, Hasher},
    mem::{size_of, ManuallyDrop},
    sync::Arc,
};

use ash::vk::{self, Handle};
use gpu_alloc::MemoryBlock;
use hashbrown::{hash_map::Entry, HashMap};
use parking_lot::Mutex;

use crate::generic::{
    ArgumentKind, Automatic, ImageExtent, ImageUsage, OutOfMemory, PixelFormat, Sampled, Storage,
    Swizzle, ViewDesc,
};

use super::{
    arguments::ArgumentsField,
    device::{Device, DeviceMemory, DeviceOwned, WeakDevice},
    from::IntoAsh,
    refs::Refs,
};

enum Flavor {
    Device {
        block: ManuallyDrop<MemoryBlock<DeviceMemory>>,
        idx: usize,
    },
    Swapchain,
}

// Contains actual `vk::Image`
struct ImageData {
    owner: WeakDevice,
    format: PixelFormat,
    usage: ImageUsage,
    extent: ImageExtent,
    layers: u32,
    levels: u32,
    flavor: Flavor,
    views: Mutex<HashMap<ViewDesc, (vk::ImageView, usize)>>,
}

impl Drop for ImageData {
    fn drop(&mut self) {
        self.owner
            .drop_image_views(self.views.get_mut().values().map(|(_, idx)| *idx));

        if let Flavor::Device { block, idx } = &mut self.flavor {
            self.owner
                .drop_image(*idx, unsafe { ManuallyDrop::take(block) });
        }
    }
}

struct Inner {
    data: Arc<ImageData>,
    desc: ViewDesc,
    usage: ImageUsage,
    extent: ImageExtent,
    owner: WeakDevice,
}

#[derive(Clone)]
pub struct Image {
    handle: vk::Image,
    view: vk::ImageView,
    inner: Arc<Inner>,
}

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

impl Eq for Image {}

impl Hash for Image {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.handle.hash(state);
        self.view.hash(state);
        Arc::as_ptr(&self.inner).hash(state);
    }
}

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

impl DeviceOwned for Image {
    #[inline]
    fn owner(&self) -> &WeakDevice {
        &self.inner.owner
    }
}

impl Image {
    fn build(
        owner: WeakDevice,
        handle: vk::Image,
        view: vk::ImageView,
        view_idx: usize,
        extent: impl Into<ImageExtent>,
        format: PixelFormat,
        usage: ImageUsage,
        layers: u32,
        levels: u32,
        flavor: Flavor,
    ) -> Self {
        let extent = extent.into();
        let desc = ViewDesc {
            format,
            base_layer: 0,
            layers,
            base_level: 0,
            levels,
            swizzle: Swizzle::IDENTITY,
        };

        let mut views = HashMap::new();
        views.insert(desc, (view, view_idx));

        Image {
            handle,
            view,
            inner: Arc::new(Inner {
                data: Arc::new(ImageData {
                    owner: owner.clone(),
                    extent,
                    format,
                    usage,
                    layers,
                    levels,
                    flavor,
                    views: Mutex::new(views),
                }),
                desc,
                extent,
                usage,
                owner,
            }),
        }
    }

    pub(super) fn new(
        owner: WeakDevice,
        handle: vk::Image,
        view: vk::ImageView,
        view_idx: usize,
        extent: ImageExtent,
        format: PixelFormat,
        usage: ImageUsage,
        layers: u32,
        levels: u32,
        block: MemoryBlock<DeviceMemory>,
        idx: usize,
    ) -> Self {
        Image::build(
            owner,
            handle,
            view,
            view_idx,
            extent,
            format,
            usage,
            layers,
            levels,
            Flavor::Device {
                block: ManuallyDrop::new(block),
                idx,
            },
        )
    }

    pub(super) fn from_swapchain_image(
        owner: WeakDevice,
        handle: vk::Image,
        view: vk::ImageView,
        view_idx: usize,
        extent: impl Into<ImageExtent>,
        format: PixelFormat,
        usage: ImageUsage,
    ) -> Self {
        Image::build(
            owner,
            handle,
            view,
            view_idx,
            extent,
            format,
            usage,
            1,
            1,
            Flavor::Swapchain,
        )
    }

    /// Creates a null/invalid Image for use when device OOM occurs.
    pub(super) fn null(
        extent: ImageExtent,
        format: PixelFormat,
        usage: ImageUsage,
        layers: u32,
        levels: u32,
    ) -> Self {
        let desc = ViewDesc {
            format,
            base_layer: 0,
            layers,
            base_level: 0,
            levels,
            swizzle: Swizzle::IDENTITY,
        };
        Image {
            handle: vk::Image::null(),
            view: vk::ImageView::null(),
            inner: Arc::new(Inner {
                data: Arc::new(ImageData {
                    owner: WeakDevice::null(),
                    format,
                    usage,
                    extent,
                    layers,
                    levels,
                    flavor: Flavor::Swapchain,
                    views: Mutex::new(hashbrown::HashMap::new()),
                }),
                desc,
                extent,
                usage,
                owner: WeakDevice::null(),
            }),
        }
    }

    #[inline]
    pub(super) fn get_view(&self, device: &Device, desc: ViewDesc) -> Image {
        if self.handle.is_null() || self.view.is_null() {
            return Image {
                handle: vk::Image::null(),
                view: vk::ImageView::null(),
                inner: Arc::new(Inner {
                    data: self.inner.data.clone(),
                    desc,
                    extent: self.inner.extent,
                    usage: self.inner.usage,
                    owner: self.inner.owner.clone(),
                }),
            };
        }

        let desc = ViewDesc {
            base_layer: desc.base_layer + self.inner.desc.base_layer,
            base_level: desc.base_level + self.inner.desc.base_level,
            ..desc
        };

        if self.inner.desc == desc {
            return self.clone();
        }

        let view = match self.inner.data.views.lock().entry(desc) {
            Entry::Occupied(entry) => entry.get().0,
            Entry::Vacant(entry) => {
                match device.new_image_view(self.handle, self.inner.extent.into_ash(), desc) {
                    Ok((view, idx)) => entry.insert((view, idx)).0,
                    Err(OutOfMemory) => {
                        device.set_oom();
                        return Image {
                            handle: vk::Image::null(),
                            view: vk::ImageView::null(),
                            inner: Arc::new(Inner {
                                data: self.inner.data.clone(),
                                desc,
                                extent: self.inner.extent,
                                usage: self.inner.usage,
                                owner: self.inner.owner.clone(),
                            }),
                        };
                    }
                }
            }
        };

        Image {
            handle: self.handle,
            view,
            inner: Arc::new(Inner {
                data: self.inner.data.clone(),
                desc,
                extent: self.inner.extent,
                usage: self.inner.usage,
                owner: self.inner.owner.clone(),
            }),
        }
    }

    #[inline]
    pub(super) fn handle(&self) -> vk::Image {
        self.handle
    }

    #[inline]
    pub(super) fn view_handle(&self) -> vk::ImageView {
        self.view
    }

    #[inline]
    pub(super) fn base_layer(&self) -> u32 {
        self.inner.desc.base_layer
    }

    #[inline]
    pub(super) fn base_level(&self) -> u32 {
        self.inner.desc.base_level
    }
}

impl crate::traits::Resource for Image {}

#[hidden_trait::expose]
impl crate::traits::Image for Image {
    #[inline]
    fn format(&self) -> PixelFormat {
        self.inner.desc.format
    }

    #[inline]
    fn extent(&self) -> ImageExtent {
        self.inner.extent
    }

    #[inline]
    fn layers(&self) -> u32 {
        self.inner.desc.layers
    }

    #[inline]
    fn levels(&self) -> u32 {
        self.inner.desc.levels
    }

    #[inline]
    fn usage(&self) -> ImageUsage {
        self.inner.usage
    }

    #[inline]
    fn view(&self, device: &Device, desc: ViewDesc) -> Image {
        self.get_view(device, desc)
    }

    #[inline]
    fn detached(&self) -> bool {
        // If strong is 1, it cannot be changed by another thread if called owns
        // mutable reference to self
        // since there are no weaks.
        debug_assert_eq!(Arc::weak_count(&self.inner), 0, "No weak refs allowed");
        debug_assert_eq!(Arc::weak_count(&self.inner.data), 0, "No weak refs allowed");
        Arc::strong_count(&self.inner) == 1 && Arc::strong_count(&self.inner.data) == 1
    }
}

impl ArgumentsField<Automatic> for Image {
    const KIND: ArgumentKind = <Self as ArgumentsField<Sampled>>::KIND;
    const SIZE: usize = <Self as ArgumentsField<Sampled>>::SIZE;
    const OFFSET: usize = <Self as ArgumentsField<Sampled>>::OFFSET;
    const STRIDE: usize = <Self as ArgumentsField<Sampled>>::STRIDE;

    type Update = <Self as ArgumentsField<Sampled>>::Update;

    #[inline]
    fn update(&self) -> <Self as ArgumentsField<Sampled>>::Update {
        <Self as ArgumentsField<Sampled>>::update(self)
    }

    #[inline]
    fn add_refs(&self, refs: &mut Refs) {
        refs.add_image(self.clone());
    }
}

impl ArgumentsField<Sampled> for Image {
    const KIND: ArgumentKind = ArgumentKind::SampledImage;
    const SIZE: usize = 1;
    const OFFSET: usize = 0;
    const STRIDE: usize = size_of::<vk::DescriptorImageInfo>();

    type Update = vk::DescriptorImageInfo;

    #[inline]
    fn update(&self) -> vk::DescriptorImageInfo {
        vk::DescriptorImageInfo {
            sampler: vk::Sampler::null(),
            image_view: self.view,
            image_layout: vk::ImageLayout::GENERAL,
        }
    }

    #[inline]
    fn add_refs(&self, refs: &mut Refs) {
        refs.add_image(self.clone());
    }
}

impl ArgumentsField<Storage> for Image {
    const KIND: ArgumentKind = ArgumentKind::StorageImage;
    const SIZE: usize = 1;
    const OFFSET: usize = 0;
    const STRIDE: usize = size_of::<vk::DescriptorImageInfo>();

    type Update = vk::DescriptorImageInfo;

    #[inline]
    fn update(&self) -> vk::DescriptorImageInfo {
        vk::DescriptorImageInfo {
            sampler: vk::Sampler::null(),
            image_view: self.view,
            image_layout: vk::ImageLayout::GENERAL,
        }
    }

    #[inline]
    fn add_refs(&self, refs: &mut Refs) {
        refs.add_image(self.clone());
    }
}