use crate::prelude::*;

use anyhow::Result;

use std::cmp;
use std::sync::{
    atomic::{AtomicU32, Ordering::SeqCst},
    Arc, Mutex,
};

pub enum NextImageSynchronization<'a> {
    Semaphore(&'a Arc<Semaphore>),
    Fence(&'a Arc<Fence>),
}

impl<'a> From<&'a Arc<Semaphore>> for NextImageSynchronization<'a> {
    fn from(value: &'a Arc<Semaphore>) -> Self {
        Self::Semaphore(value)
    }
}

impl<'a> From<&'a Arc<Fence>> for NextImageSynchronization<'a> {
    fn from(value: &'a Arc<Fence>) -> Self {
        Self::Fence(value)
    }
}

#[derive(Debug)]
pub struct Swapchain {
    width: AtomicU32,
    height: AtomicU32,
    index: AtomicU32,

    device: Arc<Device>,
    surface: Arc<Surface>,

    create_info: Mutex<VkSwapchainCreateInfoKHR>,
    swapchain: Mutex<VkSwapchainKHR>,
    usage: VkImageUsageFlagBits,

    raw: bool,
}

impl Swapchain {
    pub fn new(
        device: Arc<Device>,
        surface: &Arc<Surface>,
        vsync: bool,
        image_count: u32,
        image_usage: impl Into<VkImageUsageFlagBits>,
        prefered_format: VkFormat,
        array_layers: u32,
    ) -> Result<Arc<Swapchain>> {
        let surface_caps = surface.capabilities(&device)?;

        let extent = if surface_caps.currentExtent.width == u32::max_value() {
            return Err(anyhow::Error::msg("Surface has no extent"));
        } else {
            VkExtent2D {
                width: surface_caps.currentExtent.width,
                height: surface_caps.currentExtent.height,
            }
        };

        let mut present_mode = VK_PRESENT_MODE_FIFO_KHR;

        if !vsync {
            for present_mode_iter in surface.present_modes(&device)? {
                if present_mode_iter == VK_PRESENT_MODE_MAILBOX_KHR {
                    present_mode = VK_PRESENT_MODE_MAILBOX_KHR;
                    break;
                } else if present_mode_iter == VK_PRESENT_MODE_IMMEDIATE_KHR {
                    present_mode = VK_PRESENT_MODE_IMMEDIATE_KHR;
                }
            }
        }

        let swapchain_image_count = if surface_caps.maxImageCount < surface_caps.minImageCount {
            cmp::max(image_count, surface_caps.minImageCount)
        } else {
            cmp::max(
                cmp::min(image_count, surface_caps.maxImageCount),
                surface_caps.minImageCount,
            )
        };

        let pretransform =
            if (surface_caps.supportedTransforms & VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR) != 0 {
                VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR.into()
            } else {
                surface_caps.currentTransform
            };

        let (format, colorspace) = surface.format_colorspace(&device, prefered_format)?;

        let swapchain_ci = VkSwapchainCreateInfoKHR::new(
            0,
            surface.vk_handle(),
            swapchain_image_count,
            format,
            colorspace,
            extent,
            array_layers,
            image_usage,
            VK_SHARING_MODE_EXCLUSIVE,
            &[],
            pretransform,
            VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
            present_mode,
            device.physical_device().features().shaderClipDistance,
        );

        let swapchain = device.create_swapchain(&swapchain_ci)?;

        Ok(Arc::new(Swapchain {
            width: AtomicU32::new(extent.width),
            height: AtomicU32::new(extent.height),
            usage: swapchain_ci.imageUsage,
            index: AtomicU32::new(0),

            device,
            surface: surface.clone(),

            create_info: Mutex::new(swapchain_ci),

            swapchain: Mutex::new(swapchain),

            raw: false,
        }))
    }

    pub fn from_ci(
        device: Arc<Device>,
        swapchain_ci: &VkSwapchainCreateInfoKHR,
    ) -> Result<Arc<Self>> {
        Ok(Self::from_raw(
            device.clone(),
            swapchain_ci,
            device.create_swapchain(swapchain_ci)?,
        ))
    }

    pub fn from_raw(
        device: Arc<Device>,
        swapchain_ci: &VkSwapchainCreateInfoKHR,
        swapchain: VkSwapchainKHR,
    ) -> Arc<Self> {
        Arc::new(Swapchain {
            width: AtomicU32::new(swapchain_ci.imageExtent.width),
            height: AtomicU32::new(swapchain_ci.imageExtent.height),
            usage: swapchain_ci.imageUsage,
            index: AtomicU32::new(0),

            surface: Surface::from_vk_surface(
                swapchain_ci.surface,
                device.physical_device().instance(),
            ),
            device,

            create_info: Mutex::new(swapchain_ci.clone()),

            swapchain: Mutex::new(swapchain),

            raw: true,
        })
    }

