use sha2::{Digest, Sha256};
use std::cmp::Ordering;
use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX";
const INDEX_VERSION: u32 = 1;
const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024;
const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum StoreReason {
    Cold,
    Continued,
}

impl StoreReason {
    fn code(self) -> u8 {
        match self {
            Self::Cold => 1,
            Self::Continued => 2,
        }
    }

    fn from_code(value: u8) -> Option<Self> {
        Some(match value {
            1 => Self::Cold,
            2 => Self::Continued,
            _ => return None,
        })
    }

    fn is_anchor(self) -> bool {
        self == Self::Cold
    }
}

#[derive(Clone, Debug)]
pub(super) struct Entry {
    pub(super) checkpoint: PathBuf,
    pub(super) tag: [u8; 32],
    pub(super) tokens: u32,
    context: u32,
    hits: u32,
    created_at: u64,
    last_used: u64,
    reason: StoreReason,
    key: Vec<u8>,
    bytes: u64,
}

pub(super) struct KvStore {
    directory: PathBuf,
    budget_bytes: u64,
}

impl KvStore {
    pub(super) fn open(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
        fs::create_dir_all(directory).map_err(|error| {
            format!(
                "Could not create transient KV cache directory {}: {error}",
                directory.display()
            )
        })?;
        let store = Self {
            directory: directory.to_owned(),
            budget_bytes,
        };
        // Reclaim first, so the budget is measured against the whole directory.
        sweep_unreachable(directory);
        // A lowered budget takes effect at once, like ds4's eviction on open.
        store.evict(None);
        Ok(store)
    }

    pub(super) fn find(&self, prompt_key: &[u8], context: u32) -> Option<Entry> {
        self.entries()
            .into_iter()
            .filter(|entry| entry.context == context && prompt_key.starts_with(&entry.key))
            .max_by(|left, right| {
                left.key
                    .len()
                    .cmp(&right.key.len())
                    .then_with(|| left.tokens.cmp(&right.tokens))
            })
    }

    pub(super) fn checkpoint_path(&self, key: &[u8]) -> PathBuf {
        self.directory.join(format!("{}.bin", key_hash(key)))
    }

    pub(super) fn record(
        &self,
        checkpoint: &Path,
        key: &[u8],
        tag: [u8; 32],
        tokens: u32,
        context: u32,
        reason: StoreReason,
    ) -> Result<bool, String> {
        let checkpoint_bytes = fs::metadata(checkpoint)
            .map_err(|error| error.to_string())?
            .len();
        let index_bytes = index_size(key)?;
        if checkpoint_bytes.saturating_add(index_bytes) > self.budget_bytes {
            let _ = fs::remove_file(checkpoint);
            return Ok(false);
        }
        let now = unix_time();
        let entry = Entry {
            checkpoint: checkpoint.to_owned(),
            tag,
            tokens,
            context,
            hits: 0,
            created_at: now,
            last_used: now,
            reason,
            key: key.to_vec(),
            bytes: checkpoint_bytes.saturating_add(index_bytes),
        };
        write_entry(&entry)?;
        self.evict(Some(checkpoint));
        Ok(checkpoint.exists())
    }

    pub(super) fn touch(&self, entry: &Entry) -> Result<(), String> {
        let mut touched = entry.clone();
        touched.hits = touched.hits.saturating_add(1);
        touched.last_used = unix_time();
        write_entry(&touched)
    }

    pub(super) fn discard(&self, entry: &Entry) {
        let _ = fs::remove_file(&entry.checkpoint);
        let _ = fs::remove_file(index_path(&entry.checkpoint));
    }

    fn entries(&self) -> Vec<Entry> {
        let Ok(files) = fs::read_dir(&self.directory) else {
            return Vec::new();
        };
        files
            .filter_map(Result::ok)
            .map(|file| file.path())
            .filter(|path| path.extension().is_some_and(|value| value == "meta"))
            .filter_map(|path| read_entry(&path).ok())
            .filter(|entry| entry.checkpoint.is_file())
            .collect()
    }

    fn evict(&self, protected: Option<&Path>) {
        let mut entries = self.entries();
        let mut total = entries.iter().map(|entry| entry.bytes).sum::<u64>();
        let now = unix_time();
        while total > self.budget_bytes {
            let Some((index, _)) = entries
                .iter()
                .enumerate()
                .filter(|(_, entry)| protected != Some(entry.checkpoint.as_path()))
                .min_by(|(_, left), (_, right)| {
                    eviction_score(left, now)
                        .partial_cmp(&eviction_score(right, now))
                        .unwrap_or(Ordering::Equal)
                        .then_with(|| left.last_used.cmp(&right.last_used))
                })
            else {
                break;
            };
            let entry = entries.swap_remove(index);
            total = total.saturating_sub(entry.bytes);
            self.discard(&entry);
        }
    }
}

