vein-database 0.1.0

use anyhow::Result;
use rusqlite::Connection;
use std::fs;

use crate::paths::Paths;

struct SQLEntry {
    id: String,
    role: String,
    summary: String,
    code_style: bool,
    agent: String,
    timestamp: String,
}

pub struct SqliteDb {
    conn: Connection,
}

impl SqliteDb {
    pub fn new() -> Result<Self> {
        let path_ref = Paths::get_dialogues_db();
        if !path_ref.exists() {
            fs::create_dir_all(&path_ref)?;
        }
        let db_path = path_ref.join("dialogues.db");
        let conn = Connection::open(&db_path)?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS medium (
                id TEXT PRIMARY KEY,
                role TEXT NOT NULL,
                summary TEXT NOT NULL,
                code_style INTEGER NOT NULL,
                agent TEXT NOT NULL,
                timestamp TEXT NOT NULL
            )",
            [],
        )?;
        conn.execute(
            "CREATE UNIQUE INDEX IF NOT EXISTS idx_medium_summary ON medium(summary)",
            [],
        )?;
        Ok(Self { conn })
    }

    pub fn create(
        &self,
        id: &str,
        role: &str,
        summary: &str,
        code_style: bool,
        agent: &str,
        timestamp: &str,
    ) -> Result<String> {
        self.conn.execute(
            "INSERT INTO medium (id, role, summary, code_style, agent, timestamp) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            (id, role, summary, code_style as i32, agent, timestamp),
        )?;
        Ok(id.to_string())
    }

    pub fn upsert(
        &self,
        id: &str,
        role: &str,
        summary: &str,
        code_style: bool,
        agent: &str,
        timestamp: &str,
    ) -> Result<String> {
        self.conn.execute(
            "INSERT OR IGNORE INTO medium (id, role, summary, code_style, agent, timestamp) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
            (id, role, summary, code_style as i32, agent, timestamp),
        )?;
        Ok(id.to_string())
    }

    pub fn update(
        &self,
        id: &str,
        role: &str,
        summary: &str,
        code_style: bool,
        agent: &str,
        timestamp: &str,
    ) -> Result<bool> {
        let rows = self.conn.execute(
            "UPDATE medium SET role = ?1, summary = ?2, code_style = ?3, agent = ?4, timestamp = ?5 WHERE id = ?6",
            (role, summary, code_style as i32, agent, timestamp, id),
        )?;
        Ok(rows > 0)
    }

    pub fn delete(&self, id: &str) -> Result<bool> {
        let rows = self.conn.execute("DELETE FROM medium WHERE id = ?1", [id])?;
        Ok(rows > 0)
    }

    pub fn query(&self, query_str: &str) -> Result<String> {
        let mut stmt = self.conn.prepare(query_str)?;
        let rows = stmt.query_map([], |row| {
            Ok(SQLEntry {
                id: row.get(0)?,
                role: row.get(1)?,
                summary: row.get(2)?,
                code_style: row.get::<_, i32>(3)? != 0,
                agent: row.get(4)?,
                timestamp: row.get(5)?,
            })
        })?;

        let results: Vec<SQLEntry> = rows.filter_map(|r| r.ok()).collect();
        if results.is_empty() {
            return Ok("No results found".to_string());
        }

        let mut output = String::new();
        for entry in results {
            output.push_str(&format!(
                "id: {}\nrole: {}\nsummary: {}\ncode_style: {}\nagent: {}\ntimestamp: {}\n---\n",
                entry.id, entry.role, entry.summary, entry.code_style, entry.agent, entry.timestamp
            ));
        }
        Ok(output)
    }

    pub fn count_all(&self) -> Result<usize> {
        let count: i64 = self
            .conn
            .query_row("SELECT COUNT(*) FROM medium", [], |row| row.get(0))?;
        Ok(count as usize)
    }
}