roboticus-db 0.11.4

use roboticus_core::delegation_tools::delegation_output_tool_name_sql_in_predicate;
use roboticus_core::{Result, RoboticusError};
use rusqlite::OptionalExtension;
use serde::Serialize;
use std::collections::HashMap;

use crate::{Database, DbResultExt};

#[derive(Debug, Clone)]
pub struct SubAgentRow {
    pub id: String,
    pub name: String,
    pub display_name: Option<String>,
    pub model: String,
    pub fallback_models_json: Option<String>,
    pub role: String,
    pub description: Option<String>,
    pub skills_json: Option<String>,
    pub enabled: bool,
    pub session_count: i64,
    pub last_used_at: Option<String>,
}

fn normalized_fallback_models_json(raw: Option<&str>) -> String {
    match raw.map(str::trim) {
        Some(v) if !v.is_empty() => v.to_string(),
        _ => "[]".to_string(),
    }
}

pub fn upsert_sub_agent(db: &Database, agent: &SubAgentRow) -> Result<()> {
    let conn = db.conn();
    let fallback_models_json =
        normalized_fallback_models_json(agent.fallback_models_json.as_deref());
    conn.execute(
        "INSERT INTO sub_agents (id, name, display_name, model, fallback_models_json, role, description, skills_json, enabled, session_count, last_used_at)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)
         ON CONFLICT(name) DO UPDATE SET
           display_name = excluded.display_name,
           model = excluded.model,
           fallback_models_json = excluded.fallback_models_json,
           role = excluded.role,
           description = excluded.description,
           skills_json = excluded.skills_json,
           enabled = excluded.enabled,
           session_count = excluded.session_count,
           last_used_at = excluded.last_used_at",
        rusqlite::params![
            agent.id,
            agent.name,
            agent.display_name,
            agent.model,
            fallback_models_json,
            agent.role,
            agent.description,
            agent.skills_json,
            agent.enabled as i32,
            agent.session_count,
            agent.last_used_at,
        ],
    )
    .map_err(|e| RoboticusError::Database(format!("upsert sub_agent: {e}")))?;
    Ok(())
}

pub fn list_sub_agents(db: &Database) -> Result<Vec<SubAgentRow>> {
    let conn = db.conn();
    let mut stmt = conn
        .prepare(
            "SELECT id, name, display_name, model, fallback_models_json, role, description, skills_json, enabled, session_count, last_used_at
             FROM sub_agents ORDER BY name",
        )
        .db_err()?;

    let rows = stmt
        .query_map([], |row| {
            Ok(SubAgentRow {
                id: row.get(0)?,
                name: row.get(1)?,
                display_name: row.get(2)?,
                model: row.get(3)?,
                fallback_models_json: Some(normalized_fallback_models_json(
                    row.get::<_, Option<String>>(4)?.as_deref(),
                )),
                role: row.get(5)?,
                description: row.get(6)?,
                skills_json: row.get(7)?,
                enabled: row.get::<_, i32>(8)? != 0,
                session_count: row.get(9)?,
                last_used_at: row.get(10)?,
            })
        })
        .db_err()?
        .collect::<std::result::Result<Vec<_>, _>>()
        .db_err()?;

    Ok(rows)
}

pub fn list_enabled_sub_agents(db: &Database) -> Result<Vec<SubAgentRow>> {
    let all = list_sub_agents(db)?;
    Ok(all.into_iter().filter(|a| a.enabled).collect())
}

pub fn record_subagent_usage(db: &Database, name: &str) -> Result<()> {
    let conn = db.conn();
    conn.execute(
        "UPDATE sub_agents
         SET session_count = session_count + 1,
             last_used_at = datetime('now')
         WHERE name = ?1",
        [name],
    )
    .db_err()?;
    Ok(())
}

pub fn list_session_counts_by_agent(db: &Database) -> Result<HashMap<String, i64>> {
    let conn = db.conn();
    let mut stmt = conn
        .prepare("SELECT agent_id, COUNT(*) FROM sessions GROUP BY agent_id")
        .db_err()?;

    let rows = stmt
        .query_map([], |row| {
            let agent_id: String = row.get(0)?;
            let count: i64 = row.get(1)?;
            Ok((agent_id, count))
        })
        .db_err()?
        .collect::<std::result::Result<Vec<_>, _>>()
        .db_err()?;

    Ok(rows.into_iter().collect())
}

pub fn delete_sub_agent(db: &Database, name: &str) -> Result<bool> {
    let conn = db.conn();
    let deleted = conn
        .execute(
            "DELETE FROM sub_agents WHERE name = ?1",
            rusqlite::params![name],
        )
        .db_err()?;
    Ok(deleted > 0)
}

