neiam /dms

use async_trait::async_trait;
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
use tokio::fs;
use tokio::io::AsyncWriteExt;

use super::{Storage, StorageError, StoredObject};

#[derive(Clone)]
pub struct LocalStorage {
    root: PathBuf,
}

impl LocalStorage {
    pub fn new(root: impl Into<PathBuf>) -> Self {
        Self { root: root.into() }
    }

    fn relative_for(hash_hex: &str, extension: Option<&str>) -> String {
        // sha256/aa/bb/full-hash[.ext] — fan out to keep directory sizes sane
        let aa = &hash_hex[0..2];
        let bb = &hash_hex[2..4];
        match extension {
            Some(ext) if !ext.is_empty() => format!("sha256/{aa}/{bb}/{hash_hex}.{ext}"),
            _ => format!("sha256/{aa}/{bb}/{hash_hex}"),
        }
    }
}

#[async_trait]
impl Storage for LocalStorage {
    async fn put(&self, bytes: &[u8], extension: Option<&str>) -> Result<StoredObject, StorageError> {
        let mut hasher = Sha256::new();
        hasher.update(bytes);
        let hash_hex = hex::encode(hasher.finalize());
        let relative = Self::relative_for(&hash_hex, extension);
        let abs = self.root.join(&relative);

        if let Some(parent) = abs.parent() {
            fs::create_dir_all(parent).await?;
        }
        if !abs.exists() {
            let mut f = fs::File::create(&abs).await?;
            f.write_all(bytes).await?;
            f.sync_all().await?;
        }

        Ok(StoredObject {
            relative_path: relative,
            size_bytes: bytes.len() as u64,
            sha256_hex: hash_hex,
        })
    }

    fn absolute_path(&self, relative: &str) -> PathBuf {
        self.root.join(relative)
    }

    async fn put_derived(
        &self,
        original_relative: &str,
        suffix: &str,
        bytes: &[u8],
    ) -> Result<String, StorageError> {
        let ext = Path::new(original_relative)
            .extension()
            .and_then(|s| s.to_str())
            .unwrap_or("bin")
            .to_string();
        self.put_derived_as(original_relative, suffix, &ext, bytes).await
    }

    async fn put_derived_as(
        &self,
        original_relative: &str,
        suffix: &str,
        extension: &str,
        bytes: &[u8],
    ) -> Result<String, StorageError> {
        let derived_rel = super::derived_relative(original_relative, suffix, extension);

        let abs = self.root.join(&derived_rel);
        if let Some(p) = abs.parent() {
            fs::create_dir_all(p).await?;
        }
        let mut f = fs::File::create(&abs).await?;
        f.write_all(bytes).await?;
        f.sync_all().await?;
        Ok(derived_rel)
    }
}