//! The per-file cache: a file's finished zip entry and report data, keyed by its bytes and every
//! setting that shapes its result, so an unchanged file is never optimized twice, in any pack.
//!
//! - The key is BLAKE3 over the cache format, Packopt's version, the file's settings fingerprint
//!   and its bytes. The path is not in it, so the same texture in two packs is one entry; the
//!   time budget is not either, because the only results it changes (timed-out ones) are never
//!   stored. A result is therefore exactly what a fresh run computes, which keeps the archive
//!   byte-identical whether it came from the cache or not.
//! - Entries are written to a temporary file and renamed, so racing runs and crashes never leave
//!   a half entry; each carries a BLAKE3 of its own bytes, and anything that fails to decode or
//!   check is a miss and is overwritten.
//! - The cache is best effort: no read, write or sweep failure ever fails a run.

use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};

use crate::zip::Method;

/// Bumped whenever the entry layout or the meaning of a key changes.
const FORMAT: u32 = 3;
const MAGIC: &[u8; 8] = b"PKOPTC\0\x03";
const LAYOUT: &str = "v3";
const PARTIAL: &str = ".partial";
pub const DEFAULT_MAX_BYTES: u64 = 2 << 30;
/// Partial files older than this are left by a crashed run and are swept.
const STALE_PARTIAL: Duration = Duration::from_secs(3600);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Key([u8; 32]);

impl Key {
    fn hex(&self) -> String {
        self.0.iter().map(|b| format!("{b:02x}")).collect()
    }
}

/// The key of a file with `fingerprint` (its effective settings, see `pipeline`) and `bytes`.
pub fn key(fingerprint: &str, bytes: &[u8]) -> Key {
    let mut h = blake3::Hasher::new();
    h.update(b"packopt-cache");
    h.update(&FORMAT.to_le_bytes());
    h.update(env!("CARGO_PKG_VERSION").as_bytes());
    h.update(&[0]);
    h.update(&(fingerprint.len() as u64).to_le_bytes());
    h.update(fingerprint.as_bytes());
    h.update(bytes);
    Key(*h.finalize().as_bytes())
}

#[derive(Debug, Clone, PartialEq)]
pub struct LossyRecord {
    pub lossless_bytes: u64,
    pub lossy_bytes: u64,
    pub ssim: f64,
    pub psnr: f64,
}

#[derive(Debug, Clone, PartialEq)]
pub struct AudioRecord {
    pub lossless_bytes: u64,
    pub lossy_bytes: u64,
    pub source_kbps: f64,
    pub downmixed: bool,
    pub resampled_to: Option<u32>,
}

/// One file's result, without anything that depends on its path: the warning names the file as
/// `{path}`.
#[derive(Debug, Clone, PartialEq)]
pub struct Cached {
    pub method: Method,
    pub crc: u32,
    pub size: u64,
    pub payload: Vec<u8>,
    pub lossy: Option<LossyRecord>,
    pub audio: Option<AudioRecord>,
    pub warning: Option<String>,
}

pub const PATH_TOKEN: &str = "{path}";

fn encode(c: &Cached) -> Vec<u8> {
    let mut out = Vec::with_capacity(c.payload.len() + 96);
    out.extend_from_slice(MAGIC);
    out.push(c.method as u8);
    out.extend_from_slice(&c.crc.to_le_bytes());
    out.extend_from_slice(&c.size.to_le_bytes());
    out.extend_from_slice(&(c.payload.len() as u64).to_le_bytes());
    out.extend_from_slice(&c.payload);
    match &c.lossy {
        Some(l) => {
            out.push(1);
            out.extend_from_slice(&l.lossless_bytes.to_le_bytes());
            out.extend_from_slice(&l.lossy_bytes.to_le_bytes());
            out.extend_from_slice(&l.ssim.to_bits().to_le_bytes());
            out.extend_from_slice(&l.psnr.to_bits().to_le_bytes());
        }
        None => out.push(0),
    }
    match &c.audio {
        Some(a) => {
            out.push(1);
            out.extend_from_slice(&a.lossless_bytes.to_le_bytes());
            out.extend_from_slice(&a.lossy_bytes.to_le_bytes());
            out.extend_from_slice(&a.source_kbps.to_bits().to_le_bytes());
            out.push(u8::from(a.downmixed));
            out.extend_from_slice(&a.resampled_to.unwrap_or(0).to_le_bytes());
        }
        None => out.push(0),
    }
    match &c.warning {
        Some(w) => {
            out.push(1);
            out.extend_from_slice(&(w.len() as u32).to_le_bytes());
            out.extend_from_slice(w.as_bytes());
        }
        None => out.push(0),
    }
    let sum = *blake3::hash(&out).as_bytes();
    out.extend_from_slice(&sum);
    out
}

struct Reader<'a>(&'a [u8]);

