use std::collections::HashMap;

use crate::Database;
use ironclad_core::{IroncladError, 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 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> {
    let conn = db.conn();
    let id = uuid::Uuid::new_v4().to_string();
    conn.execute(
        "INSERT INTO tool_calls (id, turn_id, tool_name, input, output, status, duration_ms) \
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
        rusqlite::params![id, turn_id, tool_name, input, output, status, duration_ms],
    )
    .map_err(|e| IroncladError::Database(e.to_string()))?;
    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, status, duration_ms, created_at \
             FROM tool_calls WHERE turn_id = ?1 ORDER BY created_at ASC",
        )
        .map_err(|e| IroncladError::Database(e.to_string()))?;

    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)?,
                status: row.get(5)?,
                duration_ms: row.get(6)?,
                created_at: row.get(7)?,
            })
        })
        .map_err(|e| IroncladError::Database(e.to_string()))?;

    rows.collect::<std::result::Result<Vec<_>, _>>()
        .map_err(|e| IroncladError::Database(e.to_string()))
}

/// 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.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",
        )
        .map_err(|e| IroncladError::Database(e.to_string()))?;

    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)?,
                status: row.get(5)?,
                duration_ms: row.get(6)?,
                created_at: row.get(7)?,
            })
        })
        .map_err(|e| IroncladError::Database(e.to_string()))?;

    let mut map: HashMap<String, Vec<ToolCallRecord>> = HashMap::new();
    for row in rows {
        let record = row.map_err(|e| IroncladError::Database(e.to_string()))?;
        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());
    }
}
