use std::collections::BTreeMap;
use std::sync::LazyLock;
use std::{collections::HashMap, sync::Arc};

use terminator::Terminator;
use vulkano::buffer::BufferUsage;
use wgpu::util::{DeviceExt, RenderEncoder};
use wgpu::{
    Backends, BufferSlice, BufferUsages, FragmentState, IndexFormat, Instance, InstanceDescriptor,
    PipelineCompilationOptions, Surface,
};
use winit::application::ApplicationHandler;
use winit::event::{ElementState, MouseButton};
use winit::event_loop::{ActiveEventLoop, EventLoop};
use winit::window::{Window, WindowAttributes, WindowId};
use zlog::LogLevel;
use zlog::{config::LoggerConfig, query::LogQuery};
struct WindowContext<'window> {
    window: Arc<Window>,
    renderer: WgpuRenderer<'window>,
}
impl std::ops::Deref for WindowContext<'_> {
    type Target = winit::window::Window;

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

struct WgpuRenderer<'surface> {
    device: wgpu::Device,
    queue: wgpu::Queue,
    surface: wgpu::Surface<'surface>,
    surface_config: wgpu::SurfaceConfiguration,
    render_pipeline: wgpu::RenderPipeline,
}

// impl WgpuRenderer {
//     pub fn new(window: Arc<Window>) {}
struct App<'window> {
    state: WgpuState,
    windows: BTreeMap<WindowId, WindowContext<'window>>,
}

impl App<'_> {
    fn new() -> Self {
        Self {
            state: WgpuState::new(),
            windows: BTreeMap::new(),
        }
    }
    pub fn spawn_window(&mut self, event_loop: &ActiveEventLoop) {
        let attr = WindowAttributes::default()
            .with_title("Zenyx - SubWindow")
            .with_min_inner_size(winit::dpi::LogicalSize::new(1, 1));

        let window = event_loop.create_window(attr).unwrap();
        let window = Arc::new(window);
        let renderer = self.state.create_renderer(window.clone());
        let window_ctx = WindowContext {
            renderer: smol::block_on(renderer),
            window: window.clone(),
        };
        let window_id = window.id();
        self.windows.insert(window_id, window_ctx);
    }
}

struct WgpuState {
    instance: wgpu::Instance,
}
#[repr(C)]
struct Vertex {
    position: glm::Vec3,
    color: glm::Vec3,
}

impl Vertex {
    const ATTRIBS: [wgpu::VertexAttribute; 2] = [
        wgpu::VertexAttribute {
            offset: 0,
            shader_location: 0,
            format: wgpu::VertexFormat::Float32x3,
        },
        wgpu::VertexAttribute {
            offset: std::mem::offset_of!(Vertex, color) as u64,
            shader_location: 1,
            format: wgpu::VertexFormat::Float32x3,
        },
    ];

    const fn desc<'a>() -> wgpu::VertexBufferLayout<'a> {
        wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &Self::ATTRIBS,
        }
    }
}

impl WgpuState {
    fn new() -> Self {
        let backends = Backends::PRIMARY;
        let instance_descriptor = InstanceDescriptor {
            backends,
            ..Default::default()
        };
        let instance = Instance::new(&instance_descriptor);
        Self { instance }
    }

    async fn create_renderer<'surface>(&self, window: Arc<Window>) -> WgpuRenderer<'surface> {
        let surface = self.instance.create_surface(window.clone()).unwrap();
        let adapter = self
            .instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::default(),
                compatible_surface: Some(&surface),
                ..Default::default()
            })
            .await
            .unwrap();
        let device_descriptor = wgpu::DeviceDescriptor::default();
        let (device, queue) = adapter.request_device(&device_descriptor).await.unwrap();

        let size = window.inner_size();
        let width = size.width.max(1);
        let height = size.height.max(1);
        #[repr(align(4))]
        struct ShaderCode<const N: usize>([u8; N]);
        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("Pipeline Layout"),
            bind_group_layouts: &[],
            push_constant_ranges: &[],
        });

        let vert_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Vertex Shader"),
            source: unsafe {
                static shader_code: &[u8] =
                    &ShaderCode(*include_bytes!(concat!(env!("OUT_DIR"), "/vert.spv"))).0;
                // assert!(bytes.len() % 4 == 0);
                let shader = shader_code.align_to::<u32>().1;
                wgpu::ShaderSource::SpirV(std::borrow::Cow::Borrowed(shader))
            },
        });
        let frag_shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("Fragment Shader"),
            source: unsafe {
                static shader_code: &[u8] =
                    &ShaderCode(*include_bytes!(concat!(env!("OUT_DIR"), "/frag.spv"))).0;
                // assert!(bytes.len() % 4 == 0);
                let shader = shader_code.align_to::<u32>().1;
                wgpu::ShaderSource::SpirV(std::borrow::Cow::Borrowed(shader))
            },
        });
        let var_name = [Some(wgpu::ColorTargetState {
            format: surface.get_capabilities(&adapter).formats[0],
            blend: Some(wgpu::BlendState::REPLACE),
            write_mask: wgpu::ColorWrites::ALL,
        })];
        let pipeline_descriptor = wgpu::RenderPipelineDescriptor {
            label: Some("Main pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &vert_shader,
                entry_point: Some("main"),
                buffers: &[Vertex::desc()],
                compilation_options: Default::default(),
            },
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                strip_index_format: None,
                front_face: wgpu::FrontFace::Ccw,
                cull_mode: Some(wgpu::Face::Back),
                polygon_mode: wgpu::PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            fragment: Some(FragmentState {
                module: &frag_shader,
                entry_point: Some("main"),
                compilation_options: PipelineCompilationOptions::default(),
                targets: &var_name,
            }),
            multiview: None,
            cache: None,
        };
        // todo!();
        let surface_caps = surface.get_capabilities(&adapter);
        let present_mode = if surface_caps
            .present_modes
            .contains(&wgpu::PresentMode::Mailbox)
        {
            wgpu::PresentMode::Mailbox
        } else {
            wgpu::PresentMode::Fifo
        };

        let surface_config = wgpu::SurfaceConfiguration {
            width,
            height,
            format: surface_caps.formats[0],
            present_mode,
            alpha_mode: wgpu::CompositeAlphaMode::Auto,
            view_formats: vec![],
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
            desired_maximum_frame_latency: 3,
        };

        surface.configure(&device, &surface_config);
        let render_pipeline = device.create_render_pipeline(&pipeline_descriptor);
        WgpuRenderer {
            surface,
            surface_config,
            render_pipeline,
            device,
            queue,
        }
    }
}
static VERTICES: LazyLock<[Vertex; 3]> = std::sync::LazyLock::new(|| {
    [
        Vertex {
            position: glm::vec3(-0.6, -0.5, 0.0),
            color: glm::vec3(1.0, 0.0, 0.0),
        },
        Vertex {
            position: glm::vec3(0.6, -0.5, 0.0),
            color: glm::vec3(0.0, 1.0, 0.0),
        },
        Vertex {
            position: glm::vec3(0.0, 0.5, 0.0),
            color: glm::vec3(0.0, 0.0, 1.0),
        },
    ]
});