impl<'a> Reader<'a> {
    fn take(&mut self, n: usize) -> Option<&'a [u8]> {
        if self.0.len() < n {
            return None;
        }
        let (head, rest) = self.0.split_at(n);
        self.0 = rest;
        Some(head)
    }

    fn u8(&mut self) -> Option<u8> {
        Some(self.take(1)?[0])
    }

    fn u32(&mut self) -> Option<u32> {
        Some(u32::from_le_bytes(self.take(4)?.try_into().ok()?))
    }

    fn u64(&mut self) -> Option<u64> {
        Some(u64::from_le_bytes(self.take(8)?.try_into().ok()?))
    }
}

fn decode(bytes: &[u8]) -> Option<Cached> {
    let body_len = bytes.len().checked_sub(32)?;
    let (body, sum) = bytes.split_at(body_len);
    if blake3::hash(body).as_bytes() != sum {
        return None;
    }
    let mut r = Reader(body);
    if r.take(MAGIC.len())? != MAGIC {
        return None;
    }
    let method = match r.u8()? {
        0 => Method::Stored,
        8 => Method::Deflated,
        _ => return None,
    };
    let crc = r.u32()?;
    let size = r.u64()?;
    let payload_len = usize::try_from(r.u64()?).ok()?;
    let payload = r.take(payload_len)?.to_vec();
    let lossy = match r.u8()? {
        0 => None,
        1 => Some(LossyRecord {
            lossless_bytes: r.u64()?,
            lossy_bytes: r.u64()?,
            ssim: f64::from_bits(r.u64()?),
            psnr: f64::from_bits(r.u64()?),
        }),
        _ => return None,
    };
    let audio = match r.u8()? {
        0 => None,
        1 => Some(AudioRecord {
            lossless_bytes: r.u64()?,
            lossy_bytes: r.u64()?,
            source_kbps: f64::from_bits(r.u64()?),
            downmixed: match r.u8()? {
                0 => false,
                1 => true,
                _ => return None,
            },
            resampled_to: Some(r.u32()?).filter(|&rate| rate != 0),
        }),
        _ => return None,
    };
    let warning = match r.u8()? {
        0 => None,
        1 => {
            let len = r.u32()? as usize;
            Some(String::from_utf8(r.take(len)?.to_vec()).ok()?)
        }
        _ => return None,
    };
    if !r.0.is_empty() {
        return None;
    }
    Some(Cached { method, crc, size, payload, lossy, audio, warning })
}

static PARTIAL_COUNTER: AtomicU64 = AtomicU64::new(0);

pub struct Cache {
    root: PathBuf,
    max_bytes: u64,
}

impl Cache {
    pub fn new(dir: &Path, max_bytes: Option<u64>) -> Cache {
        Cache { root: dir.join(LAYOUT), max_bytes: max_bytes.unwrap_or(DEFAULT_MAX_BYTES) }
    }

    fn entry(&self, key: &Key) -> PathBuf {
        let hex = key.hex();
        self.root.join(&hex[..2]).join(hex)
    }

    /// The entry for `key`, or `None` when there is none or it does not check out. A hit
    /// refreshes the entry's age, so the sweep removes the least recently used first.
    pub fn get(&self, key: &Key) -> Option<Cached> {
        let path = self.entry(key);
        let cached = decode(&std::fs::read(&path).ok()?)?;
        if let Ok(file) = std::fs::OpenOptions::new().write(true).open(&path) {
            let _touched = file.set_modified(SystemTime::now());
        }
        Some(cached)
    }

    /// Stores `cached` under `key`. Failures are ignored: the cache only ever saves work.
    pub fn put(&self, key: &Key, cached: &Cached) {
        let path = self.entry(key);
        let Some(dir) = path.parent() else { return };
        if std::fs::create_dir_all(dir).is_err() {
            return;
        }
        let unique = PARTIAL_COUNTER.fetch_add(1, Ordering::Relaxed);
        let partial = dir.join(format!("{}.{}.{unique}{PARTIAL}", key.hex(), std::process::id()));
        if std::fs::write(&partial, encode(cached)).is_err() || std::fs::rename(&partial, &path).is_err() {
            let _removed = std::fs::remove_file(&partial);
        }
    }

    /// Every file under the cache with its size and age, oldest first (ties by path, so two
    /// sweeps of one cache agree).
    fn inventory(&self) -> Vec<(PathBuf, u64, SystemTime)> {
        let mut found = Vec::new();
        let Ok(shards) = std::fs::read_dir(&self.root) else { return found };
        for shard in shards.flatten() {
            let Ok(entries) = std::fs::read_dir(shard.path()) else { continue };
            for entry in entries.flatten() {
                let Ok(meta) = entry.metadata() else { continue };
                if meta.is_file() {
                    found.push((entry.path(), meta.len(), meta.modified().unwrap_or(SystemTime::UNIX_EPOCH)));
                }
            }
        }
        found.sort_by(|a, b| a.2.cmp(&b.2).then_with(|| a.0.cmp(&b.0)));
        found
    }

