7use std::collections::{BTreeMap, BTreeSet};

9use ::log::{debug, error, info, trace, warn};
10use rusqlite::{Connection, OptionalExtension, Transaction, params};
11use whiskers_core::{
12    Cover, Fact, IconId, Kind, Memory, MemoryDoc, MemoryError, MemorySnapshot, Mention, NewFact, PictureId, Revision, SoftAction, SoftKind, Tombstone,
13    UndoRefusal, Visibility, new_gid,
14};
15
16use crate::{DbError, Store};

19pub struct SqliteMemory {
20    store: Store,

22    device: String,
23    doc: MemoryDoc,
24}

26impl SqliteMemory {

29    pub fn open(store: Store, device: &str) -> Result<Self, MemoryError> {
30        let doc = {
31            let conn = store.lock();
32            load_doc(&conn).map_err(|e| {
33                error!("the memory cannot be read: {e}");
34                MemoryError(format!("memory is damaged: {e}"))
35            })?
36        };
37        debug!("memory opened: {} facts, {} forgotten", doc.facts().len(), doc.tombstones().len());
38        Ok(Self { store, device: device.to_owned(), doc })
39    }

42    fn change<T>(&mut self, what: &str, op: impl FnOnce(&mut MemoryDoc, &str) -> Option<T>) -> Result<T, MemoryError> {
43        let before = self.doc.clone();
44        let Some(out) = op(&mut self.doc, &self.device) else {
45            warn!("{what}: no such fact, or it has no such picture");
46            return Err(MemoryError("no such fact, or it has no such picture".into()));
47        };
48        self.keep(before, what)?;
49        Ok(out)
50    }

54    fn keep(&mut self, before: MemoryDoc, what: &str) -> Result<(), MemoryError> {
55        self.keep_with(before, what, |_| Ok(()))
56    }

59    fn keep_with(&mut self, before: MemoryDoc, what: &str, also: impl FnOnce(&Transaction<'_>) -> Result<(), DbError>) -> Result<(), MemoryError> {
60        if self.doc == before {
61            return Ok(());
62        }
63        let forgets = self.doc.tombstones() != before.tombstones();
64        let wrote = self.store.transaction(|tx| {
65            write_changes(tx, &before, &self.doc)?;
66            also(tx)
67        });
68        if let Err(e) = wrote {
69            error!("{what} not kept, the write failed: {e}");
70            self.doc = before;
71            return Err(MemoryError(e.to_string()));
72        }
73        if forgets && let Err(e) = self.store.scrub() {
74            // The memory is gone from the database; what could not be done is emptying the log of it.
75            error!("{what}: the forgotten memory may still be in the write-ahead log: {e}");
76        }
77        Ok(())
78    }
79}

81impl Memory for SqliteMemory {
82    fn facts(&self) -> Vec<Fact> {
83        trace!("memory: {} facts listed", self.doc.facts().len());
84        self.doc.facts().to_vec()
85    }
86
87    fn add(&mut self, new: NewFact, at_ms: u64) -> Result<Fact, MemoryError> {
88        let gid = new_gid();
89        let fact = self.change("add", |d, device| Some(d.add(new, at_ms, device, gid)))?;
90        info!("fact {} added ({} chars); {} facts now", fact.id, fact.text.len(), self.doc.facts().len());
91        Ok(fact)
92    }
93
94    fn set_embedding(&mut self, id: u64, embedding: Vec<f32>) -> Result<(), MemoryError> {
95        debug!("fact {id} gets an embedding of {} numbers", embedding.len());
96        self.change("set an embedding on", |d, _| d.set_embedding(id, embedding)).map(|_| ())
97    }
98
99    fn mention(&mut self, id: u64, at_ms: u64, turn: &str) -> Result<Fact, MemoryError> {
100        self.change("note a mention of", |d, device| d.mention(id, at_ms, device, turn))
101    }
102
103    fn reword(&mut self, id: u64, text: String, at_ms: u64) -> Result<Fact, MemoryError> {
104        self.change("reword", |d, device| d.reword(id, text, at_ms, device))
105    }
106
107    fn forget(&mut self, id: u64, at_ms: u64) -> Result<Fact, MemoryError> {
108        let gone = self.change("forget", |d, device| d.forget(id, at_ms, device))?;
109        info!("fact {id} forgotten; {} facts left", self.doc.facts().len());
110        Ok(gone)
111    }
112
113    fn give_icon(&mut self, id: u64, icon: IconId) -> Result<Fact, MemoryError> {
114        self.change("give an icon to", |d, _| d.give_icon(id, icon))
115    }
116
117    fn choose_cover(&mut self, id: u64, picture: &PictureId, at_ms: u64) -> Result<Fact, MemoryError> {
118        self.change("choose the card's picture of", |d, _| d.choose_cover(id, picture, at_ms))
119    }
120
121    fn hide(&mut self, id: u64, at_ms: u64) -> Result<Fact, MemoryError> {
122        let before = self.doc.clone();
123        let Some(fact) = self.doc.hide(id, at_ms) else {
124            warn!("hide: no such fact");
125            return Err(MemoryError("no such fact".into()));
126        };
127        // Putting away what is away already changes nothing, and so writes no action.
128        self.keep_with(before, "hide", |tx| {
129            tx.execute("INSERT INTO soft_action (kind, gid, at_ms) VALUES (?1, ?2, ?3)", params![SoftKind::PutAway.word(), fact.gid, at_ms as i64])?;
130            Ok(())
131        })?;
132        debug!("fact {id} put away");
133        Ok(fact)
134    }
135
136    fn soft_actions(&self) -> Vec<SoftAction> {
137        let conn = self.store.lock();
138        read_actions(&conn, None).unwrap_or_else(|e| {
139            error!("the soft actions cannot be read: {e}");
140            Vec::new()
141        })
142    }
143
144    fn undo(&mut self, action_id: u64, now_ms: u64) -> Result<Fact, UndoRefusal> {
145        let action = {
146            let conn = self.store.lock();
147            read_actions(&conn, Some(action_id)).unwrap_or_else(|e| {
148                error!("undo: the soft action cannot be read: {e}");
149                Vec::new()
150            })
151        };
152        let Some(action) = action.into_iter().next() else {
153            info!("undo of action {action_id} refused: there is none");
154            return Err(UndoRefusal::Gone);
155        };
156        if let Err(refusal) = action.permits(now_ms) {
157            info!("undo of action {action_id} refused: {refusal}");
158            return Err(refusal);
159        }
160        let Some(id) = self.doc.facts().iter().find(|f| f.gid == action.gid).map(|f| f.id) else {
161            // The row cascades away with its memory, so this is a row the schema should not have left.
162            error!("undo of action {action_id}: its memory is not here");
163            return Err(UndoRefusal::Gone);
164        };
165        let before = self.doc.clone();
166        let fact = match action.kind {
167            SoftKind::PutAway => self.doc.restore(id, now_ms),
168        };
169        let Some(fact) = fact else {
170            return Err(UndoRefusal::Gone);
171        };
172        // The memory moves back and the action is spent in one transaction (a restore of a memory already shown
173        // changes no row, and still ends the action).
174        let wrote = self.store.transaction(|tx| {
175            write_changes(tx, &before, &self.doc)?;
176            tx.execute("UPDATE soft_action SET undone_at_ms = ?2 WHERE id = ?1", params![action_id as i64, now_ms as i64])?;
177            Ok(())
178        });
179        if let Err(e) = wrote {
180            error!("undo of action {action_id} not kept, the write failed: {e}");
181            self.doc = before;
182            return Err(UndoRefusal::Gone);
183        }
184        info!("action {action_id} undone: fact {id} is back");
185        Ok(fact)
186    }
187
188    fn restore(&mut self, id: u64, at_ms: u64) -> Result<Fact, MemoryError> {
189        self.change("restore", |d, _| d.restore(id, at_ms))
190    }
191
192    fn snapshot(&self) -> MemorySnapshot {
193        self.doc.snapshot()
194    }
195
196    fn merge(&mut self, remote: MemorySnapshot) -> Result<usize, MemoryError> {
197        let before = self.doc.clone();
198        let changed = self.doc.merge(remote);
199        self.keep(before, "memory merge")?;
200        Ok(changed)
201    }
202}

205fn read_actions(conn: &Connection, only: Option<u64>) -> Result<Vec<SoftAction>, DbError> {
206    let mut st = conn.prepare("SELECT id, kind, gid, at_ms, undone_at_ms FROM soft_action WHERE (?1 IS NULL OR id = ?1) ORDER BY id")?;
207    let mut rows = st.query([only.map(|n| n as i64)])?;
208    let mut out = Vec::new();
209    while let Some(r) = rows.next()? {
210        let word: String = r.get(1)?;
211        let Some(kind) = SoftKind::from_word(&word) else {
212            return Err(DbError::Damaged("a soft action of a kind this build does not know".into()));
213        };
214        out.push(SoftAction { id: r.get::<_, i64>(0)? as u64, kind, gid: r.get(2)?, at_ms: r.get::<_, i64>(3)? as u64, undone: r.get::<_, Option<i64>>(4)?.is_some() });
215    }
216    Ok(out)
217}

219fn kind_of(word: &str) -> Kind {
220    Kind::from_word(word)
221}
222
223fn vector_bytes(v: &[f32]) -> Vec<u8> {
224    v.iter().flat_map(|x| x.to_le_bytes()).collect()
225}
226
227fn vector_of(bytes: &[u8]) -> Result<Vec<f32>, DbError> {
228    if bytes.len() % 4 != 0 {
229        return Err(DbError::Damaged("an embedding is not a whole number of floats".into()));
230    }
231    Ok(bytes.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
232}

235fn load_doc(conn: &Connection) -> Result<MemoryDoc, DbError> {
236    let mut by_gid: BTreeMap<String, Fact> = BTreeMap::new();
237    let mut order: Vec<String> = Vec::new();
238    {
239        let mut st = conn.prepare(
240            "SELECT id, gid, text, text_at_ms, learned_at_ms, kind, who, place, said_when, hidden, visibility_at_ms, icon, cover_picture, cover_at_ms
241             FROM fact ORDER BY id",
242        )?;
243        let mut rows = st.query([])?;
244        while let Some(r) = rows.next()? {
245            let gid: String = r.get(1)?;
246            let who: String = r.get(6)?;
247            let who: Vec<String> = serde_json::from_str(&who).map_err(|e| DbError::Damaged(format!("a fact's names: {e}")))?;
248            let hidden: i64 = r.get(9)?;
249            let at: i64 = r.get(10)?;
250            let visibility = if hidden == 1 { Visibility::HiddenByChild { at_ms: at as u64 } } else { Visibility::Shown { at_ms: at as u64 } };
251            let icon: Option<String> = r.get(11)?;
252            let icon = icon.and_then(|n| {
253                let icon = IconId::new(&n).ok();
254                if icon.is_none() {
255                    warn!("a fact names an icon that is not admitted; reading it as none");
256                }
257                icon
258            });
259            let cover = match (r.get::<_, Option<String>>(12)?, r.get::<_, Option<i64>>(13)?) {
260                (Some(picture), Some(at)) => Cover::Chosen { picture: PictureId(picture), at_ms: at as u64 },
261                _ => Cover::Unchosen,
262            };
263            order.push(gid.clone());
264            by_gid.insert(
265                gid.clone(),
266                Fact {
267                    id: r.get::<_, i64>(0)? as u64,
268                    text: r.get(2)?,
269                    text_at_ms: r.get::<_, i64>(3)? as u64,
270                    learned_at_ms: r.get::<_, i64>(4)? as u64,
271                    kind: kind_of(&r.get::<_, String>(5)?),
272                    who,
273                    place: r.get(7)?,
274                    when: r.get(8)?,
275                    pictures: Vec::new(),
276                    embedding: Vec::new(),
277                    gid,
278                    visibility,
279                    icon,
280                    cover,
281                    mentions: Vec::new(),
282                    revisions: Vec::new(),
283                },
284            );
285        }
286    }
287    {
288        let mut st = conn.prepare("SELECT gid, at_ms, device, turn FROM fact_mention ORDER BY gid, at_ms, device, turn")?;
289        let mut rows = st.query([])?;
290        while let Some(r) = rows.next()? {
291            if let Some(f) = by_gid.get_mut(&r.get::<_, String>(0)?) {
292                f.mentions.push(Mention { at_ms: r.get::<_, i64>(1)? as u64, device: r.get(2)?, turn: r.get(3)? });
293            }
294        }
295    }
296    {
297        let mut st = conn.prepare("SELECT gid, at_ms, text, device FROM fact_revision ORDER BY gid, at_ms, text")?;
298        let mut rows = st.query([])?;
299        while let Some(r) = rows.next()? {
300            if let Some(f) = by_gid.get_mut(&r.get::<_, String>(0)?) {
301                f.revisions.push(Revision { at_ms: r.get::<_, i64>(1)? as u64, text: r.get(2)?, device: r.get(3)? });
302            }
303        }
304    }
305    {
306        let mut st = conn.prepare("SELECT gid, vector FROM fact_embedding")?;
307        let mut rows = st.query([])?;
308        while let Some(r) = rows.next()? {
309            if let Some(f) = by_gid.get_mut(&r.get::<_, String>(0)?) {
310                f.embedding = vector_of(&r.get::<_, Vec<u8>>(1)?)?;
311            }
312        }
313    }
314    {
315        let mut st = conn.prepare("SELECT gid, picture_id FROM fact_picture ORDER BY gid, position")?;
316        let mut rows = st.query([])?;
317        while let Some(r) = rows.next()? {
318            if let Some(f) = by_gid.get_mut(&r.get::<_, String>(0)?) {
319                f.pictures.push(PictureId(r.get(1)?));
320            }
321        }
322    }
323    let mut tombstones = Vec::new();
324    {
325        let mut st = conn.prepare("SELECT gid, at_ms, device FROM tombstone ORDER BY gid")?;
326        let mut rows = st.query([])?;
327        while let Some(r) = rows.next()? {
328            tombstones.push(Tombstone { gid: r.get(0)?, at_ms: r.get::<_, i64>(1)? as u64, device: r.get(2)?, turns: Vec::new() });
329        }
330    }
331    {
332        let mut st = conn.prepare("SELECT gid, turn FROM tombstone_turn ORDER BY gid, turn")?;
333        let mut rows = st.query([])?;
334        while let Some(r) = rows.next()? {
335            let gid: String = r.get(0)?;
336            if let Some(t) = tombstones.iter_mut().find(|t| t.gid == gid) {
337                t.turns.push(r.get(1)?);
338            }
339        }
340    }
341    // Ids are never reused: the largest ever handed out is what the table's sequence remembers, after the
342    // newest fact has been forgotten as well as before.
343    let issued: Option<i64> = conn.query_row("SELECT seq FROM sqlite_sequence WHERE name = 'fact'", [], |r| r.get(0)).optional()?;
344    let facts: Vec<Fact> = order.into_iter().filter_map(|g| by_gid.remove(&g)).collect();
345    Ok(MemoryDoc::from_parts(issued.map_or(1, |s| s as u64 + 1), facts, tombstones))
346}

348fn insert_fact(tx: &Transaction<'_>, f: &Fact) -> Result<(), DbError> {
349    let (hidden, vis_at) = (i64::from(f.visibility.is_hidden()), f.visibility.at_ms() as i64);
350    let (cover_picture, cover_at) = match &f.cover {
351        Cover::Chosen { picture, at_ms } => (Some(picture.0.clone()), Some(*at_ms as i64)),
352        Cover::Unchosen => (None, None),
353    };
354    tx.execute(
355        "INSERT INTO fact (id, gid, text, text_at_ms, learned_at_ms, kind, who, place, said_when, hidden, visibility_at_ms, icon, cover_picture, cover_at_ms)
356         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)",
357        params![
358            f.id as i64,
359            f.gid,
360            f.text,
361            f.text_at_ms as i64,
362            f.learned_at_ms as i64,
363            f.kind.word(),
364            serde_json::to_string(&f.who).map_err(|e| DbError::Damaged(e.to_string()))?,
365            f.place,
366            f.when,
367            hidden,
368            vis_at,
369            f.icon.as_ref().map(IconId::as_str),
370            cover_picture,
371            cover_at,
372        ],
373    )?;
374    write_children(tx, None, f)
375}
376
377fn update_fact(tx: &Transaction<'_>, before: &Fact, f: &Fact) -> Result<(), DbError> {
378    let (hidden, vis_at) = (i64::from(f.visibility.is_hidden()), f.visibility.at_ms() as i64);
379    let (cover_picture, cover_at) = match &f.cover {
380        Cover::Chosen { picture, at_ms } => (Some(picture.0.clone()), Some(*at_ms as i64)),
381        Cover::Unchosen => (None, None),
382    };
383    tx.execute(
384        "UPDATE fact SET text = ?2, text_at_ms = ?3, kind = ?4, who = ?5, place = ?6, said_when = ?7, hidden = ?8, visibility_at_ms = ?9,
385                icon = ?10, cover_picture = ?11, cover_at_ms = ?12
386         WHERE gid = ?1",
387        params![
388            f.gid,
389            f.text,
390            f.text_at_ms as i64,
391            f.kind.word(),
392            serde_json::to_string(&f.who).map_err(|e| DbError::Damaged(e.to_string()))?,
393            f.place,
394            f.when,
395            hidden,
396            vis_at,
397            f.icon.as_ref().map(IconId::as_str),
398            cover_picture,
399            cover_at,
400        ],
401    )?;
402    write_children(tx, Some(before), f)
403}

406fn write_children(tx: &Transaction<'_>, before: Option<&Fact>, f: &Fact) -> Result<(), DbError> {
407    for m in &f.mentions {
408        tx.execute("INSERT OR IGNORE INTO fact_mention (gid, at_ms, device, turn) VALUES (?1, ?2, ?3, ?4)", params![f.gid, m.at_ms as i64, m.device, m.turn])?;
409    }
410    for r in &f.revisions {
411        tx.execute(
412            "INSERT INTO fact_revision (gid, at_ms, text, device) VALUES (?1, ?2, ?3, ?4)
413             ON CONFLICT (gid, at_ms, text) DO UPDATE SET device = excluded.device",
414            params![f.gid, r.at_ms as i64, r.text, r.device],
415        )?;
416    }
417    if f.embedding.is_empty() {
418        tx.execute("DELETE FROM fact_embedding WHERE gid = ?1", [&f.gid])?;
419    } else if before.is_none_or(|b| b.embedding != f.embedding) {
420        tx.execute("INSERT OR REPLACE INTO fact_embedding (gid, vector) VALUES (?1, ?2)", params![f.gid, vector_bytes(&f.embedding)])?;
421    }
422    let had: Vec<&PictureId> = before.map(|b| b.pictures.iter().collect()).unwrap_or_default();
423    if before.is_none_or(|b| b.pictures != f.pictures) {
424        // Only the links that went are deleted (a delete is what lets the schema drop a picture nothing else
425        // uses), and only the ones that are new are added.
426        let keep: BTreeSet<&str> = f.pictures.iter().map(|p| p.0.as_str()).collect();
427        for p in had.iter().filter(|p| !keep.contains(p.0.as_str())) {
428            tx.execute("DELETE FROM fact_picture WHERE gid = ?1 AND picture_id = ?2", params![f.gid, p.0])?;
429        }
430        let existing: BTreeSet<&str> = had.iter().map(|p| p.0.as_str()).collect();
431        let mut position: i64 = tx.query_row("SELECT COALESCE(MAX(position) + 1, 0) FROM fact_picture WHERE gid = ?1", [&f.gid], |r| r.get(0))?;
432        for p in f.pictures.iter().filter(|p| !existing.contains(p.0.as_str())) {
433            tx.execute("INSERT INTO fact_picture (gid, position, picture_id) VALUES (?1, ?2, ?3)", params![f.gid, position, p.0])?;
434            position += 1;
435        }
436    }
437    Ok(())
438}

442fn write_changes(tx: &Transaction<'_>, before: &MemoryDoc, after: &MemoryDoc) -> Result<(), DbError> {
443    let had: BTreeMap<&str, &Tombstone> = before.tombstones().iter().map(|t| (t.gid.as_str(), t)).collect();
444    for t in after.tombstones() {
445        match had.get(t.gid.as_str()) {
446            Some(old) if *old == t => {}
447            Some(_) => {
448                tx.execute("UPDATE tombstone SET at_ms = ?2, device = ?3 WHERE gid = ?1", params![t.gid, t.at_ms as i64, t.device])?;
449            }
450            None => {
451                tx.execute("INSERT INTO tombstone (gid, at_ms, device) VALUES (?1, ?2, ?3)", params![t.gid, t.at_ms as i64, t.device])?;
452            }
453        }
454        // The identities of the turns the memory was told in. Writing one empties that turn's words, here and in any
455        // line of it that arrives later (the schema's `tombstone_turn` triggers).
456        for turn in &t.turns {
457            tx.execute("INSERT OR IGNORE INTO tombstone_turn (gid, turn) VALUES (?1, ?2)", params![t.gid, turn])?;
458        }
459    }
460    let was: BTreeMap<&str, &Fact> = before.facts().iter().map(|f| (f.gid.as_str(), f)).collect();
461    let now: BTreeSet<&str> = after.facts().iter().map(|f| f.gid.as_str()).collect();
462    for gid in was.keys().filter(|g| !now.contains(**g)) {
463        tx.execute("DELETE FROM fact WHERE gid = ?1", [gid])?;
464    }
465    for f in after.facts() {
466        match was.get(f.gid.as_str()) {
467            None => insert_fact(tx, f)?,
468            Some(b) if *b != f => update_fact(tx, b, f)?,
469            Some(_) => {}
470        }
471    }
472    Ok(())
473}