3use ::log::{debug, error, trace};
4use rusqlite::{OptionalExtension, params};
5use whiskers_core::{ChatError, ChatState, ChatStore, ChatTurn, Speaker};

7use crate::{DbError, Store};

11pub struct SqliteChat {
12    store: Store,
13}

15impl SqliteChat {
16    pub fn new(store: Store) -> Self {
17        Self { store }
18    }
19}
20
21fn speaker_word(s: Speaker) -> &'static str {
22    match s {
23        Speaker::Child => "child",
24        Speaker::Whiskers => "whiskers",
25    }
26}
27
28fn speaker_of(word: &str) -> Result<Speaker, DbError> {
29    match word {
30        "child" => Ok(Speaker::Child),
31        "whiskers" => Ok(Speaker::Whiskers),
32        other => Err(DbError::Damaged(format!("a chat turn by {other:?}"))),
33    }
34}
35
36impl ChatStore for SqliteChat {
37    fn load(&self) -> Result<ChatState, ChatError> {
38        let conn = self.store.lock();
39        let read = || -> Result<ChatState, DbError> {
40            let session: Option<(String, i64, i64, String)> = conn
41                .query_row("SELECT summary, last_active_ms, version, scrubbed FROM chat_session WHERE id = 1", [], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)))
42                .optional()?;
43            let mut st = conn.prepare("SELECT speaker, text FROM chat_turn ORDER BY seq")?;
44            let mut rows = st.query([])?;
45            let mut turns = Vec::new();
46            while let Some(r) = rows.next()? {
47                turns.push(ChatTurn { speaker: speaker_of(&r.get::<_, String>(0)?)?, text: r.get(1)? });
48            }
49            Ok(match session {
50                Some((summary, last_active_ms, version, scrubbed)) => ChatState {
51                    summary,
52                    turns,
53                    last_active_ms: last_active_ms as u64,
54                    version: version as u64,
55                    scrubbed: serde_json::from_str(&scrubbed).map_err(|e| DbError::Damaged(format!("the chat's forgettings: {e}")))?,
56                },
57                None => ChatState { turns, ..ChatState::default() },
58            })
59        };
60        let state = read().map_err(|e| {
61            error!("the chat cannot be read: {e}");
62            ChatError(format!("chat is damaged: {e}"))
63        })?;
64        debug!("chat loaded: {} turns, version {}, summary {} bytes", state.turns.len(), state.version, state.summary.len());
65        Ok(state)
66    }
67
68    fn save(&mut self, state: &ChatState) -> Result<(), ChatError> {
69        debug!("saving chat: {} turns, version {}", state.turns.len(), state.version);
70        let scrubbed = serde_json::to_string(&state.scrubbed).map_err(|e| ChatError(e.to_string()))?;
71        let mut emptied = false;
72        self.store
73            .transaction(|tx| {
74                // A copy emptied for a forgetting this database had not heard of: what it held is deleted, and the
75                // old bytes must not wait in the write-ahead log (see the scrub below).
76                let before: Option<String> = tx.query_row("SELECT scrubbed FROM chat_session WHERE id = 1", [], |r| r.get(0)).optional()?;
77                emptied = before.map_or(scrubbed != "[]", |b| b != scrubbed);
78                tx.execute(
79                    "INSERT INTO chat_session (id, summary, last_active_ms, version, scrubbed) VALUES (1, ?1, ?2, ?3, ?4)
80                     ON CONFLICT (id) DO UPDATE SET summary = excluded.summary, last_active_ms = excluded.last_active_ms, version = excluded.version,
81                                                    scrubbed = excluded.scrubbed",
82                    params![state.summary, state.last_active_ms as i64, state.version as i64, scrubbed],
83                )?;
84                // The turns that are kept stay and are never rewritten: the stored turns, less some at the front
85                // (folded into the summary), are a prefix of the turns now. Of the stored turns the fewest are
86                // dropped for that to hold; what follows the prefix is appended. A copy that shares nothing
87                // with what is stored replaces it.
88                let stored: Vec<(i64, String, String)> = {
89                    let mut st = tx.prepare("SELECT seq, speaker, text FROM chat_turn ORDER BY seq")?;
90                    let rows = st.query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)))?;
91                    rows.collect::<Result<_, _>>()?
92                };
93                let same = |s: &(i64, String, String), t: &ChatTurn| s.1 == speaker_word(t.speaker) && s.2 == t.text;
94                let dropped = (0..=stored.len())
95                    .find(|d| {
96                        let kept = &stored[*d..];
97                        kept.len() <= state.turns.len() && kept.iter().zip(&state.turns).all(|(s, t)| same(s, t))
98                    })
99                    .unwrap_or(stored.len());
100                trace!("chat: {dropped} stored turn(s) dropped from the front, {} kept", stored.len() - dropped);
101                for s in &stored[..dropped] {
102                    tx.execute("DELETE FROM chat_turn WHERE seq = ?1", [s.0])?;
103                }
104                for t in &state.turns[stored.len() - dropped..] {
105                    tx.execute("INSERT INTO chat_turn (speaker, text) VALUES (?1, ?2)", params![speaker_word(t.speaker), t.text])?;
106                }
107                Ok(())
108            })
109            .map_err(|e| {
110                error!("the chat was not kept: {e}");
111                ChatError(e.to_string())
112            })?;
113        if emptied && let Err(e) = self.store.scrub() {
114            // The chat is emptied in the database; what could not be done is emptying the write-ahead log of it.
115            error!("the emptied chat may still be in the write-ahead log: {e}");
116        }
117        Ok(())
118    }
119}