8use ::log::{debug, error, trace, warn};
9use rusqlite::{OptionalExtension, params};
10use sha2::{Digest, Sha256};
11use whiskers_ports::{
12    Audio, AudioMs, CacheSize, LineTiming, Millis, SpeechCache, SpeechLine, StoreError, TimedAudio, VoiceId, speech_weight,
13};

15use crate::Store;
16
17pub struct SqliteSpeechCache {
18    store: Store,
19    cap: u64,
20}
21
22fn key_of(voice: &VoiceId, line: &SpeechLine) -> Vec<u8> {
23    let mut h = Sha256::new();
24    h.update(voice.as_str().as_bytes());
25    h.update([0]);
26    h.update(line.as_str().as_bytes());
27    h.finalize().to_vec()
28}
29
30impl SqliteSpeechCache {

32    pub fn new(store: Store, cap: u64) -> Self {
33        Self { store, cap }
34    }

37    fn unavailable(what: &str, e: impl std::fmt::Display) -> StoreError {
38        warn!("the speech cache {what}: {e}");
39        StoreError::Unavailable
40    }
41}

43fn pack(speech: &TimedAudio) -> Vec<u8> {
44    let mut out = Vec::with_capacity(speech.timing().spans().len() * 8);
45    for s in speech.timing().spans() {
46        out.extend_from_slice(&s.start().get().to_le_bytes());
47        out.extend_from_slice(&s.end().get().to_le_bytes());
48    }
49    out
50}
51
52fn unpack(blob: &[u8]) -> Option<Vec<(AudioMs, AudioMs)>> {
53    if blob.len() % 8 != 0 {
54        return None;
55    }
56    Some(
57        blob.chunks_exact(8)
58            .map(|c| {
59                let word = |b: &[u8]| u32::from_le_bytes([b[0], b[1], b[2], b[3]]);
60                (AudioMs::new(word(&c[..4])), AudioMs::new(word(&c[4..])))
61            })
62            .collect(),
63    )
64}
65
66impl SpeechCache for SqliteSpeechCache {
67    async fn find(&self, voice: &VoiceId, line: &SpeechLine, _now: Millis) -> Result<Option<TimedAudio>, StoreError> {
68        let key = key_of(voice, line);
69        let found = self
70            .store
71            .transaction(|tx| {
72                let row: Option<(Vec<u8>, Vec<u8>)> =
73                    tx.query_row("SELECT audio, spans FROM speech_cache WHERE key = ?1", [&key], |r| Ok((r.get(0)?, r.get(1)?))).optional()?;
74                if row.is_some() {
75                    tx.execute("UPDATE speech_cache SET used = (SELECT COALESCE(MAX(used), 0) + 1 FROM speech_cache) WHERE key = ?1", [&key])?;
76                }
77                Ok(row)
78            })
79            .map_err(|e| Self::unavailable("cannot be read", e))?;
80        let Some((audio, spans)) = found else { return Ok(None) };
81        let rebuilt = unpack(&spans)
82            .and_then(|spans| LineTiming::new(line.as_str().to_owned(), spans).ok())
83            .and_then(|timing| TimedAudio::new(Audio(audio), timing, line).ok());
84        match rebuilt {
85            Some(t) => {
86                trace!("speech cache: found {} for {}, {} bytes", voice.tag(), line.chars(), t.audio().0.len());
87                Ok(Some(t))
88            }
89            None => {
90                // Spans that are not of this line: an entry that cannot be used is removed, not trusted.
91                error!("a speech cache entry is damaged ({} chars); removing it", line.chars());
92                self.store.transaction(|tx| Ok(tx.execute("DELETE FROM speech_cache WHERE key = ?1", [&key])?)).map_err(|e| Self::unavailable("cannot drop an entry", e))?;
93                Ok(None)
94            }
95        }
96    }
97
98    async fn keep(&self, voice: &VoiceId, line: &SpeechLine, speech: &TimedAudio, _now: Millis) -> Result<(), StoreError> {
99        let weight = speech_weight(speech);
100        if weight > self.cap {
101            debug!("speech cache: a line of {weight} bytes is over the whole cap of {}; not kept", self.cap);
102            return Ok(());
103        }
104        let key = key_of(voice, line);
105        let (audio, spans) = (&speech.audio().0, pack(speech));
106        let evicted = self
107            .store
108            .transaction(|tx| {
109                tx.execute(
110                    "INSERT INTO speech_cache (key, audio, spans, weight, used) VALUES (?1, ?2, ?3, ?4, (SELECT COALESCE(MAX(used), 0) + 1 FROM speech_cache))
111                     ON CONFLICT (key) DO UPDATE SET audio = excluded.audio, spans = excluded.spans, weight = excluded.weight, used = excluded.used",
112                    params![key, audio, spans, weight as i64],
113                )?;
114                // Oldest out until it fits; the line just kept has the highest `used`, so it goes last.
115                let mut evicted = 0;
116                loop {
117                    let total: i64 = tx.query_row("SELECT COALESCE(SUM(weight), 0) FROM speech_cache", [], |r| r.get(0))?;
118                    if total as u64 <= self.cap {
119                        break;
120                    }
121                    evicted += tx.execute("DELETE FROM speech_cache WHERE key = (SELECT key FROM speech_cache ORDER BY used LIMIT 1)", [])?;
122                }
123                Ok(evicted)
124            })
125            .map_err(|e| Self::unavailable("cannot keep a line", e))?;
126        debug!("speech cache: kept {} bytes for {}, {evicted} older lines dropped", weight, voice.tag());
127        Ok(())
128    }
129
130    async fn size(&self) -> Result<CacheSize, StoreError> {
131        let conn = self.store.lock();
132        conn.query_row("SELECT COUNT(*), COALESCE(SUM(weight), 0) FROM speech_cache", [], |r| Ok(CacheSize { entries: r.get::<_, i64>(0)? as u64, bytes: r.get::<_, i64>(1)? as u64 }))
133            .map_err(|e| Self::unavailable("cannot be measured", e))
134    }
135}