use std::{
    collections::HashSet,
    ops::{Deref, DerefMut},
};

use bevy_ecs::{entity::EntityHashMap, prelude::*};
use bevy_window::{PrimaryWindow, RawHandleWrapper, Window, WindowClosing};
use renderer::{ash::vk, swapchain::Surface};

use crate::extract_param::Extract;

pub struct ExtractedWindow {
    pub entity: Entity,
    pub handle: RawHandleWrapper,
    pub physical_width: u32,
    pub physical_height: u32,
    pub size_changed: bool,
    /// On wayland, windows will not show unless they are presented at least once.
    pub needs_initial_present: bool,
}

#[derive(Default, Resource)]
pub struct ExtractedWindows {
    pub primary: Option<Entity>,
    pub windows: EntityHashMap<ExtractedWindow>,
}

impl Deref for ExtractedWindows {
    type Target = EntityHashMap<ExtractedWindow>;

    fn deref(&self) -> &Self::Target {
        &self.windows
    }
}

impl DerefMut for ExtractedWindows {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.windows
    }
}

#[derive(Default)]
pub struct Surfaces {
    surfaces: EntityHashMap<SurfaceData>,
    configured: HashSet<Entity>,
}

impl Surfaces {
    fn remove(&mut self, entity: &Entity) {
        self.surfaces.remove(entity);
        self.configured.remove(entity);
    }

    pub fn iter(&self) -> impl Iterator<Item = (&Entity, &SurfaceData)> {
        self.surfaces.iter()
    }
}

pub struct SurfaceData {
    renderer: renderer::Renderer2,
    pub surface: Surface,
}

impl Deref for SurfaceData {
    type Target = Surface;

    fn deref(&self) -> &Self::Target {
        &self.surface
    }
}

#[allow(clippy::type_complexity)]
pub(crate) fn extract_windows(
    mut extracted_windows: ResMut<ExtractedWindows>,
    mut closing: Extract<MessageReader<WindowClosing>>,
    windows: Extract<Query<(Entity, &RawHandleWrapper, &Window, Option<&PrimaryWindow>)>>,
    mut removed: Extract<RemovedComponents<RawHandleWrapper>>,
    mut surfaces: NonSendMut<Surfaces>,
) {
    for (entity, handle, window, primary) in windows.iter() {
        if primary.is_some() {
            extracted_windows.primary = Some(entity);
        }

        let (new_width, new_height) = (
            window.resolution.physical_width().max(1),
            window.resolution.physical_height().max(1),
        );

        let extracted_window = extracted_windows.entry(entity).or_insert(ExtractedWindow {
            entity,
            handle: handle.clone(),
            physical_width: new_width,
            physical_height: new_height,
            size_changed: false,
            needs_initial_present: true,
        });

        if extracted_window.physical_width != new_width
            || extracted_window.physical_height != new_height
        {
            extracted_window.physical_width = new_width;
            extracted_window.physical_height = new_height;
            extracted_window.size_changed = true;
        }
    }

    for closing in closing.read() {
        extracted_windows.remove(&closing.window);
        surfaces.remove(&closing.window);
    }
    for removed in removed.read() {
        extracted_windows.remove(&removed);
        surfaces.remove(&removed);
    }
}

pub(crate) fn need_surface_configuration(
    windows: Res<ExtractedWindows>,
    surfaces: NonSend<Surfaces>,
) -> bool {
    windows
        .windows
        .values()
        .any(|window| !surfaces.configured.contains(&window.entity) || window.size_changed)
}

pub(crate) fn configure_surfaces(
    windows: Res<ExtractedWindows>,
    mut surfaces: NonSendMut<Surfaces>,
) {
    for window in windows.windows.values() {
        let surface_data = surfaces.surfaces.entry(window.entity).or_insert_with(|| {
            let (display_handle, window_handle) = (
                window.handle.get_display_handle(),
                window.handle.get_window_handle(),
            );

            let renderer =
                renderer::Renderer2::new(display_handle).expect("Failed to create Vulkan renderer");
            let surface = renderer
                .create_surface(
                    window_handle,
                    vk::Extent2D {
                        width: window.physical_width,
                        height: window.physical_height,
                    },
                )
                .expect("Failed to create window surface");

            SurfaceData { renderer, surface }
        });

        if window.size_changed {
            surface_data
                .surface
                .configure(
                    surface_data.renderer.device(),
                    renderer::swapchain::SwapchainConfiguration {
                        extent: vk::Extent2D {
                            width: window.physical_width,
                            height: window.physical_height,
                        },
                        ..Default::default()
                    },
                )
                .expect("Failed to recreate surface");
        }
        surfaces.configured.insert(window.entity);
    }
}