/// Deletes what the index cannot reach: a checkpoint whose index is missing or
/// unreadable can never be found, and an index whose checkpoint is gone is
/// dead weight. Without this, an interrupted store or a checkpoint written
/// before the index existed occupies disc that eviction cannot even see, and
/// the directory grows past the budget unnoticed.
pub(crate) fn sweep_unreachable(directory: &Path) {
    let Ok(files) = fs::read_dir(directory) else {
        return;
    };
    let reachable = files
        .filter_map(Result::ok)
        .map(|file| file.path())
        .filter(|path| path.extension().is_some_and(|value| value == "meta"))
        .filter_map(|path| read_entry(&path).ok())
        .filter(|entry| entry.checkpoint.is_file())
        .map(|entry| entry.checkpoint)
        .collect::<HashSet<_>>();
    let Ok(files) = fs::read_dir(directory) else {
        return;
    };
    for file in files.flatten() {
        let path = file.path();
        let dead = match path.extension().and_then(|value| value.to_str()) {
            Some("bin") => !reachable.contains(&path),
            Some("meta") => !reachable.contains(&path.with_extension("bin")),
            // Staging files from an interrupted index write.
            Some("tmp") => true,
            _ => false,
        };
        if dead {
            let _ = fs::remove_file(&path);
        }
    }
}

fn eviction_score(entry: &Entry, now: u64) -> f64 {
    let age = now.saturating_sub(entry.last_used.max(entry.created_at)) as f64;
    let hits = f64::from(entry.hits) * f64::exp2(-age / HIT_HALF_LIFE_SECONDS);
    let anchor = if entry.reason.is_anchor() { 2.0 } else { 1.0 };
    anchor * (hits + 1.0) * f64::from(entry.tokens) / entry.bytes.max(1) as f64
}

fn write_entry(entry: &Entry) -> Result<(), String> {
    let path = index_path(&entry.checkpoint);
    let temporary = path.with_extension("meta.tmp");
    let mut file = File::create(&temporary).map_err(|error| error.to_string())?;
    file.write_all(INDEX_MAGIC)
        .map_err(|error| error.to_string())?;
    write_u32(&mut file, INDEX_VERSION)?;
    file.write_all(&entry.tag)
        .map_err(|error| error.to_string())?;
    write_u32(&mut file, entry.tokens)?;
    write_u32(&mut file, entry.context)?;
    write_u32(&mut file, entry.hits)?;
    file.write_all(&[entry.reason.code(), 0, 0, 0])
        .map_err(|error| error.to_string())?;
    write_u64(&mut file, entry.created_at)?;
    write_u64(&mut file, entry.last_used)?;
    write_u64(
        &mut file,
        u64::try_from(entry.key.len()).map_err(|_| "KV cache key is too large")?,
    )?;
    file.write_all(&entry.key)
        .map_err(|error| error.to_string())?;
    file.sync_all().map_err(|error| error.to_string())?;
    fs::rename(temporary, path).map_err(|error| error.to_string())
}

fn read_entry(path: &Path) -> Result<Entry, String> {
    let mut file = File::open(path).map_err(|error| error.to_string())?;
    let mut magic = [0; 8];
    file.read_exact(&mut magic)
        .map_err(|error| error.to_string())?;
    if &magic != INDEX_MAGIC || read_u32(&mut file)? != INDEX_VERSION {
        return Err("KV cache index has an invalid signature".into());
    }
    let mut tag = [0; 32];
    file.read_exact(&mut tag)
        .map_err(|error| error.to_string())?;
    let tokens = read_u32(&mut file)?;
    let context = read_u32(&mut file)?;
    let hits = read_u32(&mut file)?;
    let mut reason = [0; 4];
    file.read_exact(&mut reason)
        .map_err(|error| error.to_string())?;
    let reason = StoreReason::from_code(reason[0])
        .ok_or_else(|| "KV cache index has an invalid reason".to_owned())?;
    let created_at = read_u64(&mut file)?;
    let last_used = read_u64(&mut file)?;
    let key_bytes = read_u64(&mut file)?;
    if tokens == 0 || context == 0 || key_bytes > MAX_KEY_BYTES {
        return Err("KV cache index has invalid dimensions".into());
    }
    let mut key = vec![0; usize::try_from(key_bytes).map_err(|error| error.to_string())?];
    file.read_exact(&mut key)
        .map_err(|error| error.to_string())?;
    let checkpoint = path.with_extension("bin");
    if checkpoint.file_stem().and_then(|value| value.to_str()) != Some(&key_hash(&key)) {
        return Err("KV cache index key does not match its filename".into());
    }
    let bytes = fs::metadata(&checkpoint)
        .map_err(|error| error.to_string())?
        .len()
        .saturating_add(fs::metadata(path).map_err(|error| error.to_string())?.len());
    Ok(Entry {
        checkpoint,
        tag,
        tokens,
        context,
        hits,
        created_at,
        last_used,
        reason,
        key,
        bytes,
    })
}