static INDICIES: [u32; 3] = [0, 1, 2];

impl<'window> ApplicationHandler for App<'window> {
    fn window_event(
        &mut self,
        event_loop: &winit::event_loop::ActiveEventLoop,
        window_id: WindowId,
        event: winit::event::WindowEvent,
    ) {
        match event {
            winit::event::WindowEvent::RedrawRequested => {
                let window_ctx = self.windows.get(&window_id).unwrap();
                let surface_texture = window_ctx.renderer.surface.get_current_texture().unwrap();
                let view = surface_texture.texture.create_view(&Default::default());
                let mut encoder = window_ctx.renderer.device.create_command_encoder(
                    &wgpu::wgt::CommandEncoderDescriptor {
                        label: Some("Render encoder"),
                    },
                );
                {
                    let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                        label: Some("Render"),
                        color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                            view: &view,
                            resolve_target: None,
                            ops: wgpu::Operations {
                                load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                                store: wgpu::StoreOp::Store,
                            },
                        })],
                        depth_stencil_attachment: None,
                        timestamp_writes: None,
                        occlusion_query_set: None,
                    });
                    rpass.set_pipeline(&window_ctx.renderer.render_pipeline);
                    let vertex_buffer = window_ctx.renderer.device.create_buffer_init(
                        &wgpu::util::BufferInitDescriptor {
                            label: Some("Vertex buffer"),
                            contents: unsafe { VERTICES.align_to::<u8>().1 },
                            usage: BufferUsages::VERTEX,
                        },
                    );
                    let index_buffer = window_ctx.renderer.device.create_buffer_init(
                        &wgpu::util::BufferInitDescriptor {
                            label: Some("Index buffer"),
                            contents: unsafe { INDICIES.align_to::<u8>().1 },
                            usage: BufferUsages::INDEX,
                        },
                    );
                    rpass.set_vertex_buffer(0, vertex_buffer.slice(..));
                    rpass.set_index_buffer(index_buffer.slice(..), IndexFormat::Uint32);
                    rpass.draw_indexed(0..INDICIES.len() as u32, 0, 0..1);
                }
                window_ctx.renderer.queue.submit(Some(encoder.finish()));
                surface_texture.present()
            }
            winit::event::WindowEvent::CloseRequested => {
                let _ = self.windows.remove(&window_id);
                if self.windows.is_empty() {
                    event_loop.exit();
                }
            }
            winit::event::WindowEvent::MouseInput {
                device_id,
                state,
                button,
            } => {
                if button == MouseButton::Left && state == ElementState::Pressed {
                    self.spawn_window(event_loop);
                }
            }
            winit::event::WindowEvent::Resized(size) => {
                if let Some(window_ctx) = self.windows.get(&window_id) {
                    if size.width == 0 || size.height == 0 {
                        return;
                    }
                    let mut new_config = window_ctx.renderer.surface_config.clone();
                    new_config.width = size.width;
                    new_config.height = size.height;
                    window_ctx
                        .renderer
                        .surface
                        .configure(&window_ctx.renderer.device, &new_config)
                }
            }
            _ => (),
        }
    }

    fn resumed(&mut self, event_loop: &winit::event_loop::ActiveEventLoop) {
        let attr = WindowAttributes::default()
            .with_title("Zenyx")
            .with_min_inner_size(winit::dpi::LogicalSize::new(1, 1));

        if self.windows.is_empty() {
            let window = event_loop.create_window(attr).unwrap();
            let window = Arc::new(window);
            let renderer = self.state.create_renderer(window.clone());
            let window_ctx = WindowContext {
                renderer: smol::block_on(renderer),
                window: window.clone(),
            };
            let window_id = window.id();
            self.windows.insert(window_id, window_ctx);
        }
    }
}

fn main() -> Result<(), terminator::Terminator> {
    let config = LoggerConfig::default()
        .colored_stdout(true)
        .log_to_stdout(true)
        .file_include_time(true)
        .log_to_file(true)
        .level(LogLevel::Info)
        .log_path("zenyx.log");
    let _logger = zlog::Logger::new(config);
    let event_loop = EventLoop::new()?;
    event_loop.set_control_flow(winit::event_loop::ControlFlow::Poll);

    let mut app = App::new();
    event_loop.run_app(&mut app)?;
    Ok(())
}