// Jake Jensen, 2025
// Rust

use anyhow::{Result, Context};
use std::sync::Arc;
use tokio::sync::watch;
use vulkano::{
    buffer::{Buffer, BufferCreateInfo, BufferUsage, CpuAccessibleBuffer, Subbuffer},
    command_buffer::{
        allocator::{StandardCommandBufferAllocator, StandardCommandBufferAllocatorCreateInfo},
        AutoCommandBufferBuilder, CommandBufferUsage, RenderPassBeginInfo, SubpassContents,
    },
    device::{
        physical::{PhysicalDevice, PhysicalDeviceType},
        Device, DeviceCreateInfo, DeviceExtensions, QueueCreateInfo, QueueFlags,
    },
    format::Format,
    image::{view::ImageView, Image, ImageCreateInfo, ImageUsage, SwapchainImage},
    instance::{Instance, InstanceCreateInfo, InstanceExtensions, Version},
    memory::allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
    pipeline::{
        graphics::{
            input_assembly::InputAssemblyState,
            vertex_input::{Vertex, VertexInputState},
            viewport::{Viewport, ViewportState},
            GraphicsPipeline,
        },
        Pipeline, PipelineBindPoint,
    },
    render_pass::{Framebuffer, FramebufferCreateInfo, RenderPass, Subpass},
    swapchain::{
        acquire_next_image, Surface, Swapchain, SwapchainCreateInfo, SwapchainPresentInfo,
    },
    sync::{self, GpuFuture},
    VulkanLibrary,
};
use winit::{
    event::{Event, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::{Window, WindowBuilder},
};
use nalgebra::{Vector2, Vector4};

// Include the shader code
mod vs {
    vulkano_shaders::shader! {
        ty: "vertex",
        path: "src/shaders/vert.glsl",
    }
}

mod fs {
    vulkano_shaders::shader! {
        ty: "fragment",
        path: "src/shaders/frag.glsl",
    }
}

/// Represents a single circle in the simulation.
#[derive(Debug, Clone, Copy)]
pub struct Circle {
    pub position: Vector2<f32>,
    pub radius: f32,
    pub color: Vector4<f32>,
}

/// The entire physics state to be sent to the renderer.
#[derive(Debug, Clone)]
pub struct PhysicsState {
    pub circles: Vec<Circle>,
}

// Implement the `Vertex` trait for our quad vertices.
// These are the vertices that make up a unit square, which will be scaled and positioned by the shader.
#[derive(BufferContents, Vertex)]
#[repr(C)]
struct QuadVertex {
    #[format(R32G32_SFLOAT)]
    position: [f32; 2],
}

// --- Physics Simulation Task ---
async fn physics_task(
    sender: watch::Sender<PhysicsState>,
    initial_state: PhysicsState,
) -> Result<()> {
    let mut physics_state = initial_state;
    let mut last_update_time = tokio::time::Instant::now();
    let mut accumulator = tokio::time::Duration::ZERO;

    // Fixed timestep for physics updates (e.g., 60 updates per second)
    const PHYSICS_TIMESTEP: tokio::time::Duration = tokio::time::Duration::from_nanos(16_666_667); // ~1/60th second

    log::info!("Physics task started.");

    loop {
        let now = tokio::time::Instant::now();
        accumulator += now.duration_since(last_update_time);
        last_update_time = now;

        // Prevent "spiral of death" if rendering is too slow or system is overloaded
        // By clamping the accumulator, we ensure physics doesn't try to catch up infinitely.
        // It will just slow down the simulation if the system can't keep up.
        if accumulator > PHYSICS_TIMESTEP * 5 { // Cap at 5 physics steps behind
            accumulator = PHYSICS_TIMESTEP * 5;
        }

        // Perform physics updates in fixed steps
        while accumulator >= PHYSICS_TIMESTEP {
            // --- YOUR PHYSICS LOGIC GOES HERE ---
            // For now, let's just make circles bounce off walls and move randomly
            for circle in &mut physics_state.circles {
                // Simple movement
                circle.position.x += 0.001; // Move right
                circle.position.y += 0.0005; // Move up

                // Simple wall bounce (normalized device coordinates -1 to 1)
                if circle.position.x + circle.radius > 1.0 {
                    circle.position.x = 1.0 - circle.radius;
                }
                if circle.position.x - circle.radius < -1.0 {
                    circle.position.x = -1.0 + circle.radius;
                }
                if circle.position.y + circle.radius > 1.0 {
                    circle.position.y = 1.0 - circle.radius;
                }
                if circle.position.y - circle.radius < -1.0 {
                    circle.position.y = -1.0 + circle.radius;
                }
            }
            // --- END PHYSICS LOGIC ---

            accumulator -= PHYSICS_TIMESTEP;
        }

        // Send the updated physics state to the renderer.
        // `send_replace` ensures only the latest value is available, dropping older ones.
        sender.send_replace(physics_state.clone());

        // Yield control back to the Tokio runtime.
        // This is important to allow other tasks (like the winit event loop if it's also async)
        // to run, and to prevent busy-waiting.
        tokio::time::sleep(tokio::time::Duration::from_millis(1)).await;
    }
}

// --- Vulkan Renderer ---
struct Renderer {
    _library: Arc<VulkanLibrary>,
    instance: Arc<Instance>,
    surface: Arc<Surface>,
    device: Arc<Device>,
    queue: Arc<vulkano::device::Queue>,
    memory_allocator: Arc<StandardMemoryAllocator>,
    command_buffer_allocator: Arc<StandardCommandBufferAllocator>,
    render_pass: Arc<RenderPass>,
    pipeline: Arc<GraphicsPipeline>,
    swapchain: Arc<Swapchain>,
    framebuffers: Vec<Arc<Framebuffer>>,
    recreate_swapchain: bool,
    previous_frame_end: Option<Box<dyn GpuFuture>>,
    quad_vertex_buffer: Subbuffer<[QuadVertex]>, // Static buffer for the quad
}

impl Renderer {
    fn new(event_loop: &EventLoop<()>) -> Result<Self> {
        let library = VulkanLibrary::new().context("No Vulkan library found")?;
        log::info!("Vulkan library loaded: {:?}", library.name());

        let required_extensions = vulkano_win::required_extensions(&library);
        let instance = Instance::new(
            library.clone(),
            InstanceCreateInfo {
                enabled_extensions: required_extensions,
                // Enable validation layers for debugging
                enabled_layers: if cfg!(debug_assertions) {
                    vec!["VK_LAYER_KHRONOS_validation".to_string()]
                } else {
                    vec![]
                },
                max_api_version: Some(Version::V1_1), // Specify minimum Vulkan API version
                ..Default::default()
            },
        )
        .context("Failed to create Vulkan instance")?;
        log::info!("Vulkan instance created.");

        let window = Arc::new(WindowBuilder::new().build(event_loop).context("Failed to create window")?);
        let surface = vulkano_win::create_surface_from_winit(instance.clone(), window.clone())
            .context("Failed to create Vulkan surface")?;
        log::info!("Vulkan surface created.");

        let (physical_device, queue_family_index) = Self::select_physical_device(&instance, &surface)?;
        log::info!("Selected physical device: {}", physical_device.properties().device_name);

        let (device, mut queues) = Device::new(
            physical_device.clone(),
            DeviceCreateInfo {
                enabled_extensions: DeviceExtensions {
                    khr_swapchain: true, // Required for rendering to a window
                    ..DeviceExtensions::empty()
                },
                queue_create_infos: vec![QueueCreateInfo {
                    queue_family_index,
                    // Prioritize this queue for scheduling
                    queue_priorities: vec![1.0],
                    ..Default::default()
                }],
                ..Default::default()
            },
        )
        .context("Failed to create Vulkan device")?;
        let queue = queues.next().context("No queues found")?;
        log::info!("Vulkan device and queue created.");

        let memory_allocator = Arc::new(StandardMemoryAllocator::new_default(device.clone()));
        let command_buffer_allocator = Arc::new(StandardCommandBufferAllocator::new(
            device.clone(),
            StandardCommandBufferAllocatorCreateInfo::default(),
        ));
        log::info!("Memory and command buffer allocators created.");

        let (mut swapchain, images) = Self::create_swapchain(
            device.clone(),
            surface.clone(),
            physical_device.clone(),
            window.inner_size(),
        )?;
        log::info!("Swapchain created.");

        let render_pass = Self::create_render_pass(device.clone(), swapchain.image_format())?;
        log::info!("Render pass created.");

        let pipeline = Self::create_pipeline(
            device.clone(),
            render_pass.clone(),
            swapchain.image_extent(),
        )?;
        log::info!("Graphics pipeline created.");

        let framebuffers = Self::create_framebuffers(&images, render_pass.clone())?;
        log::info!("Framebuffers created.");

        // Create a static vertex buffer for a unit quad
        let quad_vertices = [
            QuadVertex { position: [-1.0, -1.0] },
            QuadVertex { position: [ 1.0, -1.0] },
            QuadVertex { position: [-1.0,  1.0] },
            QuadVertex { position: [-1.0,  1.0] },
            QuadVertex { position: [ 1.0, -1.0] },
            QuadVertex { position: [ 1.0,  1.0] },
        ];
        let quad_vertex_buffer = Buffer::from_iter(
            &memory_allocator,
            BufferCreateInfo {
                usage: BufferUsage::VERTEX_BUFFER,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                ..Default::default()
            },
            quad_vertices,
        ).context("Failed to create quad vertex buffer")?;
        log::info!("Quad vertex buffer created.");

        let previous_frame_end = Some(sync::now(device.clone()).boxed());

        Ok(Self {
            _library: library,
            instance,
            surface,
            device,
            queue,
            memory_allocator,
            command_buffer_allocator,
            render_pass,
            pipeline,
            swapchain,
            framebuffers,
            recreate_swapchain: false,
            previous_frame_end,
            quad_vertex_buffer,
        })
    }

    fn select_physical_device(
        instance: &Arc<Instance>,
        surface: &Arc<Surface>,
    ) -> Result<(Arc<PhysicalDevice>, u32)> {
        instance
            .enumerate_physical_devices()
            .context("Failed to enumerate physical devices")?
            .filter(|p| p.supported_extensions().contains(&DeviceExtensions { khr_swapchain: true, ..DeviceExtensions::empty() }))
            .filter_map(|p| {
                p.queue_family_properties()
                    .iter()
                    .enumerate()
                    .position(|(i, q)| {
                        q.queue_flags.contains(QueueFlags::GRAPHICS)
                            && p.surface_support(i as u32, surface).unwrap_or(false)
                    })
                    .map(|i| (p, i as u32))
            })
            .min_by_key(|(p, _)| match p.properties().device_type {
                PhysicalDeviceType::DiscreteGpu => 0,
                PhysicalDeviceType::IntegratedGpu => 1,
                PhysicalDeviceType::VirtualGpu => 2,
                PhysicalDeviceType::Cpu => 3,
                _ => 4,
            })
            .context("No suitable physical device found")
    }

    fn create_swapchain(
        device: Arc<Device>,
        surface: Arc<Surface>,
        physical_device: Arc<PhysicalDevice>,
        window_size: winit::dpi::PhysicalSize<u32>,
    ) -> Result<(Arc<Swapchain>, Vec<Arc<Image>>)> {
        let capabilities = physical_device
            .surface_capabilities(&surface, Default::default())
            .context("Failed to get surface capabilities")?;

        let image_format = physical_device
            .surface_formats(&surface, Default::default())
            .context("Failed to get surface formats")?
            .iter()
            .min_by_key(|(f, _)| match f {
                // Prefer sRGB if available, otherwise just pick the first one
                Format::R8G8B8A8_SRGB => 0,
                Format::B8G8R8A8_SRGB => 0,
                _ => 1,
            })
            .map(|(f, _)| *f)
            .unwrap_or(Format::B8G8R8A8_SRGB); // Fallback

        let present_mode = capabilities.present_modes.iter().next().unwrap(); // Just pick the first available
        let image_extent = capabilities.current_extent.unwrap_or([
            window_size.width,
            window_size.height,
        ]);

        Swapchain::new(
            device.clone(),
            surface.clone(),
            SwapchainCreateInfo {
                min_image_count: capabilities.min_image_count,
                image_format,
                image_extent,
                image_array_layers: 1,
                image_usage: ImageUsage::COLOR_ATTACHMENT,
                composite_alpha: capabilities.supported_composite_alpha.iter().next().unwrap(),
                present_mode,
                clipped: true,
                ..Default::default()
            },
        )
        .context("Failed to create swapchain")
    }

    fn create_render_pass(device: Arc<Device>, format: Format) -> Result<Arc<RenderPass>> {
        vulkano::single_pass_renderpass! {
            device.clone(),
            attachments: {
                color: {
                    load: Clear,
                    store: Store,
                    format: format,
                    samples: 1,
                }
            },
            pass: {
                color: [color],
                depth_stencil: {}
            }
        }
        .context("Failed to create render pass")
    }

    fn create_pipeline(
        device: Arc<Device>,
        render_pass: Arc<RenderPass>,
        image_extent: [u32; 2],
    ) -> Result<Arc<GraphicsPipeline>> {
        let vs = vs::load(device.clone()).context("Failed to load vertex shader")?;
        let fs = fs::load(device.clone()).context("Failed to load fragment shader")?;

        GraphicsPipeline::new(
            device.clone(),
            Some(vs.entry_point("main").unwrap()),
            Some(fs.entry_point("main").unwrap()),
            InputAssemblyState::new(),
            ViewportState::viewport_fixed_scissor_irrelevant([Viewport {
                offset: [0.0, 0.0],
                extent: [image_extent[0] as f32, image_extent[1] as f32],
                depth_range: 0.0..=1.0,
            }]),
            render_pass.clone().into(),
            None,
        )
        .context("Failed to create graphics pipeline")
    }

    fn create_framebuffers(
        images: &[Arc<SwapchainImage>],
        render_pass: Arc<RenderPass>,
    ) -> Result<Vec<Arc<Framebuffer>>> {
        images
            .iter()
            .map(|image| {
                let view = ImageView::new_default(image.clone()).context("Failed to create image view")?;
                Framebuffer::new(
                    render_pass.clone(),
                    FramebufferCreateInfo {
                        attachments: vec![view],
                        ..Default::default()
                    },
                )
                .context("Failed to create framebuffer")
            })
            .collect::<Result<Vec<_>>>()
    }

    fn recreate_swapchain_if_needed(&mut self, new_dimensions: [u32; 2]) -> Result<()> {
        if self.recreate_swapchain {
            self.recreate_swapchain = false;

            let (new_swapchain, new_images) = Self::create_swapchain(
                self.device.clone(),
                self.surface.clone(),
                self.device.physical_device().clone(),
                winit::dpi::PhysicalSize::new(new_dimensions[0], new_dimensions[1]),
            )?;

            self.swapchain = new_swapchain;
            self.framebuffers = Self::create_framebuffers(&new_images, self.render_pass.clone())?;
            self.pipeline = Self::create_pipeline(
                self.device.clone(),
                self.render_pass.clone(),
                new_dimensions,
            )?;
            log::info!("Swapchain recreated with new dimensions: {:?}", new_dimensions);
        }
        Ok(())
    }

    fn render(&mut self, physics_state: &PhysicsState) -> Result<()> {
        self.previous_frame_end.as_mut().unwrap().cleanup_finished();

        let image_extent = self.swapchain.image_extent();
        if image_extent[0] == 0 || image_extent[1] == 0 {
            // Window is minimized, skip rendering
            return Ok(());
        }

        self.recreate_swapchain_if_needed(image_extent)?;

        let (image_index, suboptimal, acquire_future) =
            match acquire_next_image(self.swapchain.clone(), None) {
                Ok(r) => r,
                Err(vulkano::sync::AcquireError::OutOfDate) => {
                    self.recreate_swapchain = true;
                    return Ok(());
                }
                Err(e) => return Err(anyhow::Error::from(e).context("Failed to acquire next swapchain image")),
            };

        if suboptimal {
            self.recreate_swapchain = true;
        }

        let mut builder = AutoCommandBufferBuilder::primary(
            &self.command_buffer_allocator,
            self.queue.queue_family_index(),
            CommandBufferUsage::OneTimeSubmit, // Command buffer will be submitted once
        ).context("Failed to create command buffer builder")?;

        builder
            .begin_render_pass(
                RenderPassBeginInfo {
                    clear_values: vec![Some([0.0, 0.0, 0.1, 1.0].into())], // Clear to dark blue
                    ..RenderPassBeginInfo::framebuffer(self.framebuffers[image_index as usize].clone())
                },
                SubpassContents::Inline,
            )?
            .bind_pipeline_graphics(self.pipeline.clone())?;

        // Draw each circle
        for circle in &physics_state.circles {
            let push_constants = fs::PushConstants {
                center: circle.position.into(),
                radius: circle.radius,
                color: circle.color.into(),
            };
            builder
                .push_constants(self.pipeline.layout().clone(), 0, push_constants)?
                .bind_vertex_buffers(0, self.quad_vertex_buffer.clone())?
                .draw(self.quad_vertex_buffer.len() as u32, 1, 0, 0)?;
        }

        builder.end_render_pass()?;

        let command_buffer = builder.build().context("Failed to build command buffer")?;

        let future = self
            .previous_frame_end
            .take()
            .unwrap()
            .join(acquire_future)
            .then_execute(self.queue.clone(), command_buffer)
            .context("Failed to execute command buffer")?
            .then_swapchain_present(self.queue.clone(), self.swapchain.clone(), image_index)
            .then_signal_fence_and_flush();

        match future {
            Ok(future) => {
                self.previous_frame_end = Some(future.boxed());
            }
            Err(vulkano::sync::FlushError::OutOfDate) => {
                self.recreate_swapchain = true;
                self.previous_frame_end = Some(sync::now(self.device.clone()).boxed());
            }
            Err(e) => {
                log::error!("Failed to flush future: {:?}", e);
                self.previous_frame_end = Some(sync::now(self.device.clone()).boxed());
            }
        }

        Ok(())
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging (useful for Vulkano debug messages)
    env_logger::init();

    // --- Setup communication channels ---
    // `watch` channel is great for sending the latest state, dropping older ones if not read.
    let initial_physics_state = PhysicsState {
        circles: vec![
            Circle {
                position: Vector2::new(0.0, 0.0),
                radius: 0.1,
                color: Vector4::new(1.0, 0.0, 0.0, 1.0), // Red
            },
            Circle {
                position: Vector2::new(0.5, 0.5),
                radius: 0.05,
                color: Vector4::new(0.0, 1.0, 0.0, 1.0), // Green
            },
            Circle {
                position: Vector2::new(-0.5, -0.5),
                radius: 0.08,
                color: Vector4::new(0.0, 0.0, 1.0, 1.0), // Blue
            },
        ],
    };
    let (physics_sender, physics_receiver) = watch::channel(initial_physics_state.clone());

    // --- Spawn Physics Task ---
    tokio::spawn(physics_task(physics_sender, initial_physics_state));

    // --- Setup Winit Event Loop and Renderer ---
    let event_loop = EventLoop::new().context("Failed to create event loop")?;
    let mut renderer = Renderer::new(&event_loop)?;

    // --- Main Rendering Loop ---
    event_loop.run(move |event, _, control_flow| {
        *control_flow = ControlFlow::Poll; // Continuously poll for events and redraw

        match event {
            Event::WindowEvent {
                event: WindowEvent::CloseRequested,
                ..
            } => {
                *control_flow = ControlFlow::Exit;
            }
            Event::WindowEvent {
                event: WindowEvent::Resized(_),
                ..
            } => {
                renderer.recreate_swapchain = true;
            }
            Event::MainEventsCleared => {
                // This event fires when all other events have been processed.
                // It's a good place to trigger a redraw.

                // Get the latest physics state from the watch channel
                let current_physics_state = physics_receiver.borrow().clone();

                // Render the current state
                if let Err(e) = renderer.render(&current_physics_state) {
                    log::error!("Rendering error: {:?}", e);
                    // Handle critical rendering errors, e.g., exit application
                    *control_flow = ControlFlow::Exit;
                }
            }
            _ => (),
        }
    })?;

    Ok(())
}