fn index_path(checkpoint: &Path) -> PathBuf {
    checkpoint.with_extension("meta")
}

fn index_size(key: &[u8]) -> Result<u64, String> {
    84_u64
        .checked_add(u64::try_from(key.len()).map_err(|error| error.to_string())?)
        .ok_or_else(|| "KV cache index is too large".to_owned())
}

fn key_hash(key: &[u8]) -> String {
    let digest: [u8; 32] = Sha256::digest(key).into();
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut output = String::with_capacity(64);
    for byte in digest {
        output.push(HEX[(byte >> 4) as usize] as char);
        output.push(HEX[(byte & 0x0f) as usize] as char);
    }
    output
}

fn unix_time() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |duration| duration.as_secs())
}

fn write_u32(file: &mut File, value: u32) -> Result<(), String> {
    file.write_all(&value.to_le_bytes())
        .map_err(|error| error.to_string())
}

fn write_u64(file: &mut File, value: u64) -> Result<(), String> {
    file.write_all(&value.to_le_bytes())
        .map_err(|error| error.to_string())
}

fn read_u32(file: &mut File) -> Result<u32, String> {
    let mut bytes = [0; 4];
    file.read_exact(&mut bytes)
        .map_err(|error| error.to_string())?;
    Ok(u32::from_le_bytes(bytes))
}

fn read_u64(file: &mut File) -> Result<u64, String> {
    let mut bytes = [0; 8];
    file.read_exact(&mut bytes)
        .map_err(|error| error.to_string())?;
    Ok(u64::from_le_bytes(bytes))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::settings::GIB;

    fn temporary_directory(name: &str) -> PathBuf {
        let path = std::env::temp_dir().join(format!(
            "ds4-server-kvstore-{name}-{}-{}",
            std::process::id(),
            unix_time()
        ));
        fs::create_dir_all(&path).unwrap();
        path
    }

    fn add(store: &KvStore, key: &[u8], tokens: u32, bytes: usize) -> Entry {
        let path = store.checkpoint_path(key);
        fs::write(&path, vec![0; bytes]).unwrap();
        assert!(
            store
                .record(
                    &path,
                    key,
                    [tokens as u8; 32],
                    tokens,
                    4096,
                    StoreReason::Continued,
                )
                .unwrap()
        );
        store.find(key, 4096).unwrap()
    }

    #[test]
    fn longest_compatible_prefix_wins_and_hits_survive_reopen() {
        let directory = temporary_directory("prefix");
        let store = KvStore::open(&directory, GIB).unwrap();
        add(&store, b"system", 100, 32);
        let expected = add(&store, b"system-user-one", 200, 32);

        let found = store.find(b"system-user-one-more", 4096).unwrap();
        assert_eq!(found.checkpoint, expected.checkpoint);
        store.touch(&found).unwrap();
        assert_eq!(
            KvStore::open(&directory, GIB)
                .unwrap()
                .find(b"system-user-one-more", 4096)
                .unwrap()
                .hits,
            1
        );
        assert!(store.find(b"system-user-one-more", 2048).is_none());

        fs::remove_dir_all(directory).unwrap();
    }

    #[test]
    fn opening_reclaims_checkpoints_the_index_cannot_reach() {
        let directory = temporary_directory("unreachable");
        let store = KvStore::open(&directory, GIB).unwrap();
        let indexed = add(&store, b"indexed", 100, 32);
        let stray = directory.join("deadbeef.bin");
        let widow = directory.join("widow.meta");
        let staging = directory.join("indexed.meta.tmp");
        fs::write(&stray, vec![0; 4_096]).unwrap();
        fs::write(&widow, vec![0; 64]).unwrap();
        fs::write(&staging, vec![0; 64]).unwrap();

        KvStore::open(&directory, GIB).unwrap();

        assert!(indexed.checkpoint.is_file());
        assert!(index_path(&indexed.checkpoint).is_file());
        assert!(!stray.exists());
        assert!(!widow.exists());
        assert!(!staging.exists());

        fs::remove_dir_all(directory).unwrap();
    }

    #[test]
    fn budget_evicts_the_least_valuable_unprotected_entry() {
        let directory = temporary_directory("evict");
        let store = KvStore::open(&directory, 300).unwrap();
        let first = add(&store, b"first", 10, 32);
        let second = add(&store, b"second", 100, 32);
        let third = add(&store, b"third", 200, 32);

        assert!(!first.checkpoint.exists());
        assert!(second.checkpoint.exists());
        assert!(third.checkpoint.exists());

        fs::remove_dir_all(directory).unwrap();
    }
}