use std::{
    marker::PhantomData,
    sync::{
        Arc,
        atomic::{AtomicU32, AtomicU64, Ordering},
    },
};

use ash::{
    khr,
    prelude::VkResult,
    vk::{self, Handle},
};
use parking_lot::{RawMutex, RwLock};
use raw_window_handle::{RawDisplayHandle, RawWindowHandle};

use crate::{
    Instance, Result, define_device_owned_handle,
    device::{Device, DeviceOwned},
    images, sync,
    util::RawMutexGuard,
};

use derive_more::Debug;
#[derive(Debug)]
pub struct Surface {
    raw: vk::SurfaceKHR,
    #[debug(skip)]
    functor: khr::surface::Instance,
}

impl Drop for Surface {
    fn drop(&mut self) {
        unsafe {
            self.functor.destroy_surface(self.raw, None);
        }
    }
}

impl Surface {
    #[allow(dead_code)]
    pub fn headless(instance: &Arc<Instance>) -> Result<Self> {
        let headless_instance =
            ash::ext::headless_surface::Instance::new(&instance.entry, &instance.instance);
        let functor = khr::surface::Instance::new(&instance.entry, &instance.instance);

        // SAFETY: the headless surface does not have any platform-specific requirements, and does not depend on any external handles, so it is safe to create without additional guarantees.
        // (note): ash marks this function as unsafe, likely because of
        // auto-generated bindings from vk.xml, but doesn't provide any safety
        // bounds.
        unsafe {
            let raw = headless_instance
                .create_headless_surface(&vk::HeadlessSurfaceCreateInfoEXT::default(), None)?;

            Ok(Self { raw, functor })
        }
    }

    /// # Safety
    ///
    /// The caller must ensure that the provided display and window handles are
    /// valid and remain valid for the lifetime of the surface. Namely, the
    /// window handle must refer to a valid window that is associated with the
    /// display handle, and both must not be destroyed while the surface is
    /// still in use.  Additionally, the caller must ensure that the instance
    /// was created with the appropriate platform-specific surface extensions
    /// enabled.
    pub unsafe fn new_from_raw_window_handle(
        instance: &Arc<Instance>,
        display_handle: RawDisplayHandle,
        window_handle: RawWindowHandle,
    ) -> Result<Self> {
        let functor = khr::surface::Instance::new(&instance.entry, &instance.instance);
        // SAFETY: the caller guarantees the validity of the display and window handles, and that they remain valid for the lifetime of the surface.
        let surface = unsafe {
            ash_window::create_surface(
                &instance.entry,
                &instance.instance,
                display_handle,
                window_handle,
                None,
            )?
        };

        Ok(Self {
            raw: surface,
            functor,
        })
    }
}

define_device_owned_handle! {
    pub Swapchain(vk::SwapchainKHR) {
        mutex: RawMutex,
        surface: Arc<Surface>,

        #[allow(unused)]
        present_mode: vk::PresentModeKHR,
        #[allow(unused)]
        color_space: vk::ColorSpaceKHR,
        format: vk::Format,
        images: Vec<Arc<images::Image>>,
        image_views: Vec<vk::ImageView>,
        extent: vk::Extent2D,
        min_image_count: u32,

        // sync objects:
        // we need two semaphores per each image, one acquire-semaphore and one release-semaphore.
        // semaphores must be unique to each frame and cannot be reused per swapchain.
        acquire_semaphores: Vec<vk::Semaphore>,
        release_semaphores: Vec<vk::Semaphore>,

        // one fence per in-flight frame, to synchronize image acquisition
        fences: Vec<Arc<sync::Fence>>,

        current_frame: AtomicU32,

        // for khr_present_id/khr_present_wait
        #[allow(unused)]
        present_id: AtomicU64,
    } => |this| unsafe {
        _ = this.device().wait_queue_idle(this.device().present_queue());
        tracing::debug!("dropping swapchain {:?}", this.handle());
        for view in &this.image_views {
            this.device().dev().destroy_image_view(*view, None);
        }

        this.with_locked(|swapchain| {
            this.device().swapchain().destroy_swapchain(swapchain, None)
        });

        for &semaphore in this
            .acquire_semaphores
            .iter()
            .chain(&this.release_semaphores)
        {
            this.device().dev().destroy_semaphore(semaphore, None);
        }
    }
}

