use std::sync::Arc;
use crate::vulkan::{VkDevice, VkInstance, VkPhysicalDevice, VkSurface, LOG_TARGET};
use ash::prelude::VkResult;
use ash::vk;
use crate::display::Window;

pub struct VkSwapchain {
    surface: Arc<VkSurface>,
    device: Arc<VkDevice>,
    swapchain: Option<vk::SwapchainKHR>,

    pub(super) desired_image_count: u32,
    pub(super) surface_format: vk::SurfaceFormatKHR,
    pub(super) surface_resolution: vk::Extent2D,
    pub(super) present_mode: vk::PresentModeKHR,
    pub(super) pre_transform: vk::SurfaceTransformFlagsKHR,

    pub(super) present_images: Option<Vec<vk::Image>>,
    pub(super) present_image_views: Option<Vec<vk::ImageView>>,
}

impl VkSwapchain {
    pub(super) fn new<'a, 'b>(
        window: &Window,
        surface: Arc<VkSurface>,
        device: Arc<VkDevice>,
        physical_device: &VkPhysicalDevice,
    ) -> anyhow::Result<Self> {
        log::debug!(target: LOG_TARGET, "Creating swapchain");

        let (
            surface_formats,
            surface_capabilities,
            present_modes
        ) = surface.get_physical_device_surface_infos(physical_device)?;
        log::debug!(target: LOG_TARGET, "Supported surface formats by physical device: {surface_formats:#?}");
        log::debug!(target: LOG_TARGET, "Surface capabilities: {surface_capabilities:#?}");

        let surface_format = surface_formats
            .first()
            .and_then(|f| Some(*f))
            .ok_or_else(|| anyhow::anyhow!("No available surface formats"))?;
        log::debug!(target: LOG_TARGET, "Selected surface format: {surface_format:?}");

        let mut desired_image_count = surface_capabilities.min_image_count + 1;
        if surface_capabilities.max_image_count > 0
            && desired_image_count > surface_capabilities.max_image_count
        {
            desired_image_count = surface_capabilities.max_image_count;
        }
        log::debug!(target: LOG_TARGET, "Selected surface image count: {desired_image_count}");

        let window_size = window.size()
            .ok_or_else(|| anyhow::anyhow!("Window size is not valid"))?
            .to_physical::<u32>(1.0);
        log::debug!(target: LOG_TARGET, "Window size: {window_size:?}");

        let surface_resolution = match surface_capabilities.current_extent.width {
            u32::MAX => vk::Extent2D {
                width: window_size.width,
                height: window_size.height,
            },
            _ => surface_capabilities.current_extent,
        };
        log::debug!(target: LOG_TARGET, "Surface resolution: {surface_resolution:?}");

        let pre_transform = if surface_capabilities
            .supported_transforms
            .contains(vk::SurfaceTransformFlagsKHR::IDENTITY)
        {
            vk::SurfaceTransformFlagsKHR::IDENTITY
        } else {
            surface_capabilities.current_transform
        };

        let present_mode = present_modes
            .iter()
            .cloned()
            .find(|&mode| mode == vk::PresentModeKHR::MAILBOX)
            .unwrap_or(vk::PresentModeKHR::FIFO);

        let mut swapchain = Self {
            surface,
            device,
            desired_image_count,
            surface_format,
            surface_resolution,
            present_mode,
            pre_transform,
            swapchain: None,
            present_images: None,
            present_image_views: None,
        };

        swapchain.create_swapchain()?;

        Ok(swapchain)
    }

    pub(super) fn create_swapchain(&mut self) -> VkResult<()> {
        let mut swapchain_create_info = self.create_swapchain_info(&self.surface);

        if let Some(old_swapchain) = self.swapchain {
            swapchain_create_info.old_swapchain = old_swapchain;
        }

        let swapchain = unsafe {
            self.device.swapchain_loader.create_swapchain(&swapchain_create_info, None)?
        };

        let present_images = unsafe { self.device.swapchain_loader.get_swapchain_images(swapchain)? };
        let present_images_view = present_images
            .iter()
            .map(|i| {
                self.device.create_image_view(*i, self.surface_format)
                    .expect("Failed to create image view")
            })
            .collect::<Vec<_>>();

        if log::log_enabled!(target: LOG_TARGET, log::Level::Debug) {
            let label = match self.swapchain {
                None => "Swapchain created",
                Some(_) => "Swapchain updated"
            };
            log::debug!(target: LOG_TARGET, "{label} ({swapchain:?}) : {swapchain_create_info:#?}");
        }

        self.swapchain = Some(swapchain);
        self.present_image_views = Some(present_images_view);
        self.present_images = Some(present_images);

        Ok(())
    }

    pub(super) fn update_resolution(&mut self, width: u32, height: u32) -> VkResult<()> {
        log::debug!(target: LOG_TARGET, "New resolution requested for swapchain {width}x{height}");
        self.surface_resolution = vk::Extent2D {
            width,
            height,
        };

        self.create_swapchain()?;

        Ok(())
    }

    fn create_swapchain_info(&self, surface: &VkSurface) -> vk::SwapchainCreateInfoKHR {
        vk::SwapchainCreateInfoKHR::default()
            .surface(surface.surface)
            .min_image_count(self.desired_image_count)
            .image_color_space(self.surface_format.color_space)
            .image_format(self.surface_format.format)
            .image_extent(self.surface_resolution)
            .image_usage(vk::ImageUsageFlags::COLOR_ATTACHMENT)
            .image_sharing_mode(vk::SharingMode::EXCLUSIVE)
            .pre_transform(self.pre_transform)
            .composite_alpha(vk::CompositeAlphaFlagsKHR::OPAQUE)
            .present_mode(self.present_mode)
            .clipped(true)
            .image_array_layers(1)
    }

    fn drop_swapchain(&mut self) {
        if let Some(swapchain) = self.swapchain {
            unsafe { self.device.swapchain_loader.destroy_swapchain(swapchain, None); }
            self.swapchain = None;
            log::debug!(target: LOG_TARGET, "Swapchain destroyed ({swapchain:?})");
        }
    }
}

impl Drop for VkSwapchain {
    fn drop(&mut self) {
        self.drop_swapchain()
    }
}