#![deny(unsafe_code)]
use crate::error::*;
use crate::assets::asset;
use crate::graphics::{ Permanent, ForReinit, WindowDressing, Model, Texture };
use crate::shader_data::{ Vertex, UniformBlock, VertexPushBlock };

use std::mem::size_of;

use vulkanalia::Device;
use vulkanalia::vk::{ self, Handle, HasBuilder, DeviceV1_0 };


//   Render is a state object that collects the Vulkan graphics objects which
// need to be regenerated or modified when the window changes, as with
// WindowDressing, and which are also used as part of rendering.
#[derive(Debug)]
pub struct Render {
  pub render_pass: vk::RenderPass,

  pub pipeline: vk::Pipeline,
  pub pipeline_layout: vk::PipelineLayout,

  pub model: Option<Model>,

  pub framebuffers: Vec<vk::Framebuffer>,
  pub command_buffers: Vec<vk::CommandBuffer>,
  pub descriptor_sets: Vec<vk::DescriptorSet>,
}


impl Render {
  pub fn new(permanent: &Permanent, for_reinit: &ForReinit,
             window_dressing: &WindowDressing, texture: &Texture)
      -> Result<Self>
  {
    let device = &permanent.device;
    let primary_command_pool = &permanent.primary_command_pool;
    let sample_count = for_reinit.sample_count;
    let descriptor_set_layout = &for_reinit.descriptor_set_layout;
    let swapchain = &window_dressing.swapchain;
    let depth_format = &window_dressing.depth_format;
    let color_image_view = &window_dressing.color_image_view;
    let depth_image_view = &window_dressing.depth_image_view;
    let uniform_buffers = &window_dressing.uniform_buffers;
    let descriptor_pool = &window_dressing.descriptor_pool;
    let sampler = &window_dressing.sampler;

    let render_pass = init_render_pass(device, sample_count,
                                       &swapchain.format, &depth_format)?;

    let (pipeline_layout, pipeline)
            = init_pipeline(device, descriptor_set_layout, &swapchain.extent,
                            sample_count, &render_pass)?;

    let framebuffers = init_framebuffers(
            device, &swapchain.extent, &swapchain.image_views,
            &color_image_view, &depth_image_view, &render_pass)?;

    let command_buffers = init_command_buffers(device, &framebuffers,
                                               primary_command_pool)?;

    let descriptor_sets
            = init_descriptor_sets(device, descriptor_set_layout,
                                   &uniform_buffers, &descriptor_pool,
                                   swapchain.images.len(),
                                   &texture.image_view, &sampler)?;

    let model = None;

    Ok(Render {
      render_pass,
      pipeline,
      pipeline_layout,
      model,
      framebuffers,
      command_buffers,
      descriptor_sets,
    })
  }

  //   This relies on its caller to have already waited for the device to be
  // idle.
  pub fn reinit(&mut self,
                permanent: &Permanent, for_reinit: &ForReinit,
                window_dressing: &WindowDressing, texture: &Texture)
      -> Result<()>
  {
    self.destroy_replaceable(permanent);

    let device = &permanent.device;
    let primary_command_pool = &permanent.primary_command_pool;
    let sample_count = for_reinit.sample_count;
    let descriptor_set_layout = &for_reinit.descriptor_set_layout;
    let swapchain = &window_dressing.swapchain;
    let depth_format = &window_dressing.depth_format;
    let color_image_view = &window_dressing.color_image_view;
    let depth_image_view = &window_dressing.depth_image_view;
    let uniform_buffers = &window_dressing.uniform_buffers;
    let descriptor_pool = &window_dressing.descriptor_pool;
    let sampler = &window_dressing.sampler;

    let render_pass = init_render_pass(device, sample_count,
                                       &swapchain.format, &depth_format)?;

    let (pipeline_layout, pipeline)
            = init_pipeline(device, descriptor_set_layout, &swapchain.extent,
                            sample_count, &render_pass)?;

    let framebuffers = init_framebuffers(
            device, &swapchain.extent, &swapchain.image_views,
            &color_image_view, &depth_image_view, &render_pass)?;

    // Notice that we reused the command pool.
    let command_buffers = init_command_buffers(device, &framebuffers,
                                               primary_command_pool)?;

    let descriptor_sets
            = init_descriptor_sets(device, descriptor_set_layout,
                                   &uniform_buffers, &descriptor_pool,
                                   swapchain.images.len(),
                                   &texture.image_view, sampler)?;

    self.render_pass = render_pass;
    self.pipeline = pipeline;
    self.pipeline_layout = pipeline_layout;
    self.framebuffers = framebuffers;
    self.command_buffers = command_buffers;
    self.descriptor_sets = descriptor_sets;

    Ok(())
  }