impl core::fmt::Debug for Swapchain {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Swapchain")
            .field("inner", &self.inner)
            .field("present_mode", &self.present_mode)
            .field("color_space", &self.color_space)
            .field("format", &self.format)
            .field("images", &self.images)
            .field("image_views", &self.image_views)
            .field("extent", &self.extent)
            .field("min_image_count", &self.min_image_count)
            .field("acquire_semaphores", &self.acquire_semaphores)
            .field("release_semaphores", &self.release_semaphores)
            .field("fences", &self.fences)
            .field("current_frame", &self.current_frame)
            .field("present_id", &self.present_id)
            .finish()
    }
}

impl Swapchain {
    const PREFERRED_IMAGES_IN_FLIGHT: u32 = 3;

    fn get_swapchain_params_from_surface(
        instance: &Arc<Instance>,
        surface: vk::SurfaceKHR,
        pdev: vk::PhysicalDevice,
        requested_extent: Option<vk::Extent2D>,
    ) -> Result<SwapchainParams> {
        let caps = unsafe {
            instance
                .surface
                .get_physical_device_surface_capabilities(pdev, surface)?
        };
        let formats = unsafe {
            instance
                .surface
                .get_physical_device_surface_formats(pdev, surface)?
        };
        let present_modes = unsafe {
            instance
                .surface
                .get_physical_device_surface_present_modes(pdev, surface)?
        };

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

        let format = formats
            .iter()
            .max_by_key(|&&format| {
                let is_rgba_unorm = format.format == vk::Format::R8G8B8A8_UNORM
                    || format.format == vk::Format::B8G8R8A8_UNORM;
                let is_srgb = format.color_space == vk::ColorSpaceKHR::SRGB_NONLINEAR;
                is_rgba_unorm as u8 * 10 + is_srgb as u8
            })
            .or(formats.first())
            .cloned()
            .expect("no surface format available!");

        // 0 here means no limit
        let max_image_count = core::num::NonZero::new(caps.max_image_count)
            .map(|n| n.get())
            .unwrap_or(u32::MAX);

        // we want PREFERRED_IMAGES_IN_FLIGHT images acquired at the same time,
        let image_count =
            (caps.min_image_count + Self::PREFERRED_IMAGES_IN_FLIGHT).min(max_image_count);

        let extent = current_extent_or_clamped(
            &caps,
            requested_extent.unwrap_or(vk::Extent2D::default().width(1).height(1)),
        );

        Ok(SwapchainParams {
            present_mode,
            format: format.format,
            color_space: format.color_space,
            image_count,
            extent,
            min_image_count: caps.min_image_count,
        })
    }

    pub fn new(
        device: Device,
        surface: Arc<Surface>,
        pdev: vk::PhysicalDevice,
        extent: vk::Extent2D,
    ) -> Result<Self> {
        Self::create(device, surface, pdev, Some(extent), None)
    }

