// SPDX-License-Identifier: AGPL-3.0-only

use crate::create_gpu_buffer;
use crate::sync::SyncPrimitives;
use crate::{error::RendererError, mesh::Vertex};
use ash::{Device, Instance, khr, vk};
use gpu_allocator::vulkan::{Allocation, Allocator};

/// The core renderer structure holding the Vulkan resources.
pub struct Renderer {
    /// Entry point to the Vulkan library.
    pub(crate) _entry: ash::Entry,
    /// The Vulkan instance.
    pub(crate) instance: Instance,
    /// Optional debug utility loader for validation layers.
    pub(crate) debug_utils: Option<ash::ext::debug_utils::Instance>,
    /// The debug messenger for validation layer output.
    pub(crate) debug_messenger: vk::DebugUtilsMessengerEXT,
    /// Handle to the selected physical device (GPU).
    #[expect(dead_code)]
    pub(crate) physical_device: vk::PhysicalDevice,
    /// The logical Vulkan device.
    pub(crate) device: Device,
    /// The queue used for graphics operations.
    pub(crate) graphics_queue: vk::Queue,
    /// Index of the graphics queue family.
    #[expect(dead_code)]
    pub(crate) graphics_queue_index: u32,
    /// Surface extension loader.
    pub(crate) surface_loader: khr::surface::Instance,
    /// The presentation surface.
    pub(crate) surface: vk::SurfaceKHR,
    /// Swapchain extension loader.
    pub(crate) swapchain_loader: khr::swapchain::Device,
    /// The swapchain for presenting images.
    pub(crate) swapchain: vk::SwapchainKHR,
    /// Images acquired from the swapchain.
    pub(crate) swapchain_images: Vec<vk::Image>,
    /// The pixel format of the swapchain images.
    #[expect(dead_code)]
    pub(crate) swapchain_format: vk::Format,
    /// The dimensions of the swapchain images.
    pub(crate) swapchain_extent: vk::Extent2D,
    /// Image views for each swapchain image.
    pub(crate) swapchain_image_views: Vec<vk::ImageView>,
    /// The command pool used for allocating command buffers.
    pub(crate) command_pool: vk::CommandPool,
    /// Pre-allocated command buffers for each frame in flight.
    pub(crate) command_buffers: Vec<vk::CommandBuffer>,
    /// The layout of the graphics pipeline.
    pub(crate) pipeline_layout: vk::PipelineLayout,
    /// The compiled graphics pipeline state.
    pub(crate) graphics_pipeline: vk::Pipeline,
    /// Memory manager for GPU allocations.
    pub(crate) allocator: Option<Allocator>,
    /// Buffer containing the vertex data for the initial triangle.
    pub(crate) vertex_buffer: vk::Buffer,
    /// Memory allocation for the vertex buffer.
    pub(crate) vertex_allocation: Option<Allocation>,
    /// Buffer containing the index data for indexed drawing.
    pub(crate) index_buffer: vk::Buffer,
    /// Memory allocation for the index buffer.
    pub(crate) index_allocation: Option<Allocation>,
    pub(crate) index_count: u32,
    /// The depth image used for depth testing.
    pub(crate) depth_image: vk::Image,
    /// Image view for the depth buffer.
    pub(crate) depth_image_view: vk::ImageView,
    /// Memory allocation for the depth image.
    pub(crate) depth_allocation: Option<Allocation>,
    /// Synchronization primitives for frame-by-frame execution.
    pub(crate) sync: Option<SyncPrimitives>,
    /// Index of the current frame being processed (0 to `crate::MAX_FRAMES_IN_FLIGHT` - 1).
    pub(crate) current_frame: usize,
}

impl Renderer {
    /// Renders a single frame.
    pub fn draw_frame(&mut self, camera_view: glam::Mat4) -> Result<(), RendererError> {
        let sync = self
            .sync
            .as_ref()
            .ok_or(RendererError::SyncPrimitivesMissing)?;
        let in_flight_fence = sync.in_flight[self.current_frame];
        let image_available_semaphore = sync.image_available[self.current_frame];
        let cmd = self.command_buffers[self.current_frame];

        // 1. Wait for the current frame's GPU work to finish
        unsafe {
            self.device
                .wait_for_fences(&[in_flight_fence], true, u64::MAX)?;
            self.device.reset_fences(&[in_flight_fence])?;
        }

        // 2. Acquire an image from the swapchain
        let (image_index, _is_suboptimal) = unsafe {
            self.swapchain_loader.acquire_next_image(
                self.swapchain,
                u64::MAX,
                image_available_semaphore,
                vk::Fence::null(),
            )?
        };

        // Use the semaphore tied to this specific swapchain image for rendering completion
        let render_finished_semaphore = self
            .sync
            .as_ref()
            .ok_or(RendererError::SyncPrimitivesMissing)?
            .render_finished[image_index as usize];

        // 3. Reset and begin recording the command buffer
        unsafe {
            self.device
                .reset_command_buffer(cmd, vk::CommandBufferResetFlags::empty())?;
            let begin_info = vk::CommandBufferBeginInfo::default()
                .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT);
            self.device.begin_command_buffer(cmd, &begin_info)?;
        }

