roboticus-db 0.11.4

use std::collections::HashMap;

use crate::{Database, DbResultExt};
use roboticus_core::Result;

#[derive(Debug, Clone)]
pub struct ToolCallRecord {
    pub id: String,
    pub turn_id: String,
    pub tool_name: String,
    pub input: String,
    pub output: Option<String>,
    pub skill_id: Option<String>,
    pub skill_name: Option<String>,
    pub skill_hash: Option<String>,
    pub status: String,
    pub duration_ms: Option<i64>,
    pub created_at: String,
}

pub fn record_tool_call(
    db: &Database,
    turn_id: &str,
    tool_name: &str,
    input: &str,
    output: Option<&str>,
    status: &str,
    duration_ms: Option<i64>,
) -> Result<String> {
    record_tool_call_with_skill(
        db,
        turn_id,
        tool_name,
        input,
        output,
        status,
        duration_ms,
        None,
        None,
        None,
    )
}

#[allow(clippy::too_many_arguments)]
pub fn record_tool_call_with_skill(
    db: &Database,
    turn_id: &str,
    tool_name: &str,
    input: &str,
    output: Option<&str>,
    status: &str,
    duration_ms: Option<i64>,
    skill_id: Option<&str>,
    skill_name: Option<&str>,
    skill_hash: Option<&str>,
) -> Result<String> {
    let conn = db.conn();
    let id = uuid::Uuid::new_v4().to_string();
    let tx = conn.unchecked_transaction().db_err()?;
    tx.execute(
        "INSERT INTO tool_calls (id, turn_id, tool_name, input, output, skill_id, skill_name, \
         skill_hash, status, duration_ms) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
        rusqlite::params![
            id,
            turn_id,
            tool_name,
            input,
            output,
            skill_id,
            skill_name,
            skill_hash,
            status,
            duration_ms
        ],
    )
    .db_err()?;
    if let Some(skill_id) = skill_id {
        tx.execute(
            "UPDATE skills
             SET usage_count = usage_count + 1,
                 last_used_at = datetime('now')
             WHERE id = ?1",
            [skill_id],
        )
        .db_err()?;
    }
    tx.commit().db_err()?;
    Ok(id)
}

pub fn get_tool_calls_for_turn(db: &Database, turn_id: &str) -> Result<Vec<ToolCallRecord>> {
    let conn = db.conn();
    let mut stmt = conn
        .prepare(
            "SELECT id, turn_id, tool_name, input, output, skill_id, skill_name, skill_hash, \
             status, duration_ms, created_at \
             FROM tool_calls WHERE turn_id = ?1 ORDER BY created_at ASC",
        )
        .db_err()?;

    let rows = stmt
        .query_map([turn_id], |row| {
            Ok(ToolCallRecord {
                id: row.get(0)?,
                turn_id: row.get(1)?,
                tool_name: row.get(2)?,
                input: row.get(3)?,
                output: row.get(4)?,
                skill_id: row.get(5)?,
                skill_name: row.get(6)?,
                skill_hash: row.get(7)?,
                status: row.get(8)?,
                duration_ms: row.get(9)?,
                created_at: row.get(10)?,
            })
        })
        .db_err()?;

    rows.collect::<std::result::Result<Vec<_>, _>>().db_err()
}