    fn create(
        device: Device,
        surface: Arc<Surface>,
        pdev: vk::PhysicalDevice,
        extent: Option<vk::Extent2D>,
        old_swapchain: Option<&Self>,
    ) -> Result<Self> {
        let SwapchainParams {
            present_mode,
            format,
            color_space,
            image_count,
            min_image_count,
            extent,
        } = Self::get_swapchain_params_from_surface(device.instance(), surface.raw, pdev, extent)?;

        let (swapchain, images) = {
            let lock = old_swapchain.as_ref().map(|handle| handle.lock());

            Self::create_vkswapchainkhr(
                &device,
                surface.raw,
                &device.queue_families().swapchain_family_indices(),
                extent,
                lock.as_ref().map(|lock| **lock),
                present_mode,
                format,
                color_space,
                image_count,
            )
        }?;

        let images = images
            .iter()
            .enumerate()
            .map(|(i, image)| unsafe {
                images::Image::from_swapchain_image(
                    device.clone(),
                    *image,
                    Some(format!("swapchain-{:x}-image-{i}", swapchain.as_raw()).into()),
                    vk::Extent3D {
                        width: extent.width,
                        height: extent.height,
                        depth: 1,
                    },
                    format,
                )
                .inspect(|img| {
                    _ = img.get_view(images::ImageViewDesc {
                        // TODO: make this a function that uses debug name/surface handle
                        name: Some(
                            format!("swapchain-{:x}-image-{i}-view", swapchain.as_raw()).into(),
                        ),
                        kind: vk::ImageViewType::TYPE_2D,
                        format,
                        aspect: vk::ImageAspectFlags::COLOR,
                        ..Default::default()
                    });
                })
                .map(Arc::new)
            })
            .collect::<VkResult<Vec<_>>>()?;

        let image_views = images
            .iter()
            .enumerate()
            .map(|(i, image)| {
                image.get_view(images::ImageViewDesc {
                    name: Some(format!("swapchain-{:x}-image-{i}-view", swapchain.as_raw()).into()),
                    kind: vk::ImageViewType::TYPE_2D,
                    format,
                    aspect: vk::ImageAspectFlags::COLOR,
                    ..Default::default()
                })
            })
            .collect::<VkResult<Vec<_>>>()?;

        let num_images = images.len() as u32;
        let inflight_frames = num_images - min_image_count;

        let acquire_semaphores = {
            (0..inflight_frames)
                .map(|i| unsafe {
                    device
                        .dev()
                        .create_semaphore(&vk::SemaphoreCreateInfo::default(), None)
                        .inspect(|r| {
                            #[cfg(debug_assertions)]
                            {
                                device
                                    .debug_name_object(
                                        *r,
                                        &format!("semaphore-{:x}_{i}-acquire", swapchain.as_raw()),
                                    )
                                    .unwrap();
                            }
                        })
                })
                .collect::<VkResult<Vec<_>>>()?
        };

        let release_semaphores = {
            (0..inflight_frames)
                .map(|i| unsafe {
                    device
                        .dev()
                        .create_semaphore(&vk::SemaphoreCreateInfo::default(), None)
                        .inspect(|r| {
                            #[cfg(debug_assertions)]
                            {
                                device
                                    .debug_name_object(
                                        *r,
                                        &format!("semaphore-{:x}_{i}-release", swapchain.as_raw()),
                                    )
                                    .unwrap();
                            }
                        })
                })
                .collect::<VkResult<Vec<_>>>()?
        };

        let fences = {
            (0..inflight_frames)
                .map(|i| {
                    Ok(Arc::new(sync::Fence::create(device.clone()).inspect(
                        |r| {
                            #[cfg(debug_assertions)]
                            {
                                device
                                    .debug_name_object(
                                        r.fence(),
                                        &format!("fence-{:x}_{i}", swapchain.as_raw()),
                                    )
                                    .unwrap();
                            }
                        },
                    )?))
                })
                .collect::<VkResult<Vec<_>>>()?
        };

        tracing::trace!("fences: {fences:?}");

        Ok(Self::construct(
            device,
            swapchain,
            Some(
                format!(
                    "swapchain-{}_{}",
                    surface.raw.as_raw(),
                    SWAPCHAIN_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
                )
                .into(),
            ),
            <parking_lot::RawMutex as parking_lot::lock_api::RawMutex>::INIT,
            surface,
            present_mode,
            color_space,
            format,
            images,
            image_views,
            extent,
            min_image_count,
            acquire_semaphores,
            release_semaphores,
            fences,
            AtomicU32::new(0),
            AtomicU64::new(1),
        )?)
    }

    pub fn max_in_flight_images(&self) -> u32 {
        self.num_images() - self.min_image_count
    }

    pub fn num_images(&self) -> u32 {
        self.images.len() as u32
    }

    fn recreate(&self, extent: Option<vk::Extent2D>) -> Result<Self> {
        Self::create(
            self.device().clone(),
            self.surface.clone(),
            self.device().phy(),
            extent,
            Some(self),
        )
    }

    /// returns a future yielding the frame, and true if the swapchain is
    /// suboptimal and should be recreated.
    fn acquire_image(
        self: Arc<Self>,
    ) -> impl std::future::Future<Output = VkResult<(SwapchainFrame, bool)>> {
        let frame = self
            .current_frame
            .try_update(Ordering::Release, Ordering::Relaxed, |i| {
                Some((i + 1) % self.max_in_flight_images())
            })
            .unwrap() as usize;

        tracing::trace!(frame, "acquiring image for frame {frame}");

        async move {
            let fence = self.fences[frame].clone();
            let acquire = self.acquire_semaphores[frame];
            let release = self.release_semaphores[frame];

            // spawn on threadpool because it might block.
            let (idx, suboptimal) = smol::unblock({
                let this = self.clone();
                let fence = fence.clone();
                move || unsafe {
                    this.with_locked(|swapchain| {
                        this.device().swapchain().acquire_next_image(
                            swapchain,
                            u64::MAX,
                            acquire,
                            fence.fence(),
                        )
                    })
                }
            })
            .await?;

            // wait for image to become available.
            sync::FenceFuture::new(fence.clone()).await;

            let idx = idx as usize;
            let image = self.images[idx].clone();
            let view = self.image_views[idx];

            Ok((
                SwapchainFrame {
                    index: idx as u32,
                    swapchain: self.clone(),
                    format: self.format,
                    image,
                    view,
                    acquire,
                    release,
                },
                suboptimal,
            ))
        }
    }

