use std::{error::Error, sync::Arc};

use vulkano::{
    command_buffer::{AutoCommandBufferBuilder, PrimaryAutoCommandBuffer},
    descriptor_set::{
        allocator::StandardDescriptorSetAllocator, layout::DescriptorSetLayoutCreateInfo,
    },
    device::Device,
    format::Format,
    pipeline::{
        DynamicState, GraphicsPipeline, Pipeline, PipelineBindPoint, PipelineLayout,
        PipelineShaderStageCreateInfo,
        graphics::{
            GraphicsPipelineCreateInfo,
            color_blend::{ColorBlendAttachmentState, ColorBlendState},
            depth_stencil::{DepthState, DepthStencilState},
            input_assembly::InputAssemblyState,
            multisample::MultisampleState,
            rasterization::RasterizationState,
            subpass::PipelineRenderingCreateInfo,
            vertex_input::{Vertex, VertexDefinition},
            viewport::ViewportState,
        },
        layout::{PipelineDescriptorSetLayoutCreateInfo, PipelineLayoutCreateFlags},
    },
};

use crate::core::render::{
    primitives::{
        AsDescriptorSet, AsDescriptorSetLayoutBindings, AsRecordable, AsRenderableMesh,
        AsRenderableMeshInstance, mvp::Mvp, transform::TransformRaw, vertex::Vertex3D,
    },
    resources::{pipeline::LoadableGraphicsPipeline, texture::Texture},
};

pub struct SimplePipeline;

impl LoadableGraphicsPipeline for SimplePipeline {
    fn load(
        device: &Arc<Device>,
        swapchain_format: Format,
        depth_format: Format,
    ) -> Result<Arc<GraphicsPipeline>, 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())?;

        let vertex_input_state =
            [Vertex3D::per_vertex(), TransformRaw::per_instance()].definition(&vs)?;

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

        let mvp_bindings = Mvp::as_descriptor_set_layout_bindings();
        let texture_bindings = Texture::as_descriptor_set_layout_bindings();

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

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

        let create_info = PipelineDescriptorSetLayoutCreateInfo {
            set_layouts: vec![vertex_descriptor_set_layout, fragment_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_format)],
            depth_attachment_format: Some(depth_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(),
                )),
                depth_stencil_state: Some(DepthStencilState {
                    depth: Some(DepthState::simple()),
                    ..Default::default()
                }),
                dynamic_state: [DynamicState::Viewport].into_iter().collect(),
                subpass: Some(subpass.into()),
                ..GraphicsPipelineCreateInfo::layout(layout)
            },
        )?;

        Ok(pipeline)
    }

    fn pipeline_name() -> &'static str {
        "SimplePipeline"
    }
}

impl AsRecordable for SimplePipeline {
    fn record_bind_commands(
        builder: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        descriptor_set_allocator: &Arc<StandardDescriptorSetAllocator>,
        pipeline: &Arc<GraphicsPipeline>,
        mesh: &impl AsRenderableMesh,
        instances: &impl AsRenderableMeshInstance,
        descriptor_sets: Vec<Arc<dyn AsDescriptorSet>>,
    ) -> Result<(), Box<dyn Error>> {
        builder.bind_pipeline_graphics(pipeline.clone())?;

        if !descriptor_sets.is_empty() {
            let layouts = pipeline.layout().set_layouts();

            let descriptor_sets = descriptor_sets
                .iter()
                .enumerate()
                .map(|(layout_index, data)| {
                    data.as_descriptor_set(descriptor_set_allocator, &layouts[layout_index])
                })
                .collect::<Result<Vec<_>, _>>()?;

            builder.bind_descriptor_sets(
                PipelineBindPoint::Graphics,
                pipeline.layout().clone(),
                0,
                descriptor_sets,
            )?;
        }

        builder.bind_vertex_buffers(
            0,
            (
                mesh.vertex_buffer().clone(),
                instances.instance_buffer().clone(),
            ),
        )?;

        Ok(())
    }
}

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

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