        let image = self.swapchain_images[image_index as usize];
        let view = self.swapchain_image_views[image_index as usize];

        // 4. Record the actual rendering commands
        self.record_commands(cmd, view, image, camera_view)?;

        // 5. Submit the work to the GPU
        let submit_info = vk::SubmitInfo::default()
            .wait_semaphores(std::slice::from_ref(&image_available_semaphore))
            .wait_dst_stage_mask(&[vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT])
            .command_buffers(std::slice::from_ref(&cmd))
            .signal_semaphores(std::slice::from_ref(&render_finished_semaphore));

        unsafe {
            self.device
                .queue_submit(self.graphics_queue, &[submit_info], in_flight_fence)?;
        }

        // 6. Present the result to the screen
        let present_info = vk::PresentInfoKHR::default()
            .wait_semaphores(std::slice::from_ref(&render_finished_semaphore))
            .swapchains(std::slice::from_ref(&self.swapchain))
            .image_indices(std::slice::from_ref(&image_index));

        unsafe {
            self.swapchain_loader
                .queue_present(self.graphics_queue, &present_info)?;
        }

        // Advance the frame index for the next call
        self.current_frame = (self.current_frame + 1) % crate::MAX_FRAMES_IN_FLIGHT;

        Ok(())
    }

    /// Records the drawing commands into the given command buffer.
    fn record_commands(
        &self,
        cmd: vk::CommandBuffer,
        view: vk::ImageView,
        image: vk::Image,
        camera_view: glam::Mat4,
    ) -> Result<(), RendererError> {
        // Transition layouts for drawing
        self.transition_to_draw_layout(cmd, image);

        // Begin rendering
        let color_attachment = vk::RenderingAttachmentInfo::default()
            .image_view(view)
            .image_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
            .load_op(vk::AttachmentLoadOp::CLEAR)
            .store_op(vk::AttachmentStoreOp::STORE)
            .clear_value(vk::ClearValue {
                color: vk::ClearColorValue {
                    float32: [0.1, 0.2, 0.4, 1.0],
                },
            });

        let depth_attachment = vk::RenderingAttachmentInfo::default()
            .image_view(self.depth_image_view)
            .image_layout(vk::ImageLayout::DEPTH_ATTACHMENT_OPTIMAL)
            .load_op(vk::AttachmentLoadOp::CLEAR)
            .store_op(vk::AttachmentStoreOp::STORE)
            .clear_value(vk::ClearValue {
                depth_stencil: vk::ClearDepthStencilValue {
                    depth: 1.0,
                    stencil: 0,
                },
            });

        let rendering_info = vk::RenderingInfo::default()
            .render_area(vk::Rect2D {
                offset: vk::Offset2D { x: 0, y: 0 },
                extent: self.swapchain_extent,
            })
            .layer_count(1)
            .color_attachments(std::slice::from_ref(&color_attachment))
            .depth_attachment(&depth_attachment);

        unsafe {
            self.device.cmd_begin_rendering(cmd, &rendering_info);
            self.issue_draw_calls(cmd, camera_view);
            self.device.cmd_end_rendering(cmd);
        }

        // Transition back to present
        self.transition_to_present_layout(cmd, image)?;

        Ok(())
    }

    /// Transitions the swapchain and depth images to layouts suitable for drawing.
    fn transition_to_draw_layout(&self, cmd: vk::CommandBuffer, image: vk::Image) {
        let range = vk::ImageSubresourceRange {
            aspect_mask: vk::ImageAspectFlags::COLOR,
            base_mip_level: 0,
            level_count: 1,
            base_array_layer: 0,
            layer_count: 1,
        };

        let barrier_to_draw = vk::ImageMemoryBarrier2::default()
            .image(image)
            .subresource_range(range)
            .src_stage_mask(vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT)
            .src_access_mask(vk::AccessFlags2::empty())
            .dst_stage_mask(vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT)
            .dst_access_mask(vk::AccessFlags2::COLOR_ATTACHMENT_WRITE)
            .old_layout(vk::ImageLayout::UNDEFINED)
            .new_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

        let depth_range = vk::ImageSubresourceRange {
            aspect_mask: vk::ImageAspectFlags::DEPTH,
            base_mip_level: 0,
            level_count: 1,
            base_array_layer: 0,
            layer_count: 1,
        };

        let depth_barrier = vk::ImageMemoryBarrier2::default()
            .image(self.depth_image)
            .subresource_range(depth_range)
            .src_stage_mask(vk::PipelineStageFlags2::EARLY_FRAGMENT_TESTS)
            .src_access_mask(vk::AccessFlags2::empty())
            .dst_stage_mask(vk::PipelineStageFlags2::EARLY_FRAGMENT_TESTS)
            .dst_access_mask(vk::AccessFlags2::DEPTH_STENCIL_ATTACHMENT_WRITE)
            .old_layout(vk::ImageLayout::UNDEFINED)
            .new_layout(vk::ImageLayout::DEPTH_ATTACHMENT_OPTIMAL);

        let barriers = [barrier_to_draw, depth_barrier];
        let dependency_info = vk::DependencyInfo::default().image_memory_barriers(&barriers);

        unsafe { self.device.cmd_pipeline_barrier2(cmd, &dependency_info) };
    }

    /// Issues the actual draw calls for the frame.
    fn issue_draw_calls(&self, cmd: vk::CommandBuffer, camera_view: glam::Mat4) {
        unsafe {
            self.device.cmd_bind_pipeline(
                cmd,
                vk::PipelineBindPoint::GRAPHICS,
                self.graphics_pipeline,
            );

            #[expect(clippy::cast_precision_loss)]
            let viewport = vk::Viewport {
                x: 0.0,
                y: 0.0,
                width: self.swapchain_extent.width as f32,
                height: self.swapchain_extent.height as f32,
                min_depth: 0.0,
                max_depth: 1.0,
            };
            self.device.cmd_set_viewport(cmd, 0, &[viewport]);

            let scissor = vk::Rect2D {
                offset: vk::Offset2D { x: 0, y: 0 },
                extent: self.swapchain_extent,
            };
            self.device.cmd_set_scissor(cmd, 0, &[scissor]);

            self.device
                .cmd_bind_vertex_buffers(cmd, 0, &[self.vertex_buffer], &[0]);
            self.device
                .cmd_bind_index_buffer(cmd, self.index_buffer, 0, vk::IndexType::UINT32);

            let aspect =
                f64::from(self.swapchain_extent.width) / f64::from(self.swapchain_extent.height);

            #[expect(clippy::cast_possible_truncation)]
            let projection = glam::camera::rh::proj::vulkan::perspective(
                45.0_f32.to_radians(),
                aspect as f32,
                0.1,
                500.0,
            );

            // The view matrix is supplied by the caller (the client's camera); the renderer owns only the projection, which depends on the swapchain aspect ratio it manages.
            let mvp = projection * camera_view;

            let mvp_bytes = bytemuck::cast_slice(mvp.as_ref());
            self.device.cmd_push_constants(
                cmd,
                self.pipeline_layout,
                vk::ShaderStageFlags::VERTEX,
                0,
                mvp_bytes,
            );

            self.device
                .cmd_draw_indexed(cmd, self.index_count, 1, 0, 0, 0);
        }
    }

    /// Transitions the swapchain image back to the presentation layout.
    fn transition_to_present_layout(
        &self,
        cmd: vk::CommandBuffer,
        image: vk::Image,
    ) -> Result<(), RendererError> {
        let range = vk::ImageSubresourceRange {
            aspect_mask: vk::ImageAspectFlags::COLOR,
            base_mip_level: 0,
            level_count: 1,
            base_array_layer: 0,
            layer_count: 1,
        };

        let barrier_to_present = vk::ImageMemoryBarrier2::default()
            .image(image)
            .subresource_range(range)
            .src_stage_mask(vk::PipelineStageFlags2::COLOR_ATTACHMENT_OUTPUT)
            .src_access_mask(vk::AccessFlags2::COLOR_ATTACHMENT_WRITE)
            .dst_stage_mask(vk::PipelineStageFlags2::BOTTOM_OF_PIPE)
            .dst_access_mask(vk::AccessFlags2::empty())
            .old_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)
            .new_layout(vk::ImageLayout::PRESENT_SRC_KHR);

        let dependency_info = vk::DependencyInfo::default()
            .image_memory_barriers(std::slice::from_ref(&barrier_to_present));

        unsafe {
            self.device.cmd_pipeline_barrier2(cmd, &dependency_info);
            self.device.end_command_buffer(cmd)?;
        }

        Ok(())
    }

    /// Replaces the currently rendering mesh with a new set of vertices and indices.
    ///
    /// # Errors
    /// Returns a `RendererError` if new Vulkan buffers cannot be allocated or created.
    #[expect(clippy::cast_possible_truncation)]
    pub fn update_mesh(
        &mut self,
        vertices: &[Vertex],
        indices: &[u32],
    ) -> Result<(), RendererError> {
        unsafe {
            let _ = self.device.device_wait_idle();

            let allocator = self
                .allocator
                .as_mut()
                .ok_or(RendererError::AllocatorMissing)?;

            if let Some(alloc) = self.vertex_allocation.take() {
                let _ = allocator.free(alloc);
            }
            self.device.destroy_buffer(self.vertex_buffer, None);
            self.vertex_buffer = vk::Buffer::null();

            if let Some(alloc) = self.index_allocation.take() {
                let _ = allocator.free(alloc);
            }
            self.device.destroy_buffer(self.index_buffer, None);
            self.index_buffer = vk::Buffer::null();
        }

        let allocator = self
            .allocator
            .as_mut()
            .ok_or(RendererError::AllocatorMissing)?;

        let (v_buf, v_alloc) = create_gpu_buffer(
            &self.device,
            allocator,
            bytemuck::cast_slice(vertices),
            vk::BufferUsageFlags::VERTEX_BUFFER,
            "Chunk Vertex Buffer",
        )?;

        let (i_buf, i_alloc) = crate::create_gpu_buffer(
            &self.device,
            allocator,
            bytemuck::cast_slice(indices),
            vk::BufferUsageFlags::INDEX_BUFFER,
            "Chunk Index Buffer",
        )?;

        self.vertex_buffer = v_buf;
        self.vertex_allocation = Some(v_alloc);
        self.index_buffer = i_buf;
        self.index_allocation = Some(i_alloc);
        self.index_count = indices.len() as u32;

        Ok(())
    }
}