    fn present(&self, frame: SwapchainFrame, wait: Option<vk::Semaphore>) -> Result<()> {
        let swpchain = self.lock();
        let queue = self.device().present_queue().lock();

        let wait_semaphores = wait.as_slice();

        // TODO: make this optional for devices with no support for present_wait/present_id
        // let present_id = self
        //     .present_id
        //     .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        // let mut present_id =
        //     vk::PresentIdKHR::default().present_ids(core::slice::from_ref(&present_id));

        let present_info = vk::PresentInfoKHR::default()
            .image_indices(core::slice::from_ref(&frame.index))
            .swapchains(core::slice::from_ref(&swpchain))
            .wait_semaphores(wait_semaphores);
        //.push_next(&mut present_id)

        // call winits pre_present_notify here

        unsafe {
            self.device()
                .swapchain()
                .queue_present(*queue, &present_info)?;
        }
        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    fn create_vkswapchainkhr(
        device: &Device,
        surface: vk::SurfaceKHR,
        queue_families: &[u32],
        image_extent: vk::Extent2D,
        old_swapchain: Option<vk::SwapchainKHR>,
        present_mode: vk::PresentModeKHR,
        image_format: vk::Format,
        image_color_space: vk::ColorSpaceKHR,
        image_count: u32,
    ) -> Result<(vk::SwapchainKHR, Vec<vk::Image>)> {
        let create_info = vk::SwapchainCreateInfoKHR::default()
            .surface(surface)
            .present_mode(present_mode)
            .image_color_space(image_color_space)
            .image_format(image_format)
            .min_image_count(image_count)
            .image_usage(vk::ImageUsageFlags::TRANSFER_DST | vk::ImageUsageFlags::COLOR_ATTACHMENT)
            .image_array_layers(1)
            .image_extent(image_extent)
            .image_sharing_mode(if queue_families.len() <= 1 {
                vk::SharingMode::EXCLUSIVE
            } else {
                vk::SharingMode::CONCURRENT
            })
            .queue_family_indices(queue_families)
            .pre_transform(vk::SurfaceTransformFlagsKHR::IDENTITY)
            .composite_alpha(vk::CompositeAlphaFlagsKHR::OPAQUE)
            .old_swapchain(old_swapchain.unwrap_or(vk::SwapchainKHR::null()))
            .clipped(true);

        let (swapchain, images) = unsafe {
            let swapchain = device.swapchain().create_swapchain(&create_info, None)?;

            let images = device.swapchain().get_swapchain_images(swapchain)?;

            (swapchain, images)
        };

        Ok((swapchain, images))
    }
}
static SWAPCHAIN_COUNT: AtomicU64 = AtomicU64::new(0);

#[derive(Debug)]
#[must_use = "This struct represents an acquired image from the swapchain and
 must be presented in order to free resources on the device."]
pub struct SwapchainFrame {
    pub swapchain: Arc<Swapchain>,
    pub index: u32,
    pub image: Arc<images::Image>,
    pub format: vk::Format,
    pub view: vk::ImageView,
    pub acquire: vk::Semaphore,
    pub release: vk::Semaphore,
}

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

impl SwapchainFrame {
    pub fn present(self, wait: Option<vk::Semaphore>) -> crate::Result<()> {
        self.swapchain.clone().present(self, wait)
    }
}

fn current_extent_or_clamped(
    caps: &vk::SurfaceCapabilitiesKHR,
    fallback: vk::Extent2D,
) -> vk::Extent2D {
    if caps.current_extent.width == u32::MAX {
        vk::Extent2D {
            width: fallback
                .width
                .clamp(caps.min_image_extent.width, caps.max_image_extent.width),
            height: fallback
                .height
                .clamp(caps.min_image_extent.height, caps.max_image_extent.height),
        }
    } else {
        caps.current_extent
    }
}

struct SwapchainParams {
    present_mode: vk::PresentModeKHR,
    format: vk::Format,
    color_space: vk::ColorSpaceKHR,
    /// the number of images to request from the device
    image_count: u32,
    /// the minimum number of images the surface permits
    min_image_count: u32,
    extent: vk::Extent2D,
}

impl Swapchain {
    pub fn lock(&self) -> RawMutexGuard<'_, vk::SwapchainKHR> {
        use parking_lot::lock_api::RawMutex;
        self.mutex.lock();
        RawMutexGuard {
            mutex: &self.mutex,
            value: &self.inner.object,
            _pd: PhantomData,
        }
    }

