use std::sync::Arc;

use vulkano::{
    descriptor_set::layout::DescriptorType,
    device::Device,
    image::sampler::Sampler,
    pipeline::{
        graphics::{
            color_blend::{ColorBlendAttachmentState, ColorBlendState},
            depth_stencil::{DepthState, DepthStencilState},
            input_assembly::InputAssemblyState,
            multisample::MultisampleState,
            rasterization::RasterizationState,
            vertex_input::{Vertex, VertexDefinition},
            viewport::{Viewport, ViewportState},
            GraphicsPipelineCreateInfo,
        },
        layout::PipelineDescriptorSetLayoutCreateInfo,
        GraphicsPipeline, PipelineLayout, PipelineShaderStageCreateInfo,
    },
    render_pass::{RenderPass, Subpass},
    shader::EntryPoint,
};

use crate::error::Error;

pub mod plain {
    pub mod vs {
        vulkano_shaders::shader! {
            ty: "vertex",
            path: "src/render/shaders/plain.vert",
        }
    }
    pub mod fs {
        vulkano_shaders::shader! {
            ty: "fragment",
            path: "src/render/shaders/plain.frag",
        }
    }

    pub use vs::SceneData;
}

pub fn create_two_stage_3d_pipeline<V: Vertex>(
    device: &Arc<Device>,
    render_pass: &Arc<RenderPass>,
    viewport: Viewport,
    vertex_shader: &EntryPoint,
    fragment_shader: &EntryPoint,
    samplers: Vec<Arc<Sampler>>,
) -> Result<Arc<GraphicsPipeline>, Error> {
    let vertex_input_state = V::per_vertex().definition(vertex_shader)?;
    let stages = [
        PipelineShaderStageCreateInfo::new(vertex_shader.clone()),
        PipelineShaderStageCreateInfo::new(fragment_shader.clone()),
    ];

    let layout = {
        let mut layout_create_info = PipelineDescriptorSetLayoutCreateInfo::from_stages(&stages);

        layout_create_info.set_layouts[0]
            .bindings
            .get_mut(&0)
            .unwrap()
            .descriptor_type = DescriptorType::UniformBufferDynamic;

        layout_create_info.set_layouts[1]
            .bindings
            .get_mut(&0)
            .unwrap()
            .immutable_samplers = samplers;

        PipelineLayout::new(
            device.clone(),
            layout_create_info.into_pipeline_layout_create_info(device.clone())?,
        )?
    };

    let subpass = Subpass::from(render_pass.clone(), 0).unwrap();

    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 {
                viewports: [viewport].into_iter().collect(),
                ..Default::default()
            }),
            rasterization_state: Some(RasterizationState::default()),
            multisample_state: Some(MultisampleState {
                rasterization_samples: subpass.num_samples().unwrap(),
                ..Default::default()
            }),
            color_blend_state: Some(ColorBlendState::with_attachment_states(
                subpass.num_color_attachments(),
                ColorBlendAttachmentState::default(),
            )),
            depth_stencil_state: Some(DepthStencilState {
                depth: Some(DepthState::simple()),
                ..Default::default()
            }),
            subpass: Some(subpass.into()),
            ..GraphicsPipelineCreateInfo::layout(layout)
        },
    )?;

    Ok(pipeline)
}