agent-ask 0.1.0

//! SQLite-backed artifact store (mirror of `src/store.ts`).

use std::sync::{Arc, Mutex};

use rusqlite::{params, Connection};
use serde_json::Value;

use crate::artifact::cid_of;
use crate::error::Error;

const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS artifacts (
  cid TEXT PRIMARY KEY,
  kind TEXT NOT NULL,
  created_at TEXT NOT NULL,
  body TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_artifacts_kind_created ON artifacts(kind, created_at);
CREATE INDEX IF NOT EXISTS idx_artifacts_created ON artifacts(created_at);
CREATE TABLE IF NOT EXISTS question_tags (
  cid TEXT NOT NULL,
  tag TEXT NOT NULL,
  PRIMARY KEY (cid, tag),
  FOREIGN KEY (cid) REFERENCES artifacts(cid) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_question_tags_tag ON question_tags(tag);
"#;

#[derive(Clone)]
pub struct Store {
    conn: Arc<Mutex<Connection>>,
}

impl Store {
    pub fn open(path: &str) -> Result<Self, Error> {
        let conn = Connection::open(path)?;
        conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA foreign_keys = ON;")?;
        conn.execute_batch(SCHEMA)?;
        Ok(Self { conn: Arc::new(Mutex::new(conn)) })
    }

    pub fn insert_artifact(&self, artifact: &Value) -> Result<String, Error> {
        let cid = cid_of(artifact)?;
        let kind = artifact.get("kind").and_then(|v| v.as_str())
            .ok_or_else(|| Error::Invalid("missing kind".into()))?
            .to_string();
        let created_at = artifact.get("created_at").and_then(|v| v.as_str())
            .ok_or_else(|| Error::Invalid("missing created_at".into()))?
            .to_string();
        let body = serde_json::to_string(artifact)?;
        let conn = self.conn.lock().unwrap();
        conn.execute(
            "INSERT OR IGNORE INTO artifacts(cid, kind, created_at, body) VALUES (?1, ?2, ?3, ?4)",
            params![&cid, &kind, &created_at, &body],
        )?;
        if kind == "question" {
            if let Some(tags) = artifact.get("tags").and_then(|v| v.as_array()) {
                for t in tags {
                    if let Some(tag) = t.as_str() {
                        conn.execute(
                            "INSERT OR IGNORE INTO question_tags(cid, tag) VALUES (?1, ?2)",
                            params![&cid, tag],
                        )?;
                    }
                }
            }
        }
        Ok(cid)
    }

    pub fn get_artifact(&self, cid: &str) -> Result<Option<Value>, Error> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare("SELECT body FROM artifacts WHERE cid = ?1")?;
        let mut rows = stmt.query(params![cid])?;
        if let Some(row) = rows.next()? {
            let body: String = row.get(0)?;
            Ok(Some(serde_json::from_str(&body)?))
        } else {
            Ok(None)
        }
    }

    pub fn has_artifact(&self, cid: &str) -> Result<bool, Error> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = conn.prepare("SELECT 1 FROM artifacts WHERE cid = ?1")?;
        Ok(stmt.exists(params![cid])?)
    }

    pub fn list_questions(
        &self,
        tag: Option<&str>,
        since: Option<&str>,
        limit: i64,
    ) -> Result<Vec<Value>, Error> {
        let conn = self.conn.lock().unwrap();
        let mut clauses: Vec<&str> = vec!["a.kind = 'question'"];
        let mut params_vec: Vec<rusqlite::types::Value> = Vec::new();
        if let Some(t) = tag {
            clauses.push("qt.tag = ?");
            params_vec.push(t.to_string().into());
        }
        if let Some(s) = since {
            clauses.push("a.created_at > ?");
            params_vec.push(s.to_string().into());
        }
        let join = if tag.is_some() { "JOIN question_tags qt ON qt.cid = a.cid" } else { "" };
        let sql = format!(
            "SELECT DISTINCT a.body FROM artifacts a {join} WHERE {} ORDER BY a.created_at DESC LIMIT ?",
            clauses.join(" AND ")
        );
        params_vec.push(limit.into());

        let mut stmt = conn.prepare(&sql)?;
        let rows = stmt.query_map(
            rusqlite::params_from_iter(params_vec.iter()),
            |row| row.get::<_, String>(0),
        )?;
        let mut out = Vec::new();
        for row in rows {
            out.push(serde_json::from_str(&row?)?);
        }
        Ok(out)
    }

    pub fn stream_feed(&self, since: Option<&str>, limit: i64) -> Result<Vec<Value>, Error> {
        let conn = self.conn.lock().unwrap();
        let mut params_vec: Vec<rusqlite::types::Value> = Vec::new();
        let where_clause = if let Some(s) = since {
            params_vec.push(s.to_string().into());
            "WHERE created_at > ?"
        } else {
            ""
        };
        let sql = format!(
            "SELECT body FROM artifacts {where_clause} ORDER BY created_at DESC LIMIT ?"
        );
        params_vec.push(limit.into());

        let mut stmt = conn.prepare(&sql)?;
        let rows = stmt.query_map(
            rusqlite::params_from_iter(params_vec.iter()),
            |row| row.get::<_, String>(0),
        )?;
        let mut out = Vec::new();
        for row in rows {
            out.push(serde_json::from_str(&row?)?);
        }
        Ok(out)
    }
}