    pub fn with_locked<T, F: FnOnce(vk::SwapchainKHR) -> T>(&self, f: F) -> T {
        let lock = self.lock();
        f(*lock)
    }
}

#[derive(Debug)]
pub struct WindowSurface {
    // window_handle: RawWindowHandle,
    pub surface: Arc<Surface>,
    // this mutex is for guarding the swapchain against being replaced
    // underneath WindowContext's functions
    current_swapchain: RwLock<Arc<Swapchain>>,
}

impl WindowSurface {
    pub fn new(
        device: Device,
        requested_extent: vk::Extent2D,
        window: RawWindowHandle,
        display: RawDisplayHandle,
    ) -> Result<Self> {
        let surface = Arc::new(unsafe {
            Surface::new_from_raw_window_handle(device.instance(), display, window)?
        });
        let swapchain = RwLock::new(Arc::new(Swapchain::new(
            device.clone(),
            surface.clone(),
            device.phy(),
            requested_extent,
        )?));

        Ok(Self {
            surface,
            // window_handle: window,
            current_swapchain: swapchain,
        })
    }

    /// spawns a task that continuously requests images from the current
    /// swapchain, sending them to a channel.  returns the receiver of the
    /// channel, and a handle to the task, allowing for cancellation.
    pub fn images(
        self: Arc<Self>,
    ) -> (
        smol::channel::Receiver<SwapchainFrame>,
        smol::Task<std::result::Result<(), crate::Error>>,
    ) {
        let (tx, rx) = smol::channel::bounded(8);
        let task = smol::spawn(async move {
            loop {
                let frame = self.acquire_image().await?;
                tx.send(frame)
                    .await
                    .expect("channel closed on swapchain acquiring frame");
            }
        });

        (rx, task)
    }

    pub async fn acquire_image(&self) -> Result<SwapchainFrame> {
        // clone swapchain to keep it alive
        let swapchain = self.current_swapchain.read().clone();
        let (frame, suboptimal) = swapchain.clone().acquire_image().await?;
        if suboptimal {
            let mut lock = self.current_swapchain.write();
            // only recreate our swapchain if it is still same, or else it might have already been recreated.
            if Arc::ptr_eq(&swapchain, &lock) {
                *lock = Arc::new(lock.recreate(None)?);
            }
        }

        Ok(frame)
    }

    pub fn acquire_image_blocking(&self) -> Result<SwapchainFrame> {
        smol::block_on(self.acquire_image())
    }

    pub fn recreate_with(&self, extent: Option<vk::Extent2D>) -> Result<()> {
        let mut swapchain = self.current_swapchain.write();
        *swapchain = Arc::new(swapchain.recreate(extent)?);

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::device;

    use super::*;

    fn create_headless_vk() -> Result<(Device, WindowSurface)> {
        let device = Device::new_from_default_desc(
            None,
            &[
                device::Extension {
                    name: "VK_EXT_headless_surface",
                    version: ash::ext::headless_surface::SPEC_VERSION,
                },
                device::Extension {
                    name: "VK_KHR_surface",
                    version: ash::khr::surface::SPEC_VERSION,
                },
            ],
        )?;

        let surface = Arc::new(Surface::headless(device.instance())?);

        let swapchain = Arc::new(Swapchain::new(
            device.clone(),
            surface.clone(),
            device.phy(),
            vk::Extent2D::default().width(1).height(1),
        )?);

        let window_ctx = WindowSurface {
            // window_handle: RawWindowHandle::Web(raw_window_handle::WebWindowHandle::new(0)),
            surface,
            current_swapchain: RwLock::new(swapchain),
        };

        Ok((device, window_ctx))
    }

    #[tracing_test::traced_test]
    #[test]
    fn async_swapchain_acquiring() {
        let (_dev, ctx) = create_headless_vk().expect("init");
        let ctx = Arc::new(ctx);
        let (rx, handle) = ctx.clone().images();

        eprintln!("hello world!");

        let mut count = 0;
        loop {
            let now = std::time::Instant::now();
            let frame = rx.recv_blocking().expect("recv");

            _ = frame.present(None);
            tracing::info!("mspf: {:.3}ms", now.elapsed().as_micros() as f32 / 1e3);
            count += 1;
            if count > 1000 {
                smol::block_on(handle.cancel());
                break;
            }
        }
    }
}