mcp-postgres 4.0.2

use serde_json::{json, Value};
use tokio_postgres::Client;
use crate::errors::Result as MCPResult;

const MAX_IDENTIFIER_LEN: usize = 255;

pub async fn create_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;

    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
    let valid_until = params.as_ref().and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
    let connection_limit = params.as_ref().and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));

    let mut sql = format!("CREATE USER {}", quote_ident(username));
    if let Some(pw) = password {
        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
    }
    if let Some(limit) = connection_limit {
        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
    }
    if let Some(login) = can_login {
        if login {
            sql.push_str(" LOGIN");
        } else {
            sql.push_str(" NOLOGIN");
        }
    }
    if let Some(until) = valid_until {
        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
    }

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "user": username }))
}

pub async fn alter_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;

    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
    let valid_until = params.as_ref().and_then(|p| p.get("valid_until").and_then(|v| v.as_str()));
    let connection_limit = params.as_ref().and_then(|p| p.get("connection_limit").and_then(|v| v.as_i64()));
    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
    let new_name = params.as_ref().and_then(|p| p.get("new_name").and_then(|v| v.as_str()));

    if password.is_none() && valid_until.is_none() && connection_limit.is_none() && can_login.is_none() && new_name.is_none() {
        return Err(crate::errors::MCPError::InvalidParams("No attributes specified to alter".into()));
    }

    let mut sql = format!("ALTER USER {}", quote_ident(username));
    if let Some(pw) = password {
        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
    }
    if let Some(limit) = connection_limit {
        sql.push_str(&format!(" CONNECTION LIMIT {}", limit));
    }
    if let Some(login) = can_login {
        if login { sql.push_str(" LOGIN"); } else { sql.push_str(" NOLOGIN"); }
    }
    if let Some(until) = valid_until {
        sql.push_str(&format!(" VALID UNTIL '{}'", until.replace('\'', "''")));
    }
    if let Some(name) = new_name {
        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
    }

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "user": username }))
}

pub async fn drop_user(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let username = params.as_ref().and_then(|p| p.get("username").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'username' parameter".into()))?;

    if username.is_empty() || username.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'username' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let if_exists = params.as_ref().and_then(|p| p.get("if_exists").and_then(|v| v.as_bool())).unwrap_or(false);

    let mut sql = "DROP USER".to_string();
    if if_exists { sql.push_str(" IF EXISTS"); }
    sql.push_str(&format!(" {}", quote_ident(username)));

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "user": username }))
}

pub async fn create_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;

    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let with_login = params.as_ref().and_then(|p| p.get("with_login").and_then(|v| v.as_bool())).unwrap_or(false);
    let mut sql = format!("CREATE ROLE {}", quote_ident(role_name));
    sql.push_str(if with_login { " LOGIN" } else { " NOLOGIN" });

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "role": role_name }))
}

pub async fn alter_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;

    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let password = params.as_ref().and_then(|p| p.get("password").and_then(|v| v.as_str()));
    let can_login = params.as_ref().and_then(|p| p.get("can_login").and_then(|v| v.as_bool()));
    let superuser = params.as_ref().and_then(|p| p.get("superuser").and_then(|v| v.as_bool()));
    let createdb = params.as_ref().and_then(|p| p.get("createdb").and_then(|v| v.as_bool()));
    let new_name = params.as_ref().and_then(|p| p.get("new_name").and_then(|v| v.as_str()));

    let mut sql = format!("ALTER ROLE {}", quote_ident(role_name));
    if let Some(pw) = password {
        sql.push_str(&format!(" PASSWORD '{}'", pw.replace('\'', "''")));
    }
    if let Some(login) = can_login {
        sql.push_str(if login { " LOGIN" } else { " NOLOGIN" });
    }
    if let Some(su) = superuser {
        sql.push_str(if su { " SUPERUSER" } else { " NOSUPERUSER" });
    }
    if let Some(db) = createdb {
        sql.push_str(if db { " CREATEDB" } else { " NOCREATEDB" });
    }
    if let Some(name) = new_name {
        sql.push_str(&format!(" RENAME TO {}", quote_ident(name)));
    }

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "role": role_name }))
}

pub async fn drop_role(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let role_name = params.as_ref().and_then(|p| p.get("role_name").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'role_name' parameter".into()))?;

    if role_name.is_empty() || role_name.len() > MAX_IDENTIFIER_LEN {
        return Err(crate::errors::MCPError::InvalidParams(format!("'role_name' must be 1-{MAX_IDENTIFIER_LEN} characters")));
    }

    let if_exists = params.as_ref().and_then(|p| p.get("if_exists").and_then(|v| v.as_bool())).unwrap_or(false);

    let mut sql = "DROP ROLE".to_string();
    if if_exists { sql.push_str(" IF EXISTS"); }
    sql.push_str(&format!(" {}", quote_ident(role_name)));

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "role": role_name }))
}

pub async fn grant_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let privilege = params.as_ref().and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into()))?;
    let object_type = params.as_ref().and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into()))?;
    let object_name = params.as_ref().and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into()))?;
    let grantee = params.as_ref().and_then(|p| p.get("grantee").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'grantee' parameter".into()))?;

    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");

    let valid_types = ["table", "sequence", "schema", "database", "all_tables_in_schema"];
    if !valid_types.contains(&object_type) {
        return Err(crate::errors::MCPError::InvalidParams(format!("Unsupported object_type '{}'. Use: {:?}", object_type, valid_types)));
    }

    let sql = match object_type {
        "all_tables_in_schema" => format!("GRANT {} ON ALL TABLES IN SCHEMA {} TO {}", privilege, quote_ident(schema), quote_ident(grantee)),
        _ => format!("GRANT {} ON {} {} TO {}", privilege, object_type.to_uppercase(), quote_ident(object_name), quote_ident(grantee)),
    };

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "sql": sql }))
}

pub async fn revoke_privileges(client: &Client, params: &Option<&Value>) -> MCPResult<Value> {
    let privilege = params.as_ref().and_then(|p| p.get("privilege").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'privilege' parameter".into()))?;
    let object_type = params.as_ref().and_then(|p| p.get("object_type").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_type' parameter".into()))?;
    let object_name = params.as_ref().and_then(|p| p.get("object_name").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'object_name' parameter".into()))?;
    let revokee = params.as_ref().and_then(|p| p.get("revokee").and_then(|v| v.as_str()))
        .ok_or_else(|| crate::errors::MCPError::InvalidParams("Missing 'revokee' parameter".into()))?;

    let schema = params.as_ref().and_then(|p| p.get("schema").and_then(|v| v.as_str())).unwrap_or("public");

    let valid_types = ["table", "sequence", "schema", "database", "all_tables_in_schema"];
    if !valid_types.contains(&object_type) {
        return Err(crate::errors::MCPError::InvalidParams(format!("Unsupported object_type '{}'. Use: {:?}", object_type, valid_types)));
    }

    let sql = match object_type {
        "all_tables_in_schema" => format!("REVOKE {} ON ALL TABLES IN SCHEMA {} FROM {}", privilege, quote_ident(schema), quote_ident(revokee)),
        _ => format!("REVOKE {} ON {} {} FROM {}", privilege, object_type.to_uppercase(), quote_ident(object_name), quote_ident(revokee)),
    };

    client.execute(&sql, &[]).await?;
    Ok(json!({ "success": true, "sql": sql }))
}

fn quote_ident(ident: &str) -> String {
    crate::validation::quote_ident(ident)
}