orchai/src-tauri/src/models/worktree.rs

use rusqlite::{params, Connection, Result};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Worktree {
    pub id: String,
    pub ticket_id: String,
    pub path: String,
    pub branch_name: String,
    pub status: String,
    pub created_at: String,
    pub merged_at: Option<String>,
    pub merged_into: Option<String>,
}

fn from_row(row: &rusqlite::Row) -> rusqlite::Result<Worktree> {
    Ok(Worktree {
        id: row.get(0)?,
        ticket_id: row.get(1)?,
        path: row.get(2)?,
        branch_name: row.get(3)?,
        status: row.get(4)?,
        created_at: row.get(5)?,
        merged_at: row.get(6)?,
        merged_into: row.get(7)?,
    })
}

const SELECT_ALL_COLS: &str = "SELECT id, ticket_id, path, branch_name, status, \
    created_at, merged_at, merged_into FROM worktrees";

impl Worktree {
    pub fn insert(conn: &Connection, ticket_id: &str, path: &str, branch_name: &str) -> Result<Worktree> {
        let id = Uuid::new_v4().to_string();
        let now = chrono::Utc::now().to_rfc3339();

        conn.execute(
            "INSERT INTO worktrees (id, ticket_id, path, branch_name, status, created_at) \
             VALUES (?1, ?2, ?3, ?4, 'Active', ?5)",
            params![id, ticket_id, path, branch_name, now],
        )?;

        Ok(Worktree {
            id,
            ticket_id: ticket_id.to_string(),
            path: path.to_string(),
            branch_name: branch_name.to_string(),
            status: "Active".to_string(),
            created_at: now,
            merged_at: None,
            merged_into: None,
        })
    }

    pub fn get_by_id(conn: &Connection, id: &str) -> Result<Worktree> {
        let sql = format!("{} WHERE id = ?1", SELECT_ALL_COLS);
        conn.query_row(&sql, params![id], from_row)
    }

    pub fn get_by_ticket_id(conn: &Connection, ticket_id: &str) -> Result<Option<Worktree>> {
        let sql = format!("{} WHERE ticket_id = ?1", SELECT_ALL_COLS);
        let mut stmt = conn.prepare(&sql)?;
        let mut rows = stmt.query_map(params![ticket_id], from_row)?;
        match rows.next() {
            Some(Ok(w)) => Ok(Some(w)),
            Some(Err(e)) => Err(e),
            None => Ok(None),
        }
    }

    pub fn list_by_project(conn: &Connection, project_id: &str) -> Result<Vec<Worktree>> {
        let sql = "SELECT w.id, w.ticket_id, w.path, w.branch_name, w.status, \
             w.created_at, w.merged_at, w.merged_into \
             FROM worktrees w \
             JOIN processed_tickets pt ON w.ticket_id = pt.id \
             JOIN watched_trackers wt ON pt.tracker_id = wt.id \
             WHERE wt.project_id = ?1 \
             ORDER BY w.created_at DESC";
        let mut stmt = conn.prepare(sql)?;
        let rows = stmt.query_map(params![project_id], from_row)?;
        rows.collect()
    }

    pub fn set_merged(conn: &Connection, id: &str, target_branch: &str) -> Result<()> {
        conn.execute(
            "UPDATE worktrees SET status = 'Merged', merged_at = datetime('now'), merged_into = ?1 WHERE id = ?2",
            params![target_branch, id],
        )?;
        Ok(())
    }

