gitdataai/lib/service/agent/workspace_tools/workspace.rs

use ai::error::{AiError, AiResult};
use ai::tool::tools::FunctionCall;
use async_trait::async_trait;
use db::sqlx;
use serde_json::{Value, json};
use uuid::Uuid;

use super::helpers::{arg_str, git_ctx, require_workspace_member};
use crate::agent::run::AppAgentContext;

pub struct WorkspaceInfoTool;

impl WorkspaceInfoTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for WorkspaceInfoTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for WorkspaceInfoTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "workspace_info"
    }

    fn description(&self) -> &'static str {
        "Get information about a workspace: name, description, avatar."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" }
            },
            "required": ["workspace"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let _wk_id =
            require_workspace_member(git, ctx.user_id, workspace).await?;

        let row = sqlx::query_as::<_, (String, String, String, chrono::DateTime<chrono::Utc>)>(
            "SELECT name, description, avatar_url, created_at FROM workspace WHERE name = $1",
        )
        .bind(workspace)
        .fetch_optional(git.db.reader())
        .await
        .map_err(AiError::Database)?
        .ok_or_else(|| AiError::Config(format!("workspace '{workspace}' not found")))?;

        let member_count: i64 = sqlx::query_scalar(
            "SELECT COUNT(*) FROM wk_member WHERE wk = (SELECT id FROM workspace WHERE name = $1) AND leave_at IS NULL",
        )
        .bind(workspace)
        .fetch_one(git.db.reader())
        .await
        .map_err(AiError::Database)?;

        Ok(json!({
            "name": row.0,
            "description": row.1,
            "avatar_url": row.2,
            "created_at": row.3.to_rfc3339(),
            "member_count": member_count,
        }))
    }
}

pub struct WorkspaceMembersTool;

impl WorkspaceMembersTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for WorkspaceMembersTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for WorkspaceMembersTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "workspace_members"
    }

    fn description(&self) -> &'static str {
        "List all members of a workspace with their roles."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "limit": { "type": "integer", "description": "Max results (default 50)" }
            },
            "required": ["workspace"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let wk_id =
            require_workspace_member(git, ctx.user_id, workspace).await?;
        let limit = args
            .get("limit")
            .and_then(|v| v.as_i64())
            .unwrap_or(50)
            .min(200);

        #[derive(sqlx::FromRow)]
        struct MemberRow {
            username: String,
            display_name: Option<String>,
            owner: bool,
            admin: bool,
        }

        let rows: Vec<MemberRow> = sqlx::query_as(
            "SELECT u.username, u.display_name, m.owner, m.admin \
             FROM wk_member m \
             INNER JOIN \"user\" u ON u.id = m.\"user\" \
             WHERE m.wk = $1 AND m.leave_at IS NULL \
             ORDER BY m.owner DESC, m.admin DESC \
             LIMIT $2",
        )
        .bind(wk_id)
        .bind(limit)
        .fetch_all(git.db.reader())
        .await
        .map_err(AiError::Database)?;

        let members: Vec<Value> = rows.iter().map(|r| json!({
            "username": r.username,
            "display_name": r.display_name,
            "role": if r.owner { "owner" } else if r.admin { "admin" } else { "member" },
        })).collect();

        Ok(json!({ "members": members, "count": members.len() }))
    }
}

pub struct WorkspaceGroupsTool;

impl WorkspaceGroupsTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for WorkspaceGroupsTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for WorkspaceGroupsTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "workspace_groups"
    }

    fn description(&self) -> &'static str {
        "List all user groups in a workspace."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" }
            },
            "required": ["workspace"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let wk_id =
            require_workspace_member(git, ctx.user_id, workspace).await?;

        #[derive(sqlx::FromRow)]
        struct GroupRow {
            id: Uuid,
            name: String,
            created_at: chrono::DateTime<chrono::Utc>,
        }

        let rows: Vec<GroupRow> = sqlx::query_as(
            "SELECT id, name, created_at FROM wk_group \
             WHERE wk = $1 AND is_deleted = FALSE ORDER BY name ASC",
        )
        .bind(wk_id)
        .fetch_all(git.db.reader())
        .await
        .map_err(AiError::Database)?;

        // For each group, count members
        let mut groups: Vec<Value> = Vec::new();
        for row in &rows {
            let count: i64 = sqlx::query_scalar(
                "SELECT COUNT(*) FROM wk_gp_member WHERE gp = $1",
            )
            .bind(row.id)
            .fetch_one(git.db.reader())
            .await
            .map_err(AiError::Database)?;

            groups.push(json!({
                "name": row.name,
                "member_count": count,
                "created_at": row.created_at.to_rfc3339(),
            }));
        }

        Ok(json!({ "groups": groups, "count": groups.len() }))
    }
}

pub struct WorkspaceGroupMembersTool;

impl WorkspaceGroupMembersTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for WorkspaceGroupMembersTool {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl FunctionCall for WorkspaceGroupMembersTool {
    type Context = AppAgentContext;

    fn name(&self) -> &'static str {
        "workspace_group_members"
    }

    fn description(&self) -> &'static str {
        "List members of a specific workspace group."
    }

    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "workspace": { "type": "string", "description": "Workspace name" },
                "group_name": { "type": "string", "description": "Group name" }
            },
            "required": ["workspace", "group_name"]
        })
    }

    async fn call(
        &self,
        ctx: &mut AppAgentContext,
        args: Value,
    ) -> AiResult<Value> {
        let git = git_ctx(ctx)?;
        let workspace = arg_str(&args, "workspace")?;
        let group_name = arg_str(&args, "group_name")?;
        let wk_id =
            require_workspace_member(git, ctx.user_id, workspace).await?;

        #[derive(sqlx::FromRow)]
        struct MemberRow {
            username: String,
            display_name: Option<String>,
        }

        let rows: Vec<MemberRow> = sqlx::query_as(
            "SELECT u.username, u.display_name \
             FROM wk_gp_member gm \
             INNER JOIN wk_group g ON g.id = gm.gp \
             INNER JOIN \"user\" u ON u.id = gm.\"user\" \
             WHERE g.wk = $1 AND g.name = $2 AND g.is_deleted = FALSE",
        )
        .bind(wk_id)
        .bind(group_name)
        .fetch_all(git.db.reader())
        .await
        .map_err(AiError::Database)?;

        let members: Vec<Value> = rows
            .iter()
            .map(|r| {
                json!({
                    "username": r.username,
                    "display_name": r.display_name,
                })
            })
            .collect();

        Ok(
            json!({ "group": group_name, "members": members, "count": members.len() }),
        )
    }
}