use ash::vk;

use ash::Device;

use crate::PhysicalDeviceInfo;
#[derive(Debug, Clone, Copy)]
pub struct PipelineCache {
    #[allow(dead_code)]
    key: u128,
    pub raw: vk::PipelineCache,
}

impl PipelineCache {
    pub(crate) unsafe fn destroy(self, device: &ash::Device) {
        tracing::info!("destroying pipeline cache with key {:x}", self.key);
        if let Ok(data) = self.export(device) {
            tracing::info!(
                "exported pipeline cache with key {:x} and size {} bytes",
                self.key,
                data.len()
            );
            _ = Self::write_to_disk(self.key, &data).inspect_err(|err| {
                tracing::error!("failed to write pipeline cache to disk: {err}");
            });
        }
        unsafe { device.destroy_pipeline_cache(self.raw, None) };
    }

    const MAGIC: [u8; 4] = *b"VYPC";
    const KEY_VERSION: u32 = 1;
    const PATH: &'static str = "pipeline_cache.bin";
    fn calculate_key(adapter: &PhysicalDeviceInfo) -> u128 {
        use md5::Digest;
        let mut hasher = md5::Md5::new();
        let props = &adapter.properties;
        hasher.update(bytemuck::bytes_of(&[
            props.core.vendor_id,
            props.core.api_version,
            props.core.device_id,
            props.core.driver_version,
        ]));
        u128::from_le_bytes(hasher.finalize().into())
    }

    fn load_from_disk(key: u128) -> Option<(u128, Vec<u8>)> {
        use std::io::Read;

        let file = std::fs::File::open(Self::PATH).ok()?;
        let mut reader = std::io::BufReader::new(file);
        let mut magic = [0; 4];
        reader.read_exact(&mut magic).ok()?;
        if magic != Self::MAGIC {
            return None;
        }

        let mut version = 0;
        reader
            .read_exact(bytemuck::bytes_of_mut(&mut version))
            .ok()?;
        if version != Self::KEY_VERSION {
            return None;
        }

        let mut disk_key = 0;
        reader
            .read_exact(bytemuck::bytes_of_mut(&mut disk_key))
            .ok()?;
        if disk_key != key {
            return None;
        }

        let mut data = Vec::new();
        reader.read_to_end(&mut data).ok()?;

        Some((key, data))
    }

    fn write_to_disk(key: u128, data: &[u8]) -> std::io::Result<()> {
        use std::io::Write;

        let file = std::fs::File::create(Self::PATH)?;
        let mut writer = std::io::BufWriter::new(file);
        writer.write_all(&Self::MAGIC)?;
        writer.write_all(bytemuck::bytes_of(&Self::KEY_VERSION))?;
        writer.write_all(bytemuck::bytes_of(&key))?;
        writer.write_all(data)?;
        Ok(())
    }

    pub fn new(device: &Device, adapter: &PhysicalDeviceInfo) -> crate::Result<Self> {
        let key = Self::calculate_key(adapter);
        let data = Self::load_from_disk(key).map(|(key, data)| {
            tracing::info!("loaded pipeline cache from disk with key {key:x}");
            data
        });

        let info = vk::PipelineCacheCreateInfo::default()
            // .flags(vk::PipelineCacheCreateFlags::EXTERNALLY_SYNCHRONIZED)
            .initial_data(data.as_deref().unwrap_or_default());

        let cache = unsafe { device.create_pipeline_cache(&info, None)? };

        Ok(Self { key, raw: cache })
    }

    pub fn export(&self, device: &ash::Device) -> crate::Result<Vec<u8>> {
        let data = unsafe { device.get_pipeline_cache_data(self.raw)? };
        Ok(data)
    }
}