impl Drop for Renderer {
    fn drop(&mut self) {
        unsafe {
            let _ = self.device.device_wait_idle();

            self.device.destroy_pipeline(self.graphics_pipeline, None);
            self.device
                .destroy_pipeline_layout(self.pipeline_layout, None);

            if let Some(allocator) = self.allocator.as_mut() {
                if let Some(alloc) = self.vertex_allocation.take()
                    && let Err(e) = allocator.free(alloc)
                {
                    tracing::error!("Failed to free vertex buffer allocation: {e}");
                }
                if let Some(alloc) = self.index_allocation.take()
                    && let Err(e) = allocator.free(alloc)
                {
                    tracing::error!("Failed to free index buffer allocation: {e}");
                }
                if let Some(alloc) = self.depth_allocation.take()
                    && let Err(e) = allocator.free(alloc)
                {
                    tracing::error!("Failed to free depth image allocation: {e}");
                }
            }
            self.device.destroy_buffer(self.vertex_buffer, None);
            self.device.destroy_buffer(self.index_buffer, None);
            self.device.destroy_image_view(self.depth_image_view, None);
            self.device.destroy_image(self.depth_image, None);

            // Drop the allocator before destroying the logical device so its remaining memory blocks are released while the device is still valid.
            drop(self.allocator.take());

            self.device.destroy_command_pool(self.command_pool, None);

            if let Some(sync) = self.sync.take() {
                crate::sync::destroy_sync_primitives(&self.device, sync);
            }

            // Destroy image views before the swapchain that owns the underlying images.
            for &view in &self.swapchain_image_views {
                self.device.destroy_image_view(view, None);
            }

            self.swapchain_loader
                .destroy_swapchain(self.swapchain, None);

            // Destroy the logical device
            self.device.destroy_device(None);

            // Destroy the surface
            self.surface_loader.destroy_surface(self.surface, None);

            // Destroy the debug messenger if it exists
            if let Some(debug_utils) = self.debug_utils.as_ref() {
                debug_utils.destroy_debug_utils_messenger(self.debug_messenger, None);
            }

            // Destroy the instance
            self.instance.destroy_instance(None);
        }
    }
}