    pub fn recreate(&self) -> Result<()> {
        // wait for the device to get idle
        self.device.wait_idle()?;

        let surface_caps = self.surface.capabilities(&self.device)?;

        let extent = if surface_caps.currentExtent.width == u32::max_value()
            || surface_caps.currentExtent.height == u32::max_value()
        {
            return Err(anyhow::Error::msg("Surface has no extent"));
        } else if surface_caps.currentExtent.width == 0 || surface_caps.currentExtent.height == 0 {
            // don't recreate swapchain
            return Ok(());
        } else {
            VkExtent2D {
                width: surface_caps.currentExtent.width,
                height: surface_caps.currentExtent.height,
            }
        };

        let mut swapchain_ci = self.create_info.lock().unwrap();
        swapchain_ci.imageExtent = extent;
        swapchain_ci.set_old_swapchain(*self.swapchain.lock().unwrap());

        let swapchain = self.device.create_swapchain(&swapchain_ci)?;

        // destroy the old swapchain
        self.destroy();

        // replace swapchain
        *self.swapchain.lock().unwrap() = swapchain;

        // set new surface size
        self.width.store(extent.width, SeqCst);
        self.height.store(extent.height, SeqCst);

        Ok(())
    }

    pub fn acquire_next_image<'a>(
        &self,
        time_out: u64,
        synchro: impl Into<NextImageSynchronization<'a>>,
    ) -> Result<OutOfDate<u32>> {
        let synchro = synchro.into();

        let res = self.device.acquire_next_image(
            *self.swapchain.lock().unwrap(),
            time_out,
            match synchro {
                NextImageSynchronization::Semaphore(semaphore) => Some(semaphore.vk_handle()),
                NextImageSynchronization::Fence(_) => None,
            },
            match synchro {
                NextImageSynchronization::Semaphore(_) => None,
                NextImageSynchronization::Fence(fence) => Some(fence.vk_handle()),
            },
        );

        if let Ok(r) = &res {
            if let OutOfDate::Ok(i) = r {
                self.index.store(*i, SeqCst);
            }
        }

        res
    }

    /// set current
    /// only use when externally acquired next index !!!
    pub unsafe fn set_image_index(&self, index: u32) {
        self.index.store(index, SeqCst);
    }

    pub fn current_index(&self) -> u32 {
        self.index.load(SeqCst)
    }

    pub fn vk_images(&self) -> Result<Vec<VkImage>> {
        self.device
            .swapchain_images(*self.swapchain.lock().unwrap())
    }

    pub fn wrap_images(
        &self,
        images: &[VkImage],
        queue: &Arc<Mutex<Queue>>,
        assume_layout: bool,
    ) -> Result<Vec<Arc<Image>>> {
        let format = self.format();
        let tiling = VK_IMAGE_TILING_OPTIMAL;

        if !Image::check_configuration(&self.device, tiling, format, self.usage) {
            return Err(anyhow::Error::msg(format!(
                "Image configuration not allowed (tiling: {:?}, format: {:?}, usage: {:?})",
                tiling, format, self.usage,
            )));
        }

        let mut swapchain_images = Vec::new();

        for image in images {
            swapchain_images.push(
                Image::from_preinitialized(
                    *image,
                    format,
                    self.width(),
                    self.height(),
                    VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
                    self.usage,
                    assume_layout,
                )
                .attach_sampler(Sampler::nearest_sampler().build(&self.device)?)
                .build(&self.device, queue)?,
            );
        }

        Ok(swapchain_images)
    }

    pub fn width(&self) -> u32 {
        self.width.load(SeqCst)
    }

    pub fn height(&self) -> u32 {
        self.height.load(SeqCst)
    }

    pub fn format(&self) -> VkFormat {
        self.create_info.lock().unwrap().imageFormat
    }

    #[inline]
    fn destroy(&self) {
        self.device
            .destroy_swapchain(*self.swapchain.lock().unwrap())
    }
}

impl VulkanDevice for Swapchain {
    fn device(&self) -> &Arc<Device> {
        &self.device
    }
}

impl VkHandle<VkSwapchainKHR> for Swapchain {
    fn vk_handle(&self) -> VkSwapchainKHR {
        *self.swapchain.lock().unwrap()
    }
}

impl<'a> VkHandle<VkSwapchainKHR> for &'a Swapchain {
    fn vk_handle(&self) -> VkSwapchainKHR {
        *self.swapchain.lock().unwrap()
    }
}

impl VkHandle<VkSwapchainKHR> for Arc<Swapchain> {
    fn vk_handle(&self) -> VkSwapchainKHR {
        *self.swapchain.lock().unwrap()
    }
}

impl<'a> VkHandle<VkSwapchainKHR> for &'a Arc<Swapchain> {
    fn vk_handle(&self) -> VkSwapchainKHR {
        *self.swapchain.lock().unwrap()
    }
}

impl Drop for Swapchain {
    fn drop(&mut self) {
        if !self.raw {
            self.destroy();
        }
    }
}