/// Batch-fetch all tool calls for every turn in a session, grouped by turn_id.
/// Eliminates the N+1 query pattern when analyzing sessions with many turns.
pub fn get_tool_calls_for_session(
    db: &Database,
    session_id: &str,
) -> Result<HashMap<String, Vec<ToolCallRecord>>> {
    let conn = db.conn();
    let mut stmt = conn
        .prepare(
            "SELECT tc.id, tc.turn_id, tc.tool_name, tc.input, tc.output, tc.skill_id, \
                    tc.skill_name, tc.skill_hash, tc.status, tc.duration_ms, tc.created_at \
             FROM tool_calls tc \
             INNER JOIN turns t ON tc.turn_id = t.id \
             WHERE t.session_id = ?1 \
             ORDER BY tc.created_at ASC",
        )
        .db_err()?;

    let rows = stmt
        .query_map([session_id], |row| {
            Ok(ToolCallRecord {
                id: row.get(0)?,
                turn_id: row.get(1)?,
                tool_name: row.get(2)?,
                input: row.get(3)?,
                output: row.get(4)?,
                skill_id: row.get(5)?,
                skill_name: row.get(6)?,
                skill_hash: row.get(7)?,
                status: row.get(8)?,
                duration_ms: row.get(9)?,
                created_at: row.get(10)?,
            })
        })
        .db_err()?;

    let mut map: HashMap<String, Vec<ToolCallRecord>> = HashMap::new();
    for row in rows {
        let record = row.db_err()?;
        map.entry(record.turn_id.clone()).or_default().push(record);
    }
    Ok(map)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_db() -> Database {
        let db = Database::new(":memory:").unwrap();
        // tool_calls has FK to turns, which has FK to sessions — seed parent rows
        let conn = db.conn();
        conn.execute(
            "INSERT INTO sessions (id, agent_id) VALUES ('s1', 'agent-1')",
            [],
        )
        .unwrap();
        conn.execute("INSERT INTO turns (id, session_id) VALUES ('t1', 's1')", [])
            .unwrap();
        drop(conn);
        db
    }

    #[test]
    fn record_and_retrieve_tool_call() {
        let db = test_db();
        let id = record_tool_call(
            &db,
            "t1",
            "bash",
            r#"{"cmd":"ls"}"#,
            Some("file1\nfile2"),
            "success",
            Some(42),
        )
        .unwrap();
        assert!(!id.is_empty());

        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].tool_name, "bash");
        assert_eq!(calls[0].duration_ms, Some(42));
    }

    #[test]
    fn empty_turn_returns_empty_vec() {
        let db = test_db();
        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert!(calls.is_empty());
    }

    #[test]
    fn multiple_calls_ordered_by_time() {
        let db = test_db();
        record_tool_call(&db, "t1", "read", "{}", None, "success", Some(10)).unwrap();
        record_tool_call(&db, "t1", "write", "{}", None, "success", Some(20)).unwrap();

        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].tool_name, "read");
        assert_eq!(calls[1].tool_name, "write");
    }

    #[test]
    fn record_tool_call_no_output_no_duration() {
        let db = test_db();
        let id = record_tool_call(
            &db,
            "t1",
            "search",
            r#"{"q":"test"}"#,
            None,
            "pending",
            None,
        )
        .unwrap();
        assert!(!id.is_empty());
        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert!(calls[0].output.is_none());
        assert!(calls[0].duration_ms.is_none());
        assert_eq!(calls[0].status, "pending");
    }

    #[test]
    fn record_tool_call_error_status() {
        let db = test_db();
        record_tool_call(
            &db,
            "t1",
            "bash",
            r#"{"cmd":"rm -rf /"}"#,
            Some("permission denied"),
            "error",
            Some(5),
        )
        .unwrap();
        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert_eq!(calls[0].status, "error");
        assert_eq!(calls[0].output.as_deref(), Some("permission denied"));
    }

    #[test]
    fn get_tool_calls_nonexistent_turn() {
        let db = test_db();
        let calls = get_tool_calls_for_turn(&db, "nonexistent").unwrap();
        assert!(calls.is_empty());
    }

    #[test]
    fn batch_get_tool_calls_for_session() {
        let db = test_db();
        let conn = db.conn();
        conn.execute("INSERT INTO turns (id, session_id) VALUES ('t2', 's1')", [])
            .unwrap();
        drop(conn);
        record_tool_call(&db, "t1", "read", "{}", None, "success", Some(10)).unwrap();
        record_tool_call(&db, "t1", "write", "{}", None, "success", Some(20)).unwrap();
        record_tool_call(&db, "t2", "bash", "{}", None, "error", Some(5)).unwrap();

        let map = get_tool_calls_for_session(&db, "s1").unwrap();
        assert_eq!(map.len(), 2);
        assert_eq!(map["t1"].len(), 2);
        assert_eq!(map["t2"].len(), 1);
        assert_eq!(map["t2"][0].tool_name, "bash");
    }

    #[test]
    fn batch_get_empty_session() {
        let db = test_db();
        let map = get_tool_calls_for_session(&db, "s1").unwrap();
        assert!(map.is_empty());
    }

    #[test]
    fn tool_call_fields_populated() {
        let db = test_db();
        record_tool_call(
            &db,
            "t1",
            "bash",
            r#"{"cmd":"echo hi"}"#,
            Some("hi"),
            "success",
            Some(100),
        )
        .unwrap();
        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert!(!calls[0].id.is_empty());
        assert_eq!(calls[0].turn_id, "t1");
        assert!(!calls[0].created_at.is_empty());
    }

    #[test]
    fn record_tool_call_with_skill_attribution() {
        let db = test_db();
        let conn = db.conn();
        conn.execute(
            "INSERT INTO skills (id, name, kind, source_path, content_hash, risk_level, version, author, registry_source)
             VALUES ('skill-123', 'deploy', 'instruction', '/tmp/deploy.md', 'hash-abc', 'Caution', '0.0.0', 'local', 'local')",
            [],
        )
        .unwrap();
        drop(conn);
        record_tool_call_with_skill(
            &db,
            "t1",
            "run_script",
            r#"{"path":"deploy.sh"}"#,
            Some("ok"),
            "success",
            Some(33),
            Some("skill-123"),
            Some("deploy"),
            Some("hash-abc"),
        )
        .unwrap();
        let calls = get_tool_calls_for_turn(&db, "t1").unwrap();
        assert_eq!(calls.len(), 1);
        assert_eq!(calls[0].skill_id.as_deref(), Some("skill-123"));
        assert_eq!(calls[0].skill_name.as_deref(), Some("deploy"));
        assert_eq!(calls[0].skill_hash.as_deref(), Some("hash-abc"));

        let skill = crate::skills::get_skill(&db, "skill-123")
            .unwrap()
            .expect("skill row should exist");
        assert_eq!(skill.usage_count, 1);
        assert!(skill.last_used_at.is_some());
    }
}