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

use vulkano::{
    buffer::{Buffer, BufferCreateInfo, BufferUsage, Subbuffer},
    command_buffer::{AutoCommandBufferBuilder, PrimaryAutoCommandBuffer},
    descriptor_set::{
        DescriptorSet, WriteDescriptorSet,
        allocator::StandardDescriptorSetAllocator,
        layout::{DescriptorSetLayoutBinding, DescriptorSetLayoutCreateInfo, DescriptorType},
    },
    device::Device,
    format::Format,
    memory::allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
    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},
    },
    shader::ShaderStages,
};

use crate::core::render::{
    primitives::{mvp::Mvp, transform::TransformRaw, vertex::Vertex3D},
    texture::Texture,
};

const VERTICES: [Vertex3D; 4] = [
    Vertex3D {
        position: [-0.5, -0.5, 0.0],
        uv: [0.0, 0.0],
    },
    Vertex3D {
        position: [-0.5, 0.5, 0.0],
        uv: [0.0, 1.0],
    },
    Vertex3D {
        position: [0.5, -0.5, 0.0],
        uv: [1.0, 0.0],
    },
    Vertex3D {
        position: [0.5, 0.5, 0.0],
        uv: [1.0, 1.0],
    },
];

const INDICES: [u32; 6] = [0, 2, 1, 2, 3, 1];

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

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

pub struct Square {
    vertex_buffer: Subbuffer<[Vertex3D]>,
    index_buffer: Subbuffer<[u32]>,
    pipeline: Arc<GraphicsPipeline>,
}

impl Square {
    pub fn new(
        device: &Arc<Device>,
        memory_allocator: &Arc<StandardMemoryAllocator>,
        swapchain_format: Format,
        depth_format: Format,
    ) -> Result<Self, Box<dyn Error>> {
        let vertex_buffer = Buffer::from_iter(
            memory_allocator.clone(),
            BufferCreateInfo {
                usage: BufferUsage::VERTEX_BUFFER,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
                    | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                ..Default::default()
            },
            Vec::from_iter(VERTICES),
        )?;
        let index_buffer = Buffer::from_iter(
            memory_allocator.clone(),
            BufferCreateInfo {
                usage: BufferUsage::INDEX_BUFFER,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
                    | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                ..Default::default()
            },
            Vec::from_iter(INDICES),
        )?;

        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 vertex_bindings = BTreeMap::<u32, DescriptorSetLayoutBinding>::from_iter([(
            0,
            DescriptorSetLayoutBinding {
                stages: ShaderStages::VERTEX,
                ..DescriptorSetLayoutBinding::descriptor_type(DescriptorType::UniformBuffer)
            },
        )]);
        let fragment_bindings = BTreeMap::<u32, DescriptorSetLayoutBinding>::from_iter([
            (
                0,
                DescriptorSetLayoutBinding {
                    stages: ShaderStages::FRAGMENT,
                    ..DescriptorSetLayoutBinding::descriptor_type(DescriptorType::Sampler)
                },
            ),
            (
                1,
                DescriptorSetLayoutBinding {
                    stages: ShaderStages::FRAGMENT,
                    ..DescriptorSetLayoutBinding::descriptor_type(DescriptorType::SampledImage)
                },
            ),
        ]);

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

        let fragment_descriptor_set_layout = DescriptorSetLayoutCreateInfo {
            bindings: fragment_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(Self {
            vertex_buffer,
            index_buffer,
            pipeline,
        })
    }

    pub fn render(
        &self,
        command_buffer: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        descriptor_set_allocator: &Arc<StandardDescriptorSetAllocator>,
        mvp_uniform: &Subbuffer<[Mvp]>,
        transform_uniform: &Subbuffer<[TransformRaw]>,
        texture: &Texture,
    ) -> Result<(), Box<dyn Error>> {
        let layouts = self.pipeline.layout().set_layouts();

        let uniform_descriptor_set = DescriptorSet::new(
            descriptor_set_allocator.clone(),
            layouts[0].clone(),
            [WriteDescriptorSet::buffer(0, mvp_uniform.clone())],
            [],
        )?;

        let texture_descriptor_set = DescriptorSet::new(
            descriptor_set_allocator.clone(),
            layouts[1].clone(),
            [
                WriteDescriptorSet::sampler(0, texture.get_sampler().clone()),
                WriteDescriptorSet::image_view(1, texture.get_texture().clone()),
            ],
            [],
        )?;

        command_buffer.bind_pipeline_graphics(self.pipeline.clone())?;
        command_buffer.bind_descriptor_sets(
            PipelineBindPoint::Graphics,
            self.pipeline.layout().clone(),
            0,
            vec![uniform_descriptor_set, texture_descriptor_set],
        )?;
        command_buffer
            .bind_vertex_buffers(0, (self.vertex_buffer.clone(), transform_uniform.clone()))?;
        command_buffer.bind_index_buffer(self.index_buffer.clone())?;

        unsafe {
            command_buffer.draw_indexed(
                INDICES.len() as u32,
                transform_uniform.len() as u32,
                0,
                0,
                0,
            )?;
        }

        Ok(())
    }
}