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::WindowSurface};

use crate::{RenderDevice, 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, Resource)]
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 {
    pub surface: WindowSurface,
}

impl Deref for SurfaceData {
    type Target = WindowSurface;

    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: ResMut<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: Res<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: ResMut<Surfaces>,
    device: Res<RenderDevice>,
) {
    for window in windows.windows.values() {
        let surface = surfaces.surfaces.entry(window.entity).or_insert_with(|| {
            let (display_handle, window_handle) = (
                window.handle.get_display_handle(),
                window.handle.get_window_handle(),
            );

            let surface = WindowSurface::new(
                (**device).clone(),
                vk::Extent2D {
                    width: window.physical_width,
                    height: window.physical_height,
                },
                window_handle,
                display_handle,
            )
            .expect("Failed to create window surface");

            SurfaceData { surface }
        });

        if window.size_changed {
            surface
                .surface
                .recreate_with(Some(vk::Extent2D {
                    width: window.physical_width,
                    height: window.physical_height,
                }))
                .expect("Failed to recreate surface");
        }
        surfaces.configured.insert(window.entity);
    }
}
