/*
    This file is part of jellything (https://codeberg.org/metamuffin/jellything)
    which is licensed under the GNU Affero General Public License (version 3); see /COPYING.
    Copyright (C) 2026 metamuffin <metamuffin.org>
*/
mod helper;

use anyhow::{Context, Result, anyhow};
pub use helper::{EscapeKey, HashKey};
use jellykv::BlobStorage;
use log::{info, warn};
use serde::{Deserialize, Serialize};
use std::{
    any::Any,
    collections::{BTreeMap, HashMap},
    hash::{DefaultHasher, Hash, Hasher},
    sync::{
        Arc, LazyLock, Mutex, RwLock,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    time::Instant,
};

const CACHE_GENERATION_BUCKET_COUNT: usize = 1024;
pub static CACHE_GENERATION_LOCKS: LazyLock<[Mutex<()>; CACHE_GENERATION_BUCKET_COUNT]> =
    LazyLock::new(|| [(); CACHE_GENERATION_BUCKET_COUNT].map(|_| Mutex::new(())));

thread_local! { pub static WITHIN_CACHE_FILE: AtomicBool = const { AtomicBool::new(false) }; }

pub struct Cache {
    storage: Box<dyn BlobStorage>,
    memory_cache: RwLock<HashMap<String, InMemoryCacheEntry>>,
    memory_cache_size: AtomicUsize,
    max_memory_cache_size: usize,
}

pub struct InMemoryCacheEntry {
    size: usize,
    last_access: Instant,
    object: Arc<dyn Any + Send + Sync + 'static>,
}

fn bucket(key: &str) -> usize {
    let mut h = DefaultHasher::new();
    key.hash(&mut h);
    h.finish() as usize % CACHE_GENERATION_BUCKET_COUNT
}

impl Cache {
    pub fn new(storage: Box<dyn BlobStorage>, max_memory_cache_size: usize) -> Self {
        Self {
            max_memory_cache_size,
            storage,
            memory_cache: HashMap::new().into(),
            memory_cache_size: AtomicUsize::new(0),
        }
    }
    pub fn cache(&self, key: &str, generate: impl FnOnce() -> Result<Vec<u8>>) -> Result<Vec<u8>> {
        // we need a lock even if it exists since somebody might be still in the process of writing.
        let already_within = WITHIN_CACHE_FILE.with(|a| a.swap(true, Ordering::Relaxed));
        let _guard = if already_within {
            // TODO stupid hack to avoid deadlock for nested calls; not locking is fine but might cause double-generating
            CACHE_GENERATION_LOCKS[bucket(key)].try_lock().ok()
        } else {
            CACHE_GENERATION_LOCKS[bucket(key)].lock().ok()
        };

        let out = match self.storage.read(key)? {
            Some(x) => x,
            None => {
                let value = generate()?;
                self.storage.store(key, &value)?;
                value
            }
        };

        if !already_within {
            WITHIN_CACHE_FILE.with(|a| a.swap(false, Ordering::Relaxed));
        }
        drop(_guard);
        Ok(out)
    }

    pub fn read(&self, key: &str) -> Result<Option<Vec<u8>>> {
        self.storage.read(key)
    }
    pub fn store(&self, key: String, generate: impl FnOnce() -> Result<Vec<u8>>) -> Result<String> {
        self.cache(&key, generate)?;
        Ok(key)
    }

    pub fn cache_memory<Fun, T>(
        &self,
        key: &str,
        mut generate: Fun,
    ) -> Result<Arc<T>, anyhow::Error>
    where
        Fun: FnMut() -> Result<T, anyhow::Error>,
        T: Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static,
    {
        if !key.ends_with(".json") {
            warn!("cache_memory key not ending in .json: {key:?}")
        }

        {
            let mut g = self.memory_cache.write().unwrap();
            if let Some(entry) = g.get_mut(key) {
                entry.last_access = Instant::now();
                let object = entry
                    .object
                    .clone()
                    .downcast::<T>()
                    .map_err(|_| anyhow!("inconsistent types for in-memory cache"))?;
                return Ok(object);
            }
        }

        let data = self.cache(key, move || {
            let object = generate()?;
            Ok(serde_json::to_vec(&object)?)
        })?;
        let size = data.len();
        let object = serde_json::from_slice::<T>(&data).context("decoding cache object")?;
        let object = Arc::new(object);

        {
            let mut g = self.memory_cache.write().unwrap();
            g.insert(
                key.to_owned(),
                InMemoryCacheEntry {
                    size,
                    last_access: Instant::now(),
                    object: object.clone(),
                },
            );
            self.memory_cache_size.fetch_add(size, Ordering::Relaxed);
        }

        self.cleanup_cache();

        Ok(object)
    }

    fn cleanup_cache(&self) {
        let current_size = self.memory_cache_size.load(Ordering::Relaxed);
        if current_size < self.max_memory_cache_size {
            return;
        }
        info!("running cache eviction");
        let mut g = self.memory_cache.write().unwrap();

        // TODO: if two entries have *exactly* the same size, only one of the will be remove; this is fine for now
        let mut k = BTreeMap::new();
        for (loc, entry) in g.iter() {
            k.insert(entry.last_access.elapsed(), (loc.to_owned(), entry.size));
        }
        let mut reduction = 0;
        for (loc, size) in k.values().rev().take(k.len().div_ceil(2)) {
            g.remove(loc);
            reduction += size;
        }
        self.memory_cache_size
            .fetch_sub(reduction, Ordering::Relaxed);
        drop(g);

        info!(
            "done, {} freed",
            humansize::format_size(reduction, humansize::DECIMAL)
        );
    }
}