lib.rs

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, Mutex};

pub type Bytes = Vec<u8>;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DbError {
    NotSupported,
    NotFound,
    Message(String),
}

impl fmt::Display for DbError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DbError::NotSupported => write!(f, "operation not supported"),
            DbError::NotFound => write!(f, "not found"),
            DbError::Message(s) => write!(f, "{s}"),
        }
    }
}

impl std::error::Error for DbError {}

pub type Result<T> = std::result::Result<T, DbError>;

pub const NS_DEFAULT: &str = "default";
pub const NS_IMAGES_META: &str = "images_meta";
pub const NS_IMAGES_BLOB: &str = "images_blob";

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
pub struct ImageMeta {
    pub id: String,
    pub content_type: String,
    pub width: u32,
    pub height: u32,
    pub size: u64,
    pub created_at_ms: u64,
}

#[async_trait(?Send)]
pub trait Engine {
    async fn put(&self, ns: &str, key: &str, value: Bytes) -> Result<()>;
    async fn get(&self, ns: &str, key: &str) -> Result<Option<Bytes>>;
    async fn delete(&self, ns: &str, key: &str) -> Result<()>;
    async fn scan_prefix(&self, ns: &str, prefix: &str) -> Result<Vec<String>>;
}

pub async fn put_kv<E: Engine>(engine: &E, key: &str, value: Bytes) -> Result<()> {
    engine.put(NS_DEFAULT, key, value).await
}

pub async fn get_kv<E: Engine>(engine: &E, key: &str) -> Result<Option<Bytes>> {
    engine.get(NS_DEFAULT, key).await
}

pub async fn delete_kv<E: Engine>(engine: &E, key: &str) -> Result<()> {
    engine.delete(NS_DEFAULT, key).await
}

pub async fn put_image<E: Engine>(engine: &E, meta: &ImageMeta, data: Bytes) -> Result<()> {
    let meta_key = &meta.id;
    let meta_bytes = serde_json::to_vec(meta).map_err(|e| DbError::Message(e.to_string()))?;
    engine.put(NS_IMAGES_META, meta_key, meta_bytes).await?;
    engine.put(NS_IMAGES_BLOB, meta_key, data).await
}

pub async fn get_image<E: Engine>(engine: &E, id: &str) -> Result<Option<(ImageMeta, Bytes)>> {
    let meta_bytes_opt = engine.get(NS_IMAGES_META, id).await?;
    let Some(meta_bytes) = meta_bytes_opt else {
        return Ok(None);
    };
    let data_opt = engine.get(NS_IMAGES_BLOB, id).await?;
    let Some(data) = data_opt else {
        return Err(DbError::NotFound);
    };
    let meta: ImageMeta =
        serde_json::from_slice(&meta_bytes).map_err(|e| DbError::Message(e.to_string()))?;
    Ok(Some((meta, data)))
}

pub async fn delete_image<E: Engine>(engine: &E, id: &str) -> Result<()> {
    engine.delete(NS_IMAGES_META, id).await?;
    engine.delete(NS_IMAGES_BLOB, id).await
}

pub async fn list_images<E: Engine>(engine: &E) -> Result<Vec<ImageMeta>> {
    let keys = engine.scan_prefix(NS_IMAGES_META, "").await?;
    let mut out = Vec::new();
    for k in keys {
        let meta_bytes_opt = engine.get(NS_IMAGES_META, &k).await?;
        if let Some(meta_bytes) = meta_bytes_opt {
            let meta: ImageMeta =
                serde_json::from_slice(&meta_bytes).map_err(|e| DbError::Message(e.to_string()))?;
            out.push(meta);
        }
    }
    Ok(out)
}

#[derive(Clone, Default)]
pub struct InMemoryEngine {
    inner: Arc<Mutex<HashMap<String, HashMap<String, Bytes>>>>,
}

impl InMemoryEngine {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

#[async_trait(?Send)]
impl Engine for InMemoryEngine {
    async fn put(&self, ns: &str, key: &str, value: Bytes) -> Result<()> {
        let mut guard = self.inner.lock().unwrap();
        let space = guard.entry(ns.to_string()).or_insert_with(HashMap::new);
        space.insert(key.to_string(), value);
        Ok(())
    }

    async fn get(&self, ns: &str, key: &str) -> Result<Option<Bytes>> {
        let guard = self.inner.lock().unwrap();
        let Some(space) = guard.get(ns) else {
            return Ok(None);
        };
        Ok(space.get(key).cloned())
    }

    async fn delete(&self, ns: &str, key: &str) -> Result<()> {
        let mut guard = self.inner.lock().unwrap();
        if let Some(space) = guard.get_mut(ns) {
            space.remove(key);
        }
        Ok(())
    }

    async fn scan_prefix(&self, ns: &str, prefix: &str) -> Result<Vec<String>> {
        let guard = self.inner.lock().unwrap();
        let mut out = Vec::new();
        if let Some(space) = guard.get(ns) {
            for k in space.keys() {
                if k.starts_with(prefix) {
                    out.push(k.clone());
                }
            }
        }
        Ok(out)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::executor::block_on;

    #[test]
    fn kv_put_get_roundtrip() {
        let engine = InMemoryEngine::new();
        block_on(async {
            put_kv(&engine, "a", b"1".to_vec()).await.unwrap();
            let got = get_kv(&engine, "a").await.unwrap().unwrap();
            assert_eq!(got, b"1".to_vec());
            delete_kv(&engine, "a").await.unwrap();
            let none = get_kv(&engine, "a").await.unwrap();
            assert!(none.is_none());
        });
    }

    #[test]
    fn image_meta_and_blob_roundtrip() {
        let engine = InMemoryEngine::new();
        block_on(async {
            let meta = ImageMeta {
                id: "img1".to_string(),
                content_type: "image/png".to_string(),
                width: 10,
                height: 20,
                size: 4,
                created_at_ms: 1,
            };
            put_image(&engine, &meta, vec![1, 2, 3, 4]).await.unwrap();
            let out = get_image(&engine, "img1").await.unwrap().unwrap();
            assert_eq!(out.0, meta);
            assert_eq!(out.1, vec![1, 2, 3, 4]);
            delete_image(&engine, "img1").await.unwrap();
            let none = get_image(&engine, "img1").await.unwrap();
            assert!(none.is_none());
        });
    }

    #[test]
    fn list_images_test() {
        let engine = InMemoryEngine::new();
        block_on(async {
            let meta1 = ImageMeta {
                id: "img1".to_string(),
                content_type: "image/png".to_string(),
                width: 10,
                height: 20,
                size: 4,
                created_at_ms: 1,
            };
            let meta2 = ImageMeta {
                id: "img2".to_string(),
                content_type: "image/jpeg".to_string(),
                width: 30,
                height: 40,
                size: 8,
                created_at_ms: 2,
            };
            put_image(&engine, &meta1, vec![1, 2, 3, 4]).await.unwrap();
            put_image(&engine, &meta2, vec![5, 6, 7, 8]).await.unwrap();
            let mut list = list_images(&engine).await.unwrap();
            list.sort_by(|a, b| a.id.cmp(&b.id));
            assert_eq!(list.len(), 2);
            assert_eq!(list[0], meta1);
            assert_eq!(list[1], meta2);
        });
    }
}