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

use image::DynamicImage;
use vulkano::{
    Validated, VulkanError,
    buffer::{Buffer, BufferCreateInfo, BufferUsage},
    command_buffer::{AutoCommandBufferBuilder, CopyBufferToImageInfo, PrimaryAutoCommandBuffer},
    descriptor_set::{
        DescriptorSet, WriteDescriptorSet,
        allocator::StandardDescriptorSetAllocator,
        layout::{DescriptorSetLayout, DescriptorSetLayoutBinding, DescriptorType},
    },
    device::Device,
    format::Format,
    image::{
        Image, ImageCreateInfo, ImageType, ImageUsage,
        sampler::{Filter, Sampler, SamplerAddressMode, SamplerCreateInfo},
        view::ImageView,
    },
    memory::allocator::{AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator},
    shader::ShaderStages,
};

use crate::core::render::primitives::AsBindableDescriptorSet;

pub struct Texture {
    texture: Arc<ImageView>,
    sampler: Arc<Sampler>,
}

impl Texture {
    fn new(texture: Arc<ImageView>, sampler: Arc<Sampler>) -> Self {
        Self { texture, sampler }
    }

    pub fn from_file(
        device: &Arc<Device>,
        memory_allocator: &Arc<StandardMemoryAllocator>,
        builder: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        path: &str,
    ) -> Result<Self, Box<dyn Error>> {
        let _span = tracing::info_span!("texture_load_from_file", path = path);

        let bytes = std::fs::read(path)?;
        Self::from_bytes(device, memory_allocator, builder, &bytes)
    }

    pub fn from_bytes(
        device: &Arc<Device>,
        memory_allocator: &Arc<StandardMemoryAllocator>,
        builder: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        bytes: &[u8],
    ) -> Result<Self, Box<dyn Error>> {
        let image = image::load_from_memory(bytes)?;
        Self::from_dynamic_image(device, memory_allocator, builder, image)
    }

    pub fn from_dynamic_image(
        device: &Arc<Device>,
        memory_allocator: &Arc<StandardMemoryAllocator>,
        builder: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        image: DynamicImage,
    ) -> Result<Self, Box<dyn Error>> {
        let _span = tracing::info_span!("texture_from_dynamic_image");

        let image_data = image.to_rgba8();
        let image_dimensions = image_data.dimensions();
        let image_data = image_data.into_raw();

        let upload_buffer = Buffer::new_slice(
            memory_allocator.clone(),
            BufferCreateInfo {
                usage: BufferUsage::TRANSFER_SRC,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_HOST
                    | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                ..Default::default()
            },
            image_data.len() as u64,
        )?;

        {
            let buffer_data = &mut *upload_buffer.write()?;
            buffer_data.copy_from_slice(&image_data);
        }

        let image = Image::new(
            memory_allocator.clone(),
            ImageCreateInfo {
                image_type: ImageType::Dim2d,
                format: Format::R8G8B8A8_SRGB,
                extent: [image_dimensions.0, image_dimensions.1, 1],
                array_layers: 1,
                usage: ImageUsage::TRANSFER_DST | ImageUsage::SAMPLED,
                ..Default::default()
            },
            AllocationCreateInfo::default(),
        )?;

        builder.copy_buffer_to_image(CopyBufferToImageInfo::buffer_image(
            upload_buffer,
            image.clone(),
        ))?;

        let sampler = Sampler::new(
            device.clone(),
            SamplerCreateInfo {
                mag_filter: Filter::Linear,
                min_filter: Filter::Linear,
                address_mode: [SamplerAddressMode::Repeat; 3],
                ..Default::default()
            },
        )?;

        let image_view = ImageView::new_default(image)?;

        tracing::trace!("Texture loaded with dimensions {:?}", image_dimensions);

        Ok(Self::new(image_view, sampler))
    }

    pub fn get_texture(&self) -> &Arc<ImageView> {
        &self.texture
    }

    pub fn get_sampler(&self) -> &Arc<Sampler> {
        &self.sampler
    }
}

impl AsBindableDescriptorSet<Texture> for Texture {
    fn as_descriptor_set_layout_bindings() -> BTreeMap<u32, DescriptorSetLayoutBinding> {
        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)
                },
            ),
        ])
    }

    fn as_descriptor_set(
        descriptor_set_allocator: &Arc<StandardDescriptorSetAllocator>,
        layout: &Arc<DescriptorSetLayout>,
        data: &Texture,
    ) -> Result<Arc<DescriptorSet>, Validated<VulkanError>> {
        DescriptorSet::new(
            descriptor_set_allocator.clone(),
            layout.clone(),
            [
                WriteDescriptorSet::sampler(0, data.sampler.clone()),
                WriteDescriptorSet::image_view(1, data.texture.clone()),
            ],
            [],
        )
    }
}