  //   This relies on its caller to have already waited for the device to be
  // idle.
  #[allow(unsafe_code)]
  pub fn destroy(mut self, permanent: &Permanent) {
    self.destroy_replaceable(permanent);

    if let Some(model) = self.model {
      model.destroy(&permanent.device);
    }
  }

  #[allow(unsafe_code)]
  fn destroy_replaceable(&mut self, permanent: &Permanent) {
    let device = &permanent.device;

    for framebuffer in &self.framebuffers {
      unsafe { device.destroy_framebuffer(*framebuffer, None) };
    }

    //   Notice that we free the buffers in the pool, but do not destroy the
    // pool itself. Notice also that we only do this for the primary command
    // pool, because that's the only one where we've kept track of the
    // buffers. We promise ourselves to free buffers in the transient pool
    // immediately after using them.
    unsafe {
      device.free_command_buffers(permanent.primary_command_pool,
                                  &self.command_buffers)
    };

    unsafe { device.destroy_pipeline(self.pipeline, None) };
    unsafe { device.destroy_pipeline_layout(self.pipeline_layout, None) };
    unsafe { device.destroy_render_pass(self.render_pass, None) };
  }

  pub fn set_model(&mut self, model: Model) {
    self.model = Some(model);
  }
}


#[allow(unsafe_code)]
fn init_render_pass(device: &Device, sample_count: vk::SampleCountFlags,
                    color_format: &vk::Format, depth_format: &vk::Format)
    -> Result<vk::RenderPass>
{
  let color_attachment = vk::AttachmentDescription::builder()
          .format(*color_format)
          .samples(sample_count)
          .load_op(vk::AttachmentLoadOp::CLEAR)
          .store_op(vk::AttachmentStoreOp::STORE)
          .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
          .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
          .initial_layout(vk::ImageLayout::UNDEFINED)
          .final_layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

  let color_attachment_reference = vk::AttachmentReference::builder()
          .attachment(0)
          .layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

  let depth_attachment = vk::AttachmentDescription::builder()
          .format(*depth_format)
          .samples(sample_count)
          .load_op(vk::AttachmentLoadOp::CLEAR)
          .store_op(vk::AttachmentStoreOp::DONT_CARE)
          .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
          .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
          .initial_layout(vk::ImageLayout::UNDEFINED)
          .final_layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

  let depth_attachment_reference = vk::AttachmentReference::builder()
          .attachment(1)
          .layout(vk::ImageLayout::DEPTH_STENCIL_ATTACHMENT_OPTIMAL);

  let color_resolve_attachment = vk::AttachmentDescription::builder()
          .format(*color_format)
          .samples(vk::SampleCountFlags::_1)
          .load_op(vk::AttachmentLoadOp::DONT_CARE)
          .store_op(vk::AttachmentStoreOp::STORE)
          .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
          .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
          .initial_layout(vk::ImageLayout::UNDEFINED)
          .final_layout(vk::ImageLayout::PRESENT_SRC_KHR);

  let color_resolve_attachment_reference = vk::AttachmentReference::builder()
          .attachment(2)
          .layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL);

  let color_attachments = [color_attachment_reference];
  let resolve_attachments = [color_resolve_attachment_reference];
  let subpass = vk::SubpassDescription::builder()
          .pipeline_bind_point(vk::PipelineBindPoint::GRAPHICS)
          .color_attachments(&color_attachments)
          .depth_stencil_attachment(&depth_attachment_reference)
          .resolve_attachments(&resolve_attachments);

  let dependency = vk::SubpassDependency::builder()
          .src_subpass(vk::SUBPASS_EXTERNAL)
          .src_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
                          | vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS)
          .src_access_mask(vk::AccessFlags::empty())
          .dst_subpass(0)
          .dst_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT
                          | vk::PipelineStageFlags::EARLY_FRAGMENT_TESTS)
          .dst_access_mask(vk::AccessFlags::COLOR_ATTACHMENT_WRITE
                           | vk::AccessFlags::DEPTH_STENCIL_ATTACHMENT_WRITE);

  let render_attachments = [color_attachment,
                            depth_attachment,
                            color_resolve_attachment];
  let subpasses = [subpass];
  let dependencies = [dependency];
  let render_pass_info = vk::RenderPassCreateInfo::builder()
          .attachments(&render_attachments)
          .subpasses(&subpasses)
          .dependencies(&dependencies);

  let render_pass = unsafe {
    device.create_render_pass(&render_pass_info, None)
  }?;

  Ok(render_pass)
}


