/*
    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) 2025 metamuffin <metamuffin.org>
*/
pub mod backends;
pub mod key;

use crate::backends::{CacheStorage, filesystem::Filesystem};
use anyhow::{Context, Result, anyhow};
pub use key::*;
use log::info;
use serde::{Deserialize, Serialize};
use std::{
    any::Any,
    collections::{BTreeMap, HashMap},
    path::PathBuf,
    sync::{
        Arc, LazyLock, Mutex, OnceLock, RwLock,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    time::Instant,
};

#[derive(Debug, Deserialize)]
pub struct Config {
    path: PathBuf,
    max_in_memory_cache_size: usize,
    secret: String,
}

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 static CONF_PRELOAD: std::sync::Mutex<Option<Config>> = std::sync::Mutex::new(None);
static CONF: LazyLock<Config> = LazyLock::new(|| {
    CONF_PRELOAD
        .lock()
        .unwrap()
        .take()
        .expect("cache config not preloaded. logic error")
});

static CACHE_STORE: OnceLock<Box<dyn CacheStorage>> = OnceLock::new();

pub fn init_cache() -> Result<()> {
    CACHE_STORE
        .set(Box::new(Filesystem::new(&CONF)))
        .map_err(|_| ())
        .unwrap();
    Ok(())
}

pub fn cache(key: CacheKey, 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[key.bucket()].try_lock().ok()
    } else {
        CACHE_GENERATION_LOCKS[key.bucket()].lock().ok()
    };

    let store = CACHE_STORE.get().unwrap();

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

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

pub fn cache_read(key: CacheKey) -> Result<Option<Vec<u8>>> {
    CACHE_STORE.get().unwrap().read(key)
}
pub fn cache_store(key: CacheKey, generate: impl FnOnce() -> Result<Vec<u8>>) -> Result<CacheKey> {
    cache(key, generate)?;
    Ok(key)
}

pub struct InMemoryCacheEntry {
    size: usize,
    last_access: Instant,
    object: Arc<dyn Any + Send + Sync + 'static>,
}
pub static CACHE_IN_MEMORY_OBJECTS: LazyLock<RwLock<HashMap<CacheKey, InMemoryCacheEntry>>> =
    LazyLock::new(|| RwLock::new(HashMap::new()));
pub static CACHE_IN_MEMORY_SIZE: AtomicUsize = AtomicUsize::new(0);

pub fn cache_memory<Fun, T>(key: CacheKey, mut generate: Fun) -> Result<Arc<T>, anyhow::Error>
where
    Fun: FnMut() -> Result<T, anyhow::Error>,
    T: Serialize + for<'de> Deserialize<'de> + Send + Sync + 'static,
{
    {
        let mut g = CACHE_IN_MEMORY_OBJECTS.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 = 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 = CACHE_IN_MEMORY_OBJECTS.write().unwrap();
        g.insert(
            key,
            InMemoryCacheEntry {
                size,
                last_access: Instant::now(),
                object: object.clone(),
            },
        );
        CACHE_IN_MEMORY_SIZE.fetch_add(size, Ordering::Relaxed);
    }

    cleanup_cache();

    Ok(object)
}

pub fn cleanup_cache() {
    let current_size = CACHE_IN_MEMORY_SIZE.load(Ordering::Relaxed);
    if current_size < CONF.max_in_memory_cache_size {
        return;
    }
    info!("running cache eviction");
    let mut g = CACHE_IN_MEMORY_OBJECTS.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;
    }
    CACHE_IN_MEMORY_SIZE.fetch_sub(reduction, Ordering::Relaxed);
    drop(g);

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