    /// Removes stale partial files, then the least recently used entries until the cache is at
    /// 90% of its cap. Returns the bytes removed.
    pub fn sweep(&self) -> u64 {
        let now = SystemTime::now();
        let mut removed = 0;
        let mut kept = Vec::new();
        for (path, len, modified) in self.inventory() {
            let partial = path.to_str().is_some_and(|p| p.ends_with(PARTIAL));
            let stale = now.duration_since(modified).is_ok_and(|age| age > STALE_PARTIAL);
            if partial {
                if stale && std::fs::remove_file(&path).is_ok() {
                    removed += len;
                }
                continue;
            }
            kept.push((path, len));
        }
        let mut total: u64 = kept.iter().map(|(_, len)| len).sum();
        if total <= self.max_bytes {
            return removed;
        }
        let target = self.max_bytes / 10 * 9;
        for (path, len) in kept {
            if total <= target {
                break;
            }
            if std::fs::remove_file(&path).is_ok() {
                total -= len;
                removed += len;
            }
        }
        removed
    }

    /// Removes every entry. Returns the bytes removed.
    pub fn clear(&self) -> u64 {
        let removed =
            self.inventory().iter().filter(|(p, _, _)| std::fs::remove_file(p).is_ok()).map(|(_, len, _)| len).sum();
        let _gone = std::fs::remove_dir_all(&self.root);
        removed
    }

    /// The number of entries and their total size.
    pub fn usage(&self) -> (usize, u64) {
        let entries = self.inventory();
        let bytes = entries.iter().map(|(_, len, _)| len).sum();
        (entries.len(), bytes)
    }
}

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

    fn sample() -> Cached {
        Cached {
            method: Method::Deflated,
            crc: 0xdead_beef,
            size: 1234,
            payload: b"compressed bytes".to_vec(),
            lossy: Some(LossyRecord { lossless_bytes: 900, lossy_bytes: 500, ssim: 0.991, psnr: 44.25 }),
            audio: Some(AudioRecord {
                lossless_bytes: 9000,
                lossy_bytes: 6000,
                source_kbps: 180.5,
                downmixed: true,
                resampled_to: Some(32_000),
            }),
            warning: Some(format!("{PATH_TOKEN}: not strict JSON")),
        }
    }

    #[test]
    fn entries_round_trip_and_reject_damage() {
        let bytes = encode(&sample());
        assert_eq!(decode(&bytes), Some(sample()));
        let plain = Cached { lossy: None, audio: None, warning: None, ..sample() };
        assert_eq!(decode(&encode(&plain)), Some(plain));
        for i in [0, 9, bytes.len() / 2, bytes.len() - 1] {
            let mut flipped = bytes.clone();
            flipped[i] ^= 1;
            assert_eq!(decode(&flipped), None, "flipped byte {i}");
        }
        assert_eq!(decode(&bytes[..bytes.len() - 1]), None);
        assert_eq!(decode(b""), None);
    }

    #[test]
    fn keys_depend_on_bytes_and_fingerprint() {
        assert_eq!(key("a", b"x"), key("a", b"x"));
        assert_ne!(key("a", b"x"), key("a", b"y"));
        assert_ne!(key("a", b"x"), key("b", b"x"));
        // The fingerprint's length is hashed, so no split of the same text collides.
        assert_ne!(key("ab", b"c"), key("a", b"bc"));
    }

    #[test]
    fn put_get_corrupt_and_sweep() {
        let dir = tempfile::tempdir().unwrap();
        let cache = Cache::new(dir.path(), Some(1));
        let k = key("f", b"bytes");
        assert_eq!(cache.get(&k), None);
        cache.put(&k, &sample());
        assert_eq!(cache.get(&k), Some(sample()));
        let path = cache.entry(&k);
        let mut bytes = std::fs::read(&path).unwrap();
        bytes[12] ^= 0xff;
        std::fs::write(&path, &bytes).unwrap();
        assert_eq!(cache.get(&k), None);
        cache.put(&k, &sample());
        assert_eq!(cache.get(&k), Some(sample()));
        assert_eq!(cache.usage().0, 1);
        assert!(cache.sweep() > 0);
        assert_eq!(cache.usage(), (0, 0));
    }

    #[test]
    fn the_sweep_removes_the_oldest_first() {
        let dir = tempfile::tempdir().unwrap();
        let entry_size = encode(&sample()).len() as u64;
        let cache = Cache::new(dir.path(), Some(entry_size * 3));
        let keys: Vec<Key> = (0..4u8).map(|i| key("f", &[i])).collect();
        for (i, k) in keys.iter().enumerate() {
            cache.put(k, &sample());
            let file = std::fs::OpenOptions::new().write(true).open(cache.entry(k)).unwrap();
            file.set_modified(SystemTime::UNIX_EPOCH + Duration::from_secs(1000 + i as u64)).unwrap();
        }
        cache.sweep();
        assert_eq!(cache.get(&keys[0]), None);
        assert_eq!(cache.get(&keys[1]), None);
        assert!(cache.get(&keys[2]).is_some());
        assert!(cache.get(&keys[3]).is_some());
    }
}