#[allow(unsafe_code)]
fn init_pipeline(device: &Device,
                 descriptor_set_layout: &vk::DescriptorSetLayout,
                 extent: &vk::Extent2D, sample_count: vk::SampleCountFlags,
                 render_pass: &vk::RenderPass)
    -> Result<(vk::PipelineLayout, vk::Pipeline)>
{
  let vertex_binary = asset("shaders/shader.vert.spv")?;
  let fragment_binary = asset("shaders/shader.frag.spv")?;

  let vertex_module = Permanent::load_spirv_shader_module(
          device, vertex_binary)?;
  let fragment_module = Permanent::load_spirv_shader_module(
          device, fragment_binary)?;

  let vertex_stage_info = vk::PipelineShaderStageCreateInfo::builder()
          .stage(vk::ShaderStageFlags::VERTEX)
          .module(vertex_module)
          .name(b"main\0");

  let fragment_stage_info = vk::PipelineShaderStageCreateInfo::builder()
          .stage(vk::ShaderStageFlags::FRAGMENT)
          .module(fragment_module)
          .name(b"main\0");

  let binding_descriptions = [Vertex::<f32>::binding_description()];
  let attribute_descriptions = Vertex::<f32>::attribute_descriptions();
  let vertex_input_state_info
          = vk::PipelineVertexInputStateCreateInfo::builder()
                .vertex_binding_descriptions(&binding_descriptions)
                .vertex_attribute_descriptions(&attribute_descriptions);

  let input_assembly_state_info
          = vk::PipelineInputAssemblyStateCreateInfo::builder()
                .topology(vk::PrimitiveTopology::TRIANGLE_LIST)
                .primitive_restart_enable(false);

  let viewport = vk::Viewport::builder()
          .x(0.0)
          .y(0.0)
          .width(extent.width as f32)
          .height(extent.height as f32)
          .min_depth(0.0)
          .max_depth(1.0);
  let viewports = [viewport];

  let scissor = vk::Rect2D::builder()
          .offset(vk::Offset2D { x: 0, y: 0 })
          .extent(*extent);
  let scissor_list = [scissor];

  let viewport_state_info = vk::PipelineViewportStateCreateInfo::builder()
          .viewports(&viewports)
          .scissors(&scissor_list);

  let rasterizer_state_info
          = vk::PipelineRasterizationStateCreateInfo::builder()
                .depth_clamp_enable(false)
                .rasterizer_discard_enable(false)
                .polygon_mode(vk::PolygonMode::FILL)
                .line_width(1.0)
                .cull_mode(vk::CullModeFlags::BACK)
                .front_face(vk::FrontFace::CLOCKWISE)
                .depth_bias_enable(false);

  let multisample_state_info
          = vk::PipelineMultisampleStateCreateInfo::builder()
                .sample_shading_enable(false)
                .rasterization_samples(sample_count);

  let depth_state_info = vk::PipelineDepthStencilStateCreateInfo::builder()
          .depth_test_enable(true)
          .depth_write_enable(true)
          .depth_compare_op(vk::CompareOp::LESS)
          .depth_bounds_test_enable(false)
          .min_depth_bounds(0.0)
          .max_depth_bounds(1.0)
          .stencil_test_enable(false);

  let blend_attachment_info = vk::PipelineColorBlendAttachmentState::builder()
          .color_write_mask(vk::ColorComponentFlags::all())
          .blend_enable(false)
          .src_color_blend_factor(vk::BlendFactor::ONE)
          .dst_color_blend_factor(vk::BlendFactor::ZERO)
          .color_blend_op(vk::BlendOp::ADD)
          .src_alpha_blend_factor(vk::BlendFactor::ONE)
          .dst_alpha_blend_factor(vk::BlendFactor::ZERO)
          .alpha_blend_op(vk::BlendOp::ADD);
  let blend_attachments = [blend_attachment_info];

  let blend_info = vk::PipelineColorBlendStateCreateInfo::builder()
          .logic_op_enable(false)
          .logic_op(vk::LogicOp::COPY)
          .attachments(&blend_attachments)
          .blend_constants([0.0, 0.0, 0.0, 0.0]);

  let vertex_push_constant_range = vk::PushConstantRange::builder()
          .stage_flags(vk::ShaderStageFlags::VERTEX)
          .offset(0)
          .size(size_of::<VertexPushBlock<f32>>() as u32);

  let layouts = [*descriptor_set_layout];
  let push_constant_ranges = [vertex_push_constant_range];
  let pipeline_layout_info = vk::PipelineLayoutCreateInfo::builder()
                                 .set_layouts(&layouts)
                                 .push_constant_ranges(&push_constant_ranges);

  let pipeline_layout = unsafe {
    device.create_pipeline_layout(&pipeline_layout_info, None)
  }?;

  let stages = [vertex_stage_info, fragment_stage_info];
  let pipeline_info = vk::GraphicsPipelineCreateInfo::builder()
          .stages(&stages)
          .vertex_input_state(&vertex_input_state_info)
          .input_assembly_state(&input_assembly_state_info)
          .viewport_state(&viewport_state_info)
          .rasterization_state(&rasterizer_state_info)
          .multisample_state(&multisample_state_info)
          .depth_stencil_state(&depth_state_info)
          .color_blend_state(&blend_info)
          .layout(pipeline_layout)
          .render_pass(*render_pass)
          .subpass(0);

  let pipeline = unsafe {
    device.create_graphics_pipelines(vk::PipelineCache::null(),
                                     &[pipeline_info], None)
  }?.0[0];

  unsafe {
    device.destroy_shader_module(vertex_module, None);
    device.destroy_shader_module(fragment_module, None);
  };

  Ok((pipeline_layout, pipeline))
}