    pub fn delete(conn: &Connection, id: &str) -> Result<()> {
        conn.execute("DELETE FROM worktrees WHERE id = ?1", params![id])?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::db;
    use crate::models::agent::{Agent, AgentRole, AgentTool};
    use crate::models::project::Project;
    use crate::models::ticket::ProcessedTicket;
    use crate::models::tracker::WatchedTracker;

    fn setup() -> (Connection, String) {
        let conn = db::init_in_memory().expect("db init");
        let project = Project::insert(&conn, "Test", "/path", None, "main").unwrap();
        let analyst = Agent::insert(&conn, "A", AgentRole::Analyst, AgentTool::Codex, "").unwrap();
        let developer =
            Agent::insert(&conn, "D", AgentRole::Developer, AgentTool::ClaudeCode, "").unwrap();
        let tracker = WatchedTracker::insert(
            &conn,
            &project.id,
            100,
            "Bugs",
            10,
            &analyst.id,
            &developer.id,
            vec![],
        )
        .unwrap();
        let ticket = ProcessedTicket::insert_if_new(&conn, &tracker.id, 42, "Bug 42", "{}")
            .unwrap()
            .unwrap();
        (conn, ticket.id)
    }

    #[test]
    fn test_insert_and_get_by_id() {
        let (conn, ticket_id) = setup();

        let wt = Worktree::insert(&conn, &ticket_id, "/tmp/orchai-42", "orchai/42").unwrap();
        assert_eq!(wt.status, "Active");
        assert_eq!(wt.branch_name, "orchai/42");

        let found = Worktree::get_by_id(&conn, &wt.id).unwrap();
        assert_eq!(found.id, wt.id);
        assert_eq!(found.ticket_id, ticket_id);
        assert_eq!(found.path, "/tmp/orchai-42");
    }

    #[test]
    fn test_get_by_ticket_id() {
        let (conn, ticket_id) = setup();

        let none = Worktree::get_by_ticket_id(&conn, &ticket_id).unwrap();
        assert!(none.is_none());

        Worktree::insert(&conn, &ticket_id, "/tmp/wt", "orchai/42").unwrap();

        let some = Worktree::get_by_ticket_id(&conn, &ticket_id).unwrap();
        assert!(some.is_some());
        assert_eq!(some.unwrap().ticket_id, ticket_id);
    }

    #[test]
    fn test_list_by_project() {
        let conn = db::init_in_memory().expect("db init");
        let project = Project::insert(&conn, "P1", "/path", None, "main").unwrap();
        let analyst = Agent::insert(&conn, "A", AgentRole::Analyst, AgentTool::Codex, "").unwrap();
        let developer =
            Agent::insert(&conn, "D", AgentRole::Developer, AgentTool::ClaudeCode, "").unwrap();
        let tracker = WatchedTracker::insert(
            &conn,
            &project.id,
            100,
            "Bugs",
            10,
            &analyst.id,
            &developer.id,
            vec![],
        )
        .unwrap();
        let t1 = ProcessedTicket::insert_if_new(&conn, &tracker.id, 1, "T1", "{}")
            .unwrap()
            .unwrap();
        let t2 = ProcessedTicket::insert_if_new(&conn, &tracker.id, 2, "T2", "{}")
            .unwrap()
            .unwrap();

        Worktree::insert(&conn, &t1.id, "/wt1", "orchai/1").unwrap();
        Worktree::insert(&conn, &t2.id, "/wt2", "orchai/2").unwrap();

        let worktrees = Worktree::list_by_project(&conn, &project.id).unwrap();
        assert_eq!(worktrees.len(), 2);
    }

    #[test]
    fn test_set_merged() {
        let (conn, ticket_id) = setup();
        let wt = Worktree::insert(&conn, &ticket_id, "/tmp/wt", "orchai/42").unwrap();

        Worktree::set_merged(&conn, &wt.id, "feature/login").unwrap();
        let updated = Worktree::get_by_id(&conn, &wt.id).unwrap();
        assert_eq!(updated.status, "Merged");
        assert_eq!(updated.merged_into.unwrap(), "feature/login");
        assert!(updated.merged_at.is_some());
    }

    #[test]
    fn test_delete() {
        let (conn, ticket_id) = setup();
        let wt = Worktree::insert(&conn, &ticket_id, "/tmp/wt", "orchai/42").unwrap();

        Worktree::delete(&conn, &wt.id).unwrap();
        let result = Worktree::get_by_id(&conn, &wt.id);
        assert!(result.is_err());
    }
}