use std::collections::BTreeMap;
use std::error::Error;
use std::sync::Arc;
use vulkano::descriptor_set::layout::{
    DescriptorSetLayoutBinding, DescriptorSetLayoutCreateInfo, DescriptorType,
};
use vulkano::device::Device;
use vulkano::pipeline::graphics::GraphicsPipelineCreateInfo;
use vulkano::pipeline::graphics::color_blend::{ColorBlendAttachmentState, ColorBlendState};
use vulkano::pipeline::graphics::input_assembly::InputAssemblyState;
use vulkano::pipeline::graphics::multisample::MultisampleState;
use vulkano::pipeline::graphics::rasterization::RasterizationState;
use vulkano::pipeline::graphics::subpass::PipelineRenderingCreateInfo;
use vulkano::pipeline::graphics::vertex_input::{Vertex, VertexDefinition};
use vulkano::pipeline::graphics::viewport::ViewportState;
use vulkano::pipeline::layout::{PipelineDescriptorSetLayoutCreateInfo, PipelineLayoutCreateFlags};
use vulkano::pipeline::{
    DynamicState, GraphicsPipeline, PipelineLayout, PipelineShaderStageCreateInfo,
};
use vulkano::shader::{EntryPoint, ShaderStages};
use vulkano::swapchain::Swapchain;

use crate::vulkan::Vertex2D;

pub mod shaders {
    pub mod vs {
        vulkano_shaders::shader! {
            ty: "vertex",
            path: r"res/shaders/vertex.vert",
        }
    }

    pub mod fs {
        vulkano_shaders::shader! {
            ty: "fragment",
            path: r"res/shaders/vertex.frag",
        }
    }
}

pub fn create_triangle_pipeline(
    device: &Arc<Device>,
    swapchain: &Arc<Swapchain>,
) -> Result<Arc<GraphicsPipeline>, Box<dyn Error>> {
    let (vs, fs) = load_shaders(device)?;
    let vertex_input_state = Vertex2D::per_vertex().definition(&vs)?;

    let stages = [
        PipelineShaderStageCreateInfo::new(vs),
        PipelineShaderStageCreateInfo::new(fs),
    ];

    let mut bindings = BTreeMap::<u32, DescriptorSetLayoutBinding>::new();
    let mut descriptor_set_layout_binding =
        DescriptorSetLayoutBinding::descriptor_type(DescriptorType::UniformBuffer);
    descriptor_set_layout_binding.stages = ShaderStages::VERTEX;
    bindings.insert(0, descriptor_set_layout_binding);

    let descriptor_set_layout = DescriptorSetLayoutCreateInfo {
        bindings,
        ..Default::default()
    };

    let create_info = PipelineDescriptorSetLayoutCreateInfo {
        set_layouts: vec![descriptor_set_layout],
        flags: PipelineLayoutCreateFlags::default(),
        push_constant_ranges: vec![],
    }
    .into_pipeline_layout_create_info(device.clone())?;

    let layout = PipelineLayout::new(device.clone(), create_info)?;

    let subpass = PipelineRenderingCreateInfo {
        color_attachment_formats: vec![Some(swapchain.image_format())],
        ..Default::default()
    };

    let pipeline = GraphicsPipeline::new(
        device.clone(),
        None,
        GraphicsPipelineCreateInfo {
            stages: stages.into_iter().collect(),
            vertex_input_state: Some(vertex_input_state),
            input_assembly_state: Some(InputAssemblyState::default()),
            viewport_state: Some(ViewportState::default()),
            rasterization_state: Some(RasterizationState::default()),
            multisample_state: Some(MultisampleState::default()),
            color_blend_state: Some(ColorBlendState::with_attachment_states(
                subpass.color_attachment_formats.len() as u32,
                ColorBlendAttachmentState::default(),
            )),
            dynamic_state: [DynamicState::Viewport].into_iter().collect(),
            subpass: Some(subpass.into()),
            ..GraphicsPipelineCreateInfo::layout(layout)
        },
    )?;

    Ok(pipeline)
}

fn load_shaders(device: &Arc<Device>) -> Result<(EntryPoint, EntryPoint), Box<dyn Error>> {
    let vs = shaders::vs::load(device.clone())?
        .entry_point("main")
        .ok_or("Failed find main entry point of vertex shader".to_string())?;

    let fs = shaders::fs::load(device.clone())?
        .entry_point("main")
        .ok_or("Failed find main entry point of fragment shader".to_string())?;

    Ok((vs, fs))
}