#[allow(unsafe_code)]
fn init_framebuffers(device: &Device, extent: &vk::Extent2D,
                     swapchain_image_views: &Vec<vk::ImageView>,
                     color_image_view: &vk::ImageView,
                     depth_image_view: &vk::ImageView,
                     render_pass: &vk::RenderPass)
    -> Result<Vec<vk::Framebuffer>>
{
  let mut framebuffers = Vec::new();

  for color_resolve_image_view in swapchain_image_views {
    let attachments = [*color_image_view,
                       *depth_image_view,
                       *color_resolve_image_view];

    let framebuffer_info = vk::FramebufferCreateInfo::builder()
            .render_pass(*render_pass)
            .attachments(&attachments)
            .width(extent.width)
            .height(extent.height)
            .layers(1);

    let framebuffer = unsafe {
      device.create_framebuffer(&framebuffer_info, None)
    }?;

    framebuffers.push(framebuffer);
  }

  Ok(framebuffers)
}


#[allow(unsafe_code)]
fn init_command_buffers(device: &Device,
                        framebuffers: &Vec<vk::Framebuffer>,
                        command_pool: &vk::CommandPool)
    -> Result<Vec<vk::CommandBuffer>>
{
  let command_buffer_allocation_info
          = vk::CommandBufferAllocateInfo::builder()
                .command_pool(*command_pool)
                .level(vk::CommandBufferLevel::PRIMARY)
                .command_buffer_count(framebuffers.len() as u32);
  let command_buffers = unsafe {
    device.allocate_command_buffers(&command_buffer_allocation_info)
  }?;

  Ok(command_buffers)
}


#[allow(unsafe_code)]
fn init_descriptor_sets(device: &Device, layout: &vk::DescriptorSetLayout,
                        buffers: &Vec<vk::Buffer>, pool: &vk::DescriptorPool,
                        count: usize, texture_image_view: &vk::ImageView,
                        sampler: &vk::Sampler)
    -> Result<Vec<vk::DescriptorSet>>
{
  let layouts = vec![*layout; count];
  let set_info = vk::DescriptorSetAllocateInfo::builder()
          .descriptor_pool(*pool)
          .set_layouts(&layouts);
  let sets = unsafe { device.allocate_descriptor_sets(&set_info) }?;

  for index in 0 .. count {
    let buffer_info = vk::DescriptorBufferInfo::builder()
            .buffer(buffers[index])
            .offset(0)
            .range(size_of::<UniformBlock<f32>>() as vk::DeviceSize);

    let buffer_info_list = [buffer_info];
    let uniform_block_write_info = vk::WriteDescriptorSet::builder()
            .dst_set(sets[index])
            .dst_binding(0)
            .dst_array_element(0)
            .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER)
            .buffer_info(&buffer_info_list);

    let image_info = vk::DescriptorImageInfo::builder()
            .image_layout(vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL)
            .image_view(*texture_image_view)
            .sampler(*sampler);
    let image_info_list = [image_info];
    let sampler_write_info = vk::WriteDescriptorSet::builder()
            .dst_set(sets[index])
            .dst_binding(1)
            .dst_array_element(0)
            .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
            .image_info(&image_info_list);

    let write_info_list = [uniform_block_write_info, sampler_write_info];
    let copy_info_list: [vk::CopyDescriptorSet; 0] = [];

    unsafe {
      device.update_descriptor_sets(&write_info_list, &copy_info_list)
    };
  }

  Ok(sets)
}