use egui_winit_vulkano::{Gui, GuiConfig};
use render::render_thread;
use setup::{get_compute_pipeline, get_framebuffers};
use std::sync::Arc;
use vulkano::{
    buffer::{Buffer, BufferCreateInfo, BufferUsage},
    command_buffer::{allocator::StandardCommandBufferAllocator, AutoCommandBufferBuilder},
    descriptor_set::{
        allocator::StandardDescriptorSetAllocator, PersistentDescriptorSet, WriteDescriptorSet,
    },
    memory::allocator::{
        AllocationCreateInfo, MemoryAllocator, MemoryTypeFilter, StandardMemoryAllocator,
    },
    pipeline::Pipeline,
    swapchain::{self, SwapchainCreateInfo, SwapchainPresentInfo},
    sync::{self, future::FenceSignalFuture, GpuFuture},
    Validated, VulkanError,
};
use winit::{
    event::{Event, WindowEvent},
    event_loop::ControlFlow,
};

mod render;
mod setup;

mod cs {
    vulkano_shaders::shader! {
        ty: "compute",
        src: "
            #version 460
            layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
            layout(set = 0, binding = 0) uniform writeonly image2D outputImage;

            layout(set = 0, binding = 1) uniform GPUSettings {
                uint samples;
            } settings;

            layout(set = 0, binding = 2) buffer Data {
                float data[];
            } buf;


            void main() {
                uvec2 size = imageSize(outputImage);
                uvec2 pos = gl_GlobalInvocationID.xy;
                if (pos.x < size.x && pos.y < size.y) {
                    uint index = (pos.y * size.x + pos.x) * 3;
                    vec3 color = vec3(buf.data[index], buf.data[index + 1], buf.data[index + 2]);

                    vec4 to_write = vec4(color / settings.samples, 1.0);
                    imageStore(outputImage, ivec2(pos), to_write);
                }
            }
        "
    }
}