/// Count successful delegated inferences that targeted this subagent by name.
pub fn count_successful_delegations_to_subagent(db: &Database, agent_name: &str) -> Result<i64> {
    let conn = db.conn();
    let needle_sp = format!("delegated_subagent={agent_name} ");
    let needle_nl = format!("delegated_subagent={agent_name}\n");
    let count: i64 = conn
        .query_row(
            &format!(
                "SELECT COUNT(*) FROM tool_calls WHERE status = 'success' \
                 AND {} \
                 AND (instr(COALESCE(output,''), ?1) > 0 OR instr(COALESCE(output,''), ?2) > 0)",
                delegation_output_tool_name_sql_in_predicate()
            ),
            rusqlite::params![needle_sp, needle_nl],
            |row| row.get(0),
        )
        .map_err(|e| RoboticusError::Database(format!("count delegations to subagent: {e}")))?;
    Ok(count)
}

pub fn get_subagent_created_at(db: &Database, name: &str) -> Result<Option<String>> {
    let conn = db.conn();
    conn.query_row(
        "SELECT created_at FROM sub_agents WHERE name = ?1",
        [name],
        |row| row.get::<_, String>(0),
    )
    .optional()
    .map_err(|e| RoboticusError::Database(format!("get_subagent_created_at: {e}")))
}

fn subagent_age_days(created_at: Option<&str>) -> Option<i64> {
    let s = created_at?;
    if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) {
        let created = dt.with_timezone(&chrono::Utc).date_naive();
        let now = chrono::Utc::now().date_naive();
        return Some((now - created).num_days());
    }
    if let Ok(nd) = chrono::NaiveDate::parse_from_str(s, "%Y-%m-%d") {
        let now = chrono::Utc::now().date_naive();
        return Some((now - nd).num_days());
    }
    if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") {
        let now = chrono::Utc::now().date_naive();
        return Some((now - ndt.date()).num_days());
    }
    None
}

#[derive(Debug, Clone, Serialize)]
pub struct SubagentRetirementCandidate {
    pub name: String,
    pub role: String,
    pub created_at: Option<String>,
    pub session_count: i64,
    pub delegation_invocations: i64,
    pub age_days: Option<i64>,
    pub meets_age_threshold: bool,
    pub unused: bool,
    /// Ready to soft-retire (`enabled` can be set false): taskable subagent, unused, meets age.
    pub eligible: bool,
    pub reason: String,
}

/// Candidates for soft retirement: enabled **subagent** rows with zero chat sessions for that
/// `agent_id`, zero recorded delegation invocations, and age ≥ `min_age_days` (unless `min_age_days` ≤ 0).
pub fn list_subagent_retirement_candidates(
    db: &Database,
    min_age_days: i32,
) -> Result<Vec<SubagentRetirementCandidate>> {
    let agents = list_sub_agents(db)?;
    let session_counts = list_session_counts_by_agent(db)?;
    let mut out = Vec::new();
    for agent in agents {
        let normalized = agent.role.trim().to_ascii_lowercase();
        if normalized != "subagent" && normalized != "specialist" {
            continue;
        }
        if !agent.enabled {
            continue;
        }
        let session_count = session_counts.get(&agent.name).copied().unwrap_or(0);
        let delegation_invocations = count_successful_delegations_to_subagent(db, &agent.name)?;
        let created_at = get_subagent_created_at(db, &agent.name)?;
        let age_days = subagent_age_days(created_at.as_deref());
        let unused = session_count == 0 && delegation_invocations == 0;
        let meets_age_threshold = if min_age_days <= 0 {
            true
        } else {
            age_days
                .map(|d| d >= i64::from(min_age_days))
                .unwrap_or(false)
        };
        let eligible = unused && meets_age_threshold;
        let reason = if !unused {
            if session_count > 0 && delegation_invocations > 0 {
                format!("has {session_count} session(s) and {delegation_invocations} delegation(s)")
            } else if session_count > 0 {
                format!("has {session_count} chat session(s)")
            } else {
                format!("has {delegation_invocations} delegation invocation(s)")
            }
        } else if !meets_age_threshold {
            format!(
                "below min_age_days ({min_age_days}); age_days={age_days:?}, created_at={created_at:?}"
            )
        } else {
            "unused and meets age threshold".to_string()
        };
        out.push(SubagentRetirementCandidate {
            name: agent.name,
            role: agent.role,
            created_at,
            session_count,
            delegation_invocations,
            age_days,
            meets_age_threshold,
            unused,
            eligible,
            reason,
        });
    }
    out.sort_by(|a, b| a.name.cmp(&b.name));
    Ok(out)
}

