pg-api 0.1.0

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tokio_postgres::{Client, NoTls};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatabaseInfo {
    pub name: String,
    pub owner: String,
    pub encoding: String,
    pub size: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserInfo {
    pub username: String,
    pub is_superuser: bool,
    pub can_create_db: bool,
    pub can_create_role: bool,
    pub databases: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PostgresIntrospection {
    pub databases: Vec<DatabaseInfo>,
    pub users: Vec<UserInfo>,
    pub version: String,
}

pub async fn introspect_postgres(connection_string: &str) -> Result<PostgresIntrospection> {
    let (client, connection) = tokio_postgres::connect(connection_string, NoTls).await?;
    
    tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("connection error: {e}");
        }
    });

    let version = get_postgres_version(&client).await?;
    let databases = get_databases(&client).await?;
    let users = get_users(&client).await?;

    Ok(PostgresIntrospection {
        databases,
        users,
        version,
    })
}

async fn get_postgres_version(client: &Client) -> Result<String> {
    let row = client
        .query_one("SELECT version()", &[])
        .await?;
    
    let version: String = row.get(0);
    Ok(version)
}

async fn get_databases(client: &Client) -> Result<Vec<DatabaseInfo>> {
    let query = r#"
        SELECT 
            datname as name,
            pg_catalog.pg_get_userbyid(datdba) as owner,
            pg_catalog.pg_encoding_to_char(encoding) as encoding,
            pg_catalog.pg_database_size(datname) as size
        FROM pg_catalog.pg_database
        WHERE datistemplate = false
        ORDER BY datname
    "#;
    
    let rows = client.query(query, &[]).await?;
    
    let mut databases = Vec::new();
    for row in rows {
        let size: i64 = row.get("size");
        databases.push(DatabaseInfo {
            name: row.get("name"),
            owner: row.get("owner"),
            encoding: row.get("encoding"),
            size: Some(format_bytes(size)),
        });
    }
    
    Ok(databases)
}

async fn get_users(client: &Client) -> Result<Vec<UserInfo>> {
    let query = r#"
        SELECT 
            usename as username,
            usesuper as is_superuser,
            usecreatedb as can_create_db,
            usecreaterole as can_create_role
        FROM pg_catalog.pg_user
        ORDER BY usename
    "#;
    
    let rows = client.query(query, &[]).await?;
    
    let mut users = Vec::new();
    for row in rows {
        let username: String = row.get("username");
        
        // Get databases this user has permissions on
        let db_query = r#"
            SELECT DISTINCT datname 
            FROM pg_database d
            WHERE d.datdba = (SELECT usesysid FROM pg_user WHERE usename = $1)
            OR EXISTS (
                SELECT 1 FROM pg_catalog.pg_namespace n
                JOIN pg_catalog.pg_auth_members m ON (m.member = (SELECT usesysid FROM pg_user WHERE usename = $1))
                WHERE n.nspname = 'public'
            )
            ORDER BY datname
        "#;
        
        let db_rows = client.query(db_query, &[&username]).await?;
        let databases: Vec<String> = db_rows.iter().map(|r| r.get(0)).collect();
        
        users.push(UserInfo {
            username,
            is_superuser: row.get("is_superuser"),
            can_create_db: row.get("can_create_db"),
            can_create_role: row.get("can_create_role"),
            databases,
        });
    }
    
    Ok(users)
}

fn format_bytes(bytes: i64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut size = bytes as f64;
    let mut unit_index = 0;
    
    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
        size /= 1024.0;
        unit_index += 1;
    }
    
    format!("{:.2} {}", size, UNITS[unit_index])
}

pub async fn introspect_with_sudo() -> Result<PostgresIntrospection> {
    // Try to connect using sudo -u postgres
    let connection_string = "host=/var/run/postgresql user=postgres dbname=postgres";
    introspect_postgres(connection_string).await
}

#[allow(dead_code)]
pub async fn get_database_permissions(client: &Client, database: &str, username: &str) -> Result<Vec<String>> {
    // This is a simplified version - in a real implementation, you'd want to check:
    // - Schema permissions
    // - Table permissions
    // - Function permissions
    // - etc.
    
    let query = r#"
        SELECT has_database_privilege($1, $2, 'CONNECT') as can_connect,
               has_database_privilege($1, $2, 'CREATE') as can_create,
               has_database_privilege($1, $2, 'TEMP') as can_temp
    "#;
    
    let row = client.query_one(query, &[&username, &database]).await?;
    
    let mut permissions = Vec::new();
    if row.get::<_, bool>("can_connect") {
        permissions.push("CONNECT".to_string());
    }
    if row.get::<_, bool>("can_create") {
        permissions.push("CREATE".to_string());
    }
    if row.get::<_, bool>("can_temp") {
        permissions.push("TEMP".to_string());
    }
    
    Ok(permissions)
}