fn main() {
    let (event_loop, physical_device, device, queue, window, surface) = setup::setup();

    let (mut swapchain, mut images) = setup::create_swapchain(
        physical_device.clone(),
        device.clone(),
        surface.clone(),
        window.inner_size(),
    );

    let mut gui = Gui::new(
        &event_loop,
        surface.clone(),
        queue.clone(),
        swapchain.image_format(),
        GuiConfig {
            is_overlay: true,
            ..Default::default()
        },
    );

    let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));

    let command_buffer_allocator =
        StandardCommandBufferAllocator::new(device.clone(), Default::default());
    let descriptor_set_allocator =
        StandardDescriptorSetAllocator::new(device.clone(), Default::default());

    let (control_tx, control_rx) = std::sync::mpsc::channel();
    let (data_tx, data_rx) = std::sync::mpsc::channel();

    let mut settings = render::RenderSettings {
        width: window.inner_size().width,
        height: window.inner_size().height,
    };

    let send_settings = settings.clone();
    let send_memory_allocator = memory_allocator.clone();

    std::thread::spawn(move || {
        render_thread(send_settings, control_rx, data_tx, send_memory_allocator)
    });

    let render_pass = setup::get_render_pass(device.clone(), &swapchain);
    let mut framebuffers = setup::get_framebuffers(&images, &render_pass);

    let mut recreate_swapchain = false;

    let mut fences: Vec<Option<Arc<FenceSignalFuture<_>>>> = vec![None; images.len()];
    let mut previous_fence_i = 0;

    let mut buffer = None;

    let cs = cs::load(device.clone()).unwrap();

    let cpu_buffers: Vec<_> = framebuffers
        .iter()
        .map(|_| {
            Buffer::from_data(
                memory_allocator.clone(),
                BufferCreateInfo {
                    usage: BufferUsage::UNIFORM_BUFFER,
                    ..Default::default()
                },
                AllocationCreateInfo {
                    memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
                        | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                    ..Default::default()
                },
                cs::GPUSettings { samples: 1 },
            )
            .unwrap()
        })
        .collect();

    let mut pipeline = get_compute_pipeline(device.clone(), cs.clone());

    event_loop.run(move |event, _, control_flow| match event {
        Event::WindowEvent { window_id, event } => {
            gui.update(&event);
            match event {
                WindowEvent::Resized(_) => {
                    recreate_swapchain = true;
                    window.request_redraw();
                    settings.width = window.inner_size().width;
                    settings.height = window.inner_size().height;
                    let _ = control_tx.send(settings.clone());
                }
                WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                _ => (),
            }
        }
        Event::MainEventsCleared => {
            if recreate_swapchain {
                let (new_swapchain, new_images) = swapchain
                    .recreate(SwapchainCreateInfo {
                        image_extent: window.inner_size().into(),
                        ..swapchain.create_info()
                    })
                    .expect("failed to recreate swapchain: {e}");
                swapchain = new_swapchain;
                framebuffers = get_framebuffers(&new_images, &render_pass);
                pipeline = get_compute_pipeline(device.clone(), cs.clone());
            }

            gui.immediate_ui(|gui| {
                let ctx = gui.context();
                egui::Window::new("panel").show(&ctx, |ui| {
                    ui.heading("fdjkslf");

                    ui.collapsing("Render Settings:", |ui| {
                        ui.label(format!("{:#?}", &settings));
                    });
                });
            });

            let (image_i, suboptimal, acquire_future) =
                match swapchain::acquire_next_image(swapchain.clone(), None)
                    .map_err(Validated::unwrap)
                {
                    Ok(r) => r,
                    Err(e) => panic!("failed to acquire next image: {e}"),
                };

            let previous_future = match fences[previous_fence_i as usize].clone() {
                // Create a NowFuture
                None => {
                    let mut now = sync::now(device.clone());
                    now.cleanup_finished();

                    now.boxed()
                }
                // Use the existing FenceSignalFuture
                Some(fence) => fence.boxed(),
            };

            if suboptimal {
                recreate_swapchain = true;
            }

            if let Some(image_fence) = &fences[image_i as usize] {
                image_fence.wait(None).unwrap();
            }

            while let Ok(b) = data_rx.try_recv() {
                buffer = Some(b);
            }

            if buffer.as_ref().is_some_and(|b| {
                b.width != window.inner_size().width || b.height != window.inner_size().height
            }) {
                buffer = None;
            }

            let compute_future = match &buffer {
                Some(b) => {
                    match cpu_buffers[image_i as usize].write() {
                        Ok(mut local) => {
                            local.samples = b.samples;
                        }
                        Err(e) => {
                            println!("{}", e);
                        }
                    }

                    let mut builder = AutoCommandBufferBuilder::primary(
                        &command_buffer_allocator,
                        queue.queue_family_index(),
                        vulkano::command_buffer::CommandBufferUsage::OneTimeSubmit,
                    )
                    .unwrap();

                    let pipeline_layout = pipeline.layout();
                    let descriptor_set_layouts = pipeline_layout.set_layouts();

                    let descriptor_set_layout = descriptor_set_layouts.first().unwrap();

                    let descriptor_set = PersistentDescriptorSet::new(
                        &descriptor_set_allocator,
                        descriptor_set_layout.clone(),
                        [
                            WriteDescriptorSet::image_view(
                                0,
                                framebuffers[image_i as usize].attachments()[0].clone(),
                            ),
                            WriteDescriptorSet::buffer(1, cpu_buffers[image_i as usize].clone()),
                            WriteDescriptorSet::buffer(2, b.buffer.clone()),
                        ],
                        [],
                    )
                    .unwrap();

                    builder
                        .bind_pipeline_compute(pipeline.clone())
                        .unwrap()
                        .bind_descriptor_sets(
                            vulkano::pipeline::PipelineBindPoint::Compute,
                            pipeline_layout.clone(),
                            0,
                            descriptor_set,
                        )
                        .unwrap()
                        .dispatch([b.width.div_ceil(8), b.height.div_ceil(8), 1])
                        .unwrap();

                    let command = builder.build().unwrap();

                    acquire_future
                        .join(previous_future)
                        .then_execute(queue.clone(), command)
                        .unwrap()
                        .then_signal_fence()
                        .boxed()
                }
                None => acquire_future.join(previous_future).boxed(),
            };

            let gui_future = gui.draw_on_image(
                compute_future,
                framebuffers[image_i as usize].attachments()[0].clone(),
            );

            let future = gui_future
                .then_swapchain_present(
                    queue.clone(),
                    SwapchainPresentInfo::swapchain_image_index(swapchain.clone(), image_i),
                )
                .then_signal_fence_and_flush();

            fences[image_i as usize] = match future.map_err(Validated::unwrap) {
                Ok(value) => Some(Arc::new(value)),
                Err(VulkanError::OutOfDate) => {
                    recreate_swapchain = true;
                    None
                }
                Err(e) => {
                    println!("failed to flush future: {e}");
                    None
                }
            };

            previous_fence_i = image_i;
        }
        _ => (),
    })
}