/// Disable (`enabled = 0`) named subagents. Only **`subagent` / `specialist`** rows are updated.
/// Returns the number of rows matched by `UPDATE`.
pub fn disable_subagents_by_name(db: &Database, names: &[String]) -> Result<u64> {
    if names.is_empty() {
        return Ok(0);
    }
    let conn = db.conn();
    let mut total = 0u64;
    for name in names {
        let n = conn
            .execute(
                "UPDATE sub_agents SET enabled = 0 WHERE name = ?1 AND enabled = 1 \
                 AND (role = 'subagent' OR role = 'specialist')",
                [name.as_str()],
            )
            .map_err(|e| RoboticusError::Database(format!("disable_subagents_by_name: {e}")))?;
        total += n as u64;
    }
    Ok(total)
}

/// Update a subagent's `last_used_at` timestamp to now.
/// Used by post-turn observer dispatch to keep workspace activity current.
pub fn touch_sub_agent(db: &Database, name: &str) -> Result<()> {
    db.conn()
        .execute(
            "UPDATE sub_agents SET last_used_at = datetime('now') WHERE name = ?1",
            [name],
        )
        .map_err(|e| RoboticusError::Database(format!("touch_sub_agent: {e}")))?;
    Ok(())
}

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

    fn test_db() -> Database {
        Database::new(":memory:").unwrap()
    }

    fn sample_agent(name: &str) -> SubAgentRow {
        SubAgentRow {
            id: uuid::Uuid::new_v4().to_string(),
            name: name.to_string(),
            display_name: Some(name.replace('-', " ")),
            model: "test-model".into(),
            fallback_models_json: Some("[]".into()),
            role: "specialist".into(),
            description: Some("Test agent".into()),
            skills_json: None,
            enabled: true,
            session_count: 0,
            last_used_at: None,
        }
    }

    #[test]
    fn upsert_and_list() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("alpha")).unwrap();
        upsert_sub_agent(&db, &sample_agent("bravo")).unwrap();
        let agents = list_sub_agents(&db).unwrap();
        assert_eq!(agents.len(), 2);
        assert_eq!(agents[0].name, "alpha");
        assert_eq!(agents[1].name, "bravo");
    }

    #[test]
    fn record_subagent_usage_updates_counter_and_recency() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("alpha")).unwrap();

        record_subagent_usage(&db, "alpha").unwrap();

        let agent = list_sub_agents(&db)
            .unwrap()
            .into_iter()
            .find(|agent| agent.name == "alpha")
            .unwrap();
        assert_eq!(agent.session_count, 1);
        assert!(agent.last_used_at.is_some());
    }

    #[test]
    fn upsert_updates_existing() {
        let db = test_db();
        let mut agent = sample_agent("alpha");
        upsert_sub_agent(&db, &agent).unwrap();
        agent.model = "updated-model".into();
        agent.session_count = 42;
        upsert_sub_agent(&db, &agent).unwrap();
        let agents = list_sub_agents(&db).unwrap();
        assert_eq!(agents.len(), 1);
        assert_eq!(agents[0].model, "updated-model");
        assert_eq!(agents[0].session_count, 42);
    }

    #[test]
    fn list_enabled_filters() {
        let db = test_db();
        let mut a = sample_agent("enabled-one");
        upsert_sub_agent(&db, &a).unwrap();
        a = sample_agent("disabled-one");
        a.enabled = false;
        upsert_sub_agent(&db, &a).unwrap();
        let enabled = list_enabled_sub_agents(&db).unwrap();
        assert_eq!(enabled.len(), 1);
        assert_eq!(enabled[0].name, "enabled-one");
    }

    #[test]
    fn delete_works() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("doomed")).unwrap();
        assert!(delete_sub_agent(&db, "doomed").unwrap());
        assert!(!delete_sub_agent(&db, "doomed").unwrap());
        assert!(list_sub_agents(&db).unwrap().is_empty());
    }

    #[test]
    fn session_counts_by_agent_reads_sessions_table() {
        let db = test_db();
        {
            let conn = db.conn();
            conn.execute(
                "INSERT INTO sessions (id, agent_id, scope_key, status) VALUES (?1, ?2, 'agent', 'active')",
                rusqlite::params!["s1", "alpha"],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO sessions (id, agent_id, scope_key, status) VALUES (?1, ?2, 'agent', 'archived')",
                rusqlite::params!["s2", "alpha"],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO sessions (id, agent_id, scope_key, status) VALUES (?1, ?2, 'agent', 'active')",
                rusqlite::params!["s3", "bravo"],
            )
            .unwrap();
        }

        let counts = list_session_counts_by_agent(&db).unwrap();
        assert_eq!(counts.get("alpha"), Some(&2));
        assert_eq!(counts.get("bravo"), Some(&1));
    }

    #[test]
    fn upsert_normalizes_missing_fallback_models() {
        let db = test_db();
        let mut agent = sample_agent("fallback-default");
        agent.fallback_models_json = None;
        upsert_sub_agent(&db, &agent).unwrap();
        let stored = list_sub_agents(&db).unwrap();
        assert_eq!(
            stored[0].fallback_models_json.as_deref(),
            Some("[]"),
            "missing fallback models should normalize to JSON empty array"
        );
    }

    #[test]
    fn count_delegations_detects_successful_delegate_output() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("worker")).unwrap();
        assert_eq!(
            count_successful_delegations_to_subagent(&db, "worker").unwrap(),
            0
        );
        {
            let conn = db.conn();
            conn.execute(
                "INSERT INTO sessions (id, agent_id, scope_key, status) VALUES ('s0', 'orch', 'agent', 'active')",
                [],
            )
            .unwrap();
            conn.execute("INSERT INTO turns (id, session_id) VALUES ('t1', 's0')", [])
                .unwrap();
            conn.execute(
                "INSERT INTO tool_calls (id, turn_id, tool_name, input, output, status, duration_ms) \
                 VALUES ('tc1', 't1', 'delegate-subagent', '{}', 'delegated_subagent=worker model=x\nok', 'success', 1)",
                [],
            )
            .unwrap();
        }
        assert_eq!(
            count_successful_delegations_to_subagent(&db, "worker").unwrap(),
            1
        );
    }

    #[test]
    fn retirement_eligible_when_old_and_never_used() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("dusty")).unwrap();
        {
            let conn = db.conn();
            conn.execute(
                "UPDATE sub_agents SET created_at = '2000-01-01 00:00:00' WHERE name = 'dusty'",
                [],
            )
            .unwrap();
        }
        let c = list_subagent_retirement_candidates(&db, 30).unwrap();
        let dusty = c.iter().find(|x| x.name == "dusty").unwrap();
        assert!(dusty.eligible, "{}", dusty.reason);
    }

    #[test]
    fn retirement_not_eligible_after_delegation_record() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("worker")).unwrap();
        {
            let conn = db.conn();
            conn.execute(
                "UPDATE sub_agents SET created_at = '2000-01-01 00:00:00' WHERE name = 'worker'",
                [],
            )
            .unwrap();
            conn.execute(
                "INSERT INTO sessions (id, agent_id, scope_key, status) VALUES ('s0', 'orch', 'agent', 'active')",
                [],
            )
            .unwrap();
            conn.execute("INSERT INTO turns (id, session_id) VALUES ('t1', 's0')", [])
                .unwrap();
            conn.execute(
                "INSERT INTO tool_calls (id, turn_id, tool_name, input, output, status, duration_ms) \
                 VALUES ('tc1', 't1', 'delegate-subagent', '{}', 'delegated_subagent=worker model=x\n', 'success', 1)",
                [],
            )
            .unwrap();
        }
        let c = list_subagent_retirement_candidates(&db, 0).unwrap();
        let w = c.iter().find(|x| x.name == "worker").unwrap();
        assert!(!w.eligible);
    }

    #[test]
    fn disable_subagents_by_name_matches_upsert_enabled_false() {
        let db = test_db();
        let a = sample_agent("parity");
        upsert_sub_agent(&db, &a).unwrap();
        assert!(
            list_sub_agents(&db)
                .unwrap()
                .iter()
                .find(|r| r.name == "parity")
                .unwrap()
                .enabled
        );

        disable_subagents_by_name(&db, &["parity".into()]).unwrap();

        let after_disable = list_sub_agents(&db).unwrap();
        let row1 = after_disable.iter().find(|r| r.name == "parity").unwrap();
        assert!(!row1.enabled);

        // Same end state as explicit upsert — retirement + admin paths must stay equivalent.
        let mut b = sample_agent("parity2");
        upsert_sub_agent(&db, &b).unwrap();
        b.enabled = false;
        upsert_sub_agent(&db, &b).unwrap();
        let agents_b = list_sub_agents(&db).unwrap();
        let row2 = agents_b.iter().find(|r| r.name == "parity2").unwrap();
        assert!(!row2.enabled);
        assert_eq!(row1.enabled, row2.enabled);
    }

    #[test]
    fn disable_subagents_by_name_updates_row() {
        let db = test_db();
        upsert_sub_agent(&db, &sample_agent("gone")).unwrap();
        assert_eq!(disable_subagents_by_name(&db, &["gone".into()]).unwrap(), 1);
        let agents = list_sub_agents(&db).unwrap();
        assert!(!agents[0].enabled);
    }
}