use std::{
    collections::BTreeMap,
    sync::Arc,
    time::{Duration, Instant},
};

use async_trait::async_trait;
use tokio::sync::RwLock;

use crate::cache::{CacheKey, CacheResult, CacheStore};

#[derive(Debug, Clone)]
struct Entry {
    value: Vec<u8>,
    expires_at: Option<Instant>,
}

/// In-memory cache backend for tests and local services.
#[derive(Debug, Clone, Default)]
pub struct MemoryCacheStore {
    entries: Arc<RwLock<BTreeMap<String, Entry>>>,
}

impl MemoryCacheStore {
    /// Creates an empty in-memory store.
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl CacheStore for MemoryCacheStore {
    async fn get_raw(&self, key: &CacheKey) -> CacheResult<Option<Vec<u8>>> {
        let rendered = key.render();
        let mut entries = self.entries.write().await;
        let Some(entry) = entries.get(&rendered) else {
            return Ok(None);
        };
        if entry
            .expires_at
            .is_some_and(|deadline| deadline <= Instant::now())
        {
            entries.remove(&rendered);
            return Ok(None);
        }
        Ok(Some(entry.value.clone()))
    }

    async fn set_raw(
        &self,
        key: &CacheKey,
        value: Vec<u8>,
        ttl: Option<Duration>,
    ) -> CacheResult<()> {
        let expires_at = ttl.map(|ttl| Instant::now() + ttl);
        self.entries
            .write()
            .await
            .insert(key.render(), Entry { value, expires_at });
        Ok(())
    }

    async fn delete(&self, key: &CacheKey) -> CacheResult<()> {
        self.entries.write().await.remove(&key.render());
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use crate::cache::{CacheKey, CacheStore, MemoryCacheStore};

    #[tokio::test]
    async fn memory_cache_stores_json_with_ttl() {
        let store = MemoryCacheStore::new();
        let key = CacheKey::new("test", ["user", "1"]);
        store
            .set_json(
                &key,
                &serde_json::json!({"name":"Ada"}),
                Some(Duration::from_secs(1)),
            )
            .await
            .expect("set");
        let value: serde_json::Value = store.get_json(&key).await.expect("get").expect("value");
        assert_eq!(value["name"], "Ada");
        store.delete(&key).await.expect("delete");
        assert!(store.get_raw(&key).await.expect("get").is_none());
    }
}
