use crate::vulkan::{
    VkCommandPool, VkDevice, VkFramebuffer, VkGraphicsPipeline, VkInstance, VkPhysicalDevice,
    VkRenderPass, VkSemaphore, VkSurface, VkSwapchain,
};
use ash::vk;
use std::sync::Arc;

pub struct VkRenderContext {
    instance: Arc<VkInstance>,
    surface: Arc<VkSurface>,
    device: Arc<VkDevice>,

    swapchain: Arc<VkSwapchain>,
    render_pass: Arc<VkRenderPass>,
    framebuffers: Vec<Arc<VkFramebuffer>>,
    pipeline: Arc<VkGraphicsPipeline>,
    command_pool: VkCommandPool,
    command_buffers: Vec<vk::CommandBuffer>,
    image_available_semaphore: VkSemaphore,
    render_finished_semaphore: VkSemaphore,
}

impl VkRenderContext {
    pub fn init(window: &crate::display::Window) -> anyhow::Result<Self> {
        let required_extensions = window.required_extensions()?;

        let instance = Arc::new(VkInstance::new(&required_extensions));
        let surface = Arc::new(VkSurface::new(&window, instance.clone())?);

        let mut physical_devices = instance.get_physical_devices();
        physical_devices.sort_by(|a, b| b.priority().cmp(&a.priority()));

        let (physical_device, queue_family_index, properties) =
            VkPhysicalDevice::pick_physical_device_and_queue_by(
                &physical_devices,
                Some(vk::QueueFlags::GRAPHICS),
                Some(&surface),
            )
            .ok_or_else(|| anyhow::anyhow!("Unable to find physical device"))?;
        log::debug!(
            "Selected queue {properties:#?} for physical device {:?}",
            physical_device.properties.device_name_as_c_str()
        );

        let device = Arc::new(VkDevice::new_graphics_device(
            &instance,
            &physical_device,
            queue_family_index,
        )?);

        let swapchain = Arc::new(VkSwapchain::new(
            &window,
            &surface,
            &device,
            &physical_device,
        )?);

        let render_pass = Arc::new(VkRenderPass::new(&device, &swapchain)?);

        let framebuffers = swapchain
            .create_framebuffers(&render_pass)
            .ok_or_else(|| anyhow::anyhow!("Failed to get framebuffers"))?;

        let pipeline = Arc::new(VkGraphicsPipeline::new(&device, &swapchain, &render_pass)?);

        let command_pool = VkCommandPool::new(&device)?;

        let command_buffer_info = vk::CommandBufferAllocateInfo::default()
            .command_pool(command_pool.handle)
            .level(vk::CommandBufferLevel::PRIMARY)
            .command_buffer_count(framebuffers.len() as u32);

        // Destroyed with command pool
        let command_buffers = unsafe {
            device
                .handle
                .allocate_command_buffers(&command_buffer_info)?
        };

        // Same in VkGraphicsPipeline (TODO: Refactor this)
        let render_area = vk::Rect2D::default().extent(swapchain.surface_resolution);
        let clear_value = vk::ClearValue::default();
        for (index, command_buffer) in command_buffers.iter().enumerate() {
            let command_buffer_begin_info = vk::CommandBufferBeginInfo::default();
            unsafe {
                device
                    .handle
                    .begin_command_buffer(*command_buffer, &command_buffer_begin_info)?
            };

            let clear_values = [clear_value];
            let framebuffer = framebuffers[index].as_ref();
            let render_pass_begin_info = vk::RenderPassBeginInfo::default()
                .render_pass(render_pass.handle)
                .framebuffer(framebuffer.handle)
                .render_area(render_area)
                .clear_values(&clear_values);

            unsafe {
                device.handle.cmd_begin_render_pass(
                    *command_buffer,
                    &render_pass_begin_info,
                    vk::SubpassContents::INLINE,
                );
            };

            unsafe {
                device.handle.cmd_bind_pipeline(
                    *command_buffer,
                    vk::PipelineBindPoint::GRAPHICS,
                    pipeline.pipeline,
                )
            };

            unsafe { device.handle.cmd_draw(*command_buffer, 3, 1, 0, 0) };

            unsafe { device.handle.cmd_end_render_pass(*command_buffer) };

            unsafe { device.handle.end_command_buffer(*command_buffer)? };
        }

        let image_available_semaphore = VkSemaphore::new(device.clone())?;
        let render_finished_semaphore = VkSemaphore::new(device.clone())?;

        Ok(Self {
            instance,
            surface,
            device,

            swapchain,
            render_pass,
            framebuffers,
            pipeline,

            command_pool,
            command_buffers,

            image_available_semaphore,
            render_finished_semaphore,
        })
    }

    pub fn render(&mut self) -> anyhow::Result<()> {
        let queue = self
            .device
            .get_device_queue(0)
            .ok_or_else(|| anyhow::anyhow!("Failed to get a queue"))?;

        let (index, _) = self
            .swapchain
            .acquire_next_image(&self.image_available_semaphore)?;

        let wait_semaphores = [self.image_available_semaphore.handle];
        let signal_semaphores = [self.render_finished_semaphore.handle];
        let wait_stages = [vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT];
        let command_buffers_to_submit = [self.command_buffers[index as usize]];
        let submit_info = vk::SubmitInfo::default()
            .wait_semaphores(&wait_semaphores)
            .wait_dst_stage_mask(&wait_stages)
            .command_buffers(&command_buffers_to_submit)
            .signal_semaphores(&signal_semaphores);

        unsafe {
            self.device
                .handle
                .queue_submit(*queue, &[submit_info], vk::Fence::null())?
        };

        Ok(())
    }

    pub fn update_resolution(&mut self, width: u32, height: u32) -> anyhow::Result<()> {
        if let Some(swapchain) = Arc::get_mut(&mut self.swapchain) {
            swapchain.update_resolution(width, height)?;
        }

        Ok(())
    }

    pub fn exit(&self) {
        unsafe { self.device.handle.device_wait_idle().unwrap() }
    }
}