autocache 0.5.0

use std::sync::Arc;

use anyhow::Result;
use arc_swap::ArcSwapOption;
use bytes::Bytes;
use redis::AsyncCommands;
use serde::{de::DeserializeOwned, Serialize};

use crate::{cache::Cache, Codec, Entry};

pub struct RedisCache<K, V> {
    namespace: ArcSwapOption<String>,

    ttl_sec: usize,
    redis_cli: redis::Client,

    _marker: std::marker::PhantomData<fn() -> (K, V)>,
}

impl<K, V> RedisCache<K, V> {
    pub fn new(cli: redis::Client) -> Self {
        Self::new_with_ttl(cli, 0)
    }
    pub fn new_with_ttl(cli: redis::Client, ttl_sec: usize) -> Self {
        Self {
            namespace: None.into(),
            ttl_sec,

            redis_cli: cli,
            _marker: std::marker::PhantomData,
        }
    }
}

impl<K, V> RedisCache<K, V>
where
    K: Sync + AsRef<str>,
{
    fn generate_redis_key(&self, k: &K) -> String {
        if let Some(ns) = self.namespace.load().as_ref() {
            if ns.is_empty() {
                return k.as_ref().to_string();
            }

            let mut key = String::new();
            key.push_str(ns);
            key.push(':');

            key.push_str(k.as_ref());

            key
        } else {
            k.as_ref().to_string()
        }
    }
}

impl<K, V> Cache for RedisCache<K, V>
where
    K: Clone + Sync + AsRef<str> + Serialize + DeserializeOwned,
    V: Sync + Codec,
{
    type Key = K;
    type Value = V;

    async fn mget(&self, keys: &[Self::Key]) -> Result<Vec<Entry<Self::Key, Self::Value>>> {
        if keys.is_empty() {
            return Ok(vec![]);
        }

        let mut conn = self.redis_cli.get_multiplexed_async_connection().await?;

        if let [key] = keys {
            let data: Option<Bytes> = conn.get(self.generate_redis_key(key)).await?;

            return match data {
                Some(data) => Ok(vec![Entry::<K, V>::decode(data)?]),
                None => Ok(vec![]),
            };
        }

        let res: Vec<Option<Bytes>> = conn
            .mget(
                keys.iter()
                    .map(|key| self.generate_redis_key(key))
                    .collect::<Vec<_>>(),
            )
            .await?;

        res.into_iter()
            .flatten()
            .map(|data| Entry::<K, V>::decode(data).map_err(anyhow::Error::from))
            .collect()
    }

    async fn mset(&self, entries: &[Entry<Self::Key, Self::Value>]) -> Result<()> {
        if entries.is_empty() {
            return Ok(());
        }

        let encoded_kvs = entries
            .iter()
            .map(|entry| {
                Ok((
                    self.generate_redis_key(&entry.key),
                    entry.encode()?.to_vec(),
                ))
            })
            .collect::<Result<Vec<_>>>()?;

        let mut conn = self.redis_cli.get_multiplexed_async_connection().await?;

        if let [(key, value)] = encoded_kvs.as_slice() {
            if self.ttl_sec == 0 {
                conn.set::<_, _, ()>(key, value).await?;
            } else {
                conn.set_ex::<_, _, ()>(key, value, self.ttl_sec.try_into()?)
                    .await?;
            }
            return Ok(());
        }

        if self.ttl_sec == 0 {
            conn.mset::<_, _, ()>(&encoded_kvs).await?;
        } else {
            let mut pipe = redis::Pipeline::new();
            pipe.mset(&encoded_kvs);
            let ttl_sec = self.ttl_sec.try_into()?;
            for (key, _) in &encoded_kvs {
                pipe.expire(key, ttl_sec);
            }

            pipe.query_async::<()>(&mut conn).await?;
        }

        Ok(())
    }

    async fn mdel(&self, keys: &[Self::Key]) -> Result<()> {
        if keys.is_empty() {
            return Ok(());
        }

        let mut conn = self.redis_cli.get_multiplexed_async_connection().await?;
        conn.del::<_, ()>(
            keys.iter()
                .map(|key| self.generate_redis_key(key))
                .collect::<Vec<_>>(),
        )
        .await?;

        Ok(())
    }

    fn name(&self) -> &'static str {
        "rediscache"
    }

    fn set_ns(&self, ns: String) {
        self.namespace.store(Some(Arc::new(ns)));
    }
}