pg-api 0.3.4

use axum::{
    extract::{Json, Path, State},
    http::{HeaderMap, StatusCode, header},
    response::{IntoResponse, Html, Response},
    Extension,
};
use chrono::Utc;
use serde_json::{json, Value};
use std::time::Instant;
use uuid::Uuid;

use crate::{
    database::{
        execute_query_with_pool, execute_transaction_with_pool, get_or_create_pool,
        validate_query_permissions, MAX_BATCH_ITEMS,
    },
    models::*,
};

/// Map gate/validator error prefixes to stable API codes (clients match on these).
fn gate_error_code(msg: &str) -> &'static str {
    if msg.starts_with("MULTI_STATEMENT") {
        "MULTI_STATEMENT_NOT_ALLOWED"
    } else if msg.starts_with("STATEMENT_NOT_ALLOWED") {
        "STATEMENT_NOT_ALLOWED"
    } else if msg.starts_with("STATEMENT_TIMEOUT") {
        "STATEMENT_TIMEOUT"
    } else if msg.starts_with("MISSING_PERMISSION") {
        "PERMISSION_DENIED"
    } else {
        "PERMISSION_DENIED"
    }
}

pub async fn health_check() -> impl IntoResponse {
    Json(json!({
        "status": "healthy",
        "service": env!("CARGO_PKG_NAME"),
        "version": env!("CARGO_PKG_VERSION")
    }))
}

pub async fn status_handler(State(state): State<AppState>) -> impl IntoResponse {
    let instances = state.instances.read().await;
    let connections = state.connections.len();
    
    Json(json!({
        "instances": instances.len(),
        "active_connections": connections,
        "status": "operational"
    }))
}

pub async fn query_handler(
    State(state): State<AppState>,
    Extension(account): Extension<Account>,
    headers: HeaderMap,
    Json(request): Json<QueryRequest>,
) -> Result<Json<ApiResponse<QueryResult>>, StatusCode> {
    let start = Instant::now();
    let request_id = headers
        .get("x-request-id")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("unknown")
        .to_string();
    
    let mut metadata = ResponseMetadata {
        request_id: request_id.clone(),
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    // Check database access
    let db_access = account.databases.iter()
        .find(|db| db.database == request.database)
        .ok_or(StatusCode::FORBIDDEN)?;
    
    // Validate query permissions
    if let Err(e) = validate_query_permissions(&request.query, &db_access.permissions, &account.role) {
        metadata.execution_time_ms = start.elapsed().as_millis();
        let code = gate_error_code(&e);
        return Ok(Json(ApiResponse::error(code, e, metadata)));
    }
    
    // Get connection pool
    let pool_key = format!("{}-{}-{}", account.instance_id, db_access.username, request.database);
    let pool = match get_or_create_pool(&state, &pool_key, &account, db_access).await {
        Ok(p) => p,
        Err(e) => {
            metadata.execution_time_ms = start.elapsed().as_millis();
            return Ok(Json(ApiResponse::error("CONNECTION_ERROR", e, metadata)));
        }
    };
    
    // Execute query (options now enforced: statement_timeout, read_only)
    match execute_query_with_pool(pool, request.query, request.params, &request.options).await {
        Ok(result) => {
            metadata.execution_time_ms = start.elapsed().as_millis();
            // v0.3.4: usa linhas afetadas reais quando o executor informa
            // (escritas sem RETURNING); senao deriva de rows.len().
            metadata.rows_affected =
                Some(result.rows_affected.unwrap_or(result.rows.len() as u64));

            Ok(Json(ApiResponse::success(result, metadata)))
        }
        Err(e) => {
            metadata.execution_time_ms = start.elapsed().as_millis();
            let code = if e.starts_with("STATEMENT_TIMEOUT") {
                "STATEMENT_TIMEOUT"
            } else {
                "QUERY_ERROR"
            };
            Ok(Json(ApiResponse::error(code, e, metadata)))
        }
    }
}

pub async fn batch_query_handler(
    State(state): State<AppState>,
    Extension(account): Extension<Account>,
    headers: HeaderMap,
    Json(requests): Json<Vec<QueryRequest>>,
) -> Result<Json<ApiResponse<Vec<QueryResult>>>, StatusCode> {
    let start = Instant::now();
    let request_id = Uuid::new_v4().to_string();

    // v0.3: bound the batch so one request cannot fan out unbounded work.
    if requests.len() > MAX_BATCH_ITEMS {
        let metadata = ResponseMetadata {
            request_id,
            execution_time_ms: start.elapsed().as_millis(),
            rows_affected: None,
            instance_id: Some(account.instance_id.clone()),
            timestamp: Utc::now(),
        };
        return Ok(Json(ApiResponse::error(
            "STATEMENT_NOT_ALLOWED",
            format!("batch exceeds {MAX_BATCH_ITEMS} statements"),
            metadata,
        )));
    }
    if requests.is_empty() {
        let metadata = ResponseMetadata {
            request_id,
            execution_time_ms: start.elapsed().as_millis(),
            rows_affected: None,
            instance_id: Some(account.instance_id.clone()),
            timestamp: Utc::now(),
        };
        return Ok(Json(ApiResponse::error(
            "STATEMENT_NOT_ALLOWED",
            "empty batch".to_string(),
            metadata,
        )));
    }
    
    let mut results = Vec::new();
    let mut total_rows = 0u64;
    
    for request in requests {
        match query_handler(
            State(state.clone()),
            Extension(account.clone()),
            headers.clone(),
            Json(request),
        ).await {
            Ok(Json(response)) => {
                if let Some(data) = response.data {
                    total_rows += data.rows_affected.unwrap_or(data.rows.len() as u64);
                    results.push(data);
                }
            }
            Err(_) => {
                // Handle error
            }
        }
    }
    
    let metadata = ResponseMetadata {
        request_id,
        execution_time_ms: start.elapsed().as_millis(),
        rows_affected: Some(total_rows),
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    Ok(Json(ApiResponse::success(results, metadata)))
}

pub async fn transaction_handler(
    State(state): State<AppState>,
    Extension(account): Extension<Account>,
    Json(requests): Json<Vec<QueryRequest>>,
) -> Result<Json<ApiResponse<Vec<QueryResult>>>, StatusCode> {
    // v0.3: real atomic transaction on a single pooled connection.
    // Every statement is permission-checked up front; any failure rolls back.
    let start = Instant::now();
    let request_id = Uuid::new_v4().to_string();
    let mkmeta = |rows_affected: Option<u64>| ResponseMetadata {
        request_id: request_id.clone(),
        execution_time_ms: start.elapsed().as_millis(),
        rows_affected,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };

    if requests.is_empty() || requests.len() > MAX_BATCH_ITEMS {
        return Ok(Json(ApiResponse::error(
            "STATEMENT_NOT_ALLOWED",
            format!("transaction must contain 1..={MAX_BATCH_ITEMS} statements"),
            mkmeta(None),
        )));
    }

    // All statements in a transaction must target the same database.
    let database = requests[0].database.clone();
    if requests.iter().any(|r| r.database != database) {
        return Ok(Json(ApiResponse::error(
            "STATEMENT_NOT_ALLOWED",
            "all transaction statements must target the same database".to_string(),
            mkmeta(None),
        )));
    }

    let db_access = match account
        .databases
        .iter()
        .find(|db| db.database == database)
    {
        Some(db) => db,
        None => return Err(StatusCode::FORBIDDEN),
    };

    for req in &requests {
        if let Err(e) = validate_query_permissions(&req.query, &db_access.permissions, &account.role)
        {
            let code = gate_error_code(&e);
            return Ok(Json(ApiResponse::error(code, e, mkmeta(None))));
        }
    }

    let pool_key = format!("{}-{}-{}", account.instance_id, db_access.username, database);
    let pool = match get_or_create_pool(&state, &pool_key, &account, db_access).await {
        Ok(p) => p,
        Err(e) => {
            return Ok(Json(ApiResponse::error("CONNECTION_ERROR", e, mkmeta(None))));
        }
    };

    let stmts: Vec<(String, Vec<serde_json::Value>, QueryOptions)> = requests
        .into_iter()
        .map(|r| (r.query, r.params, r.options))
        .collect();
    match execute_transaction_with_pool(pool, stmts).await {
        Ok(results) => {
            let total = results
                .iter()
                .map(|r| r.rows_affected.unwrap_or(r.rows.len() as u64))
                .sum();
            Ok(Json(ApiResponse::success(results, mkmeta(Some(total)))))
        }
        Err(e) => {
            let code = if e.starts_with("STATEMENT_TIMEOUT") {
                "STATEMENT_TIMEOUT"
            } else {
                "QUERY_ERROR"
            };
            Ok(Json(ApiResponse::error(code, e, mkmeta(None))))
        }
    }
}

pub async fn list_databases(
    State(_state): State<AppState>,
    Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Vec<String>>>, StatusCode> {
    let request_id = Uuid::new_v4().to_string();
    let databases: Vec<String> = account.databases.iter()
        .map(|db| db.database.clone())
        .collect();
    
    let metadata = ResponseMetadata {
        request_id,
        execution_time_ms: 0,
        rows_affected: Some(databases.len() as u64),
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    Ok(Json(ApiResponse::success(databases, metadata)))
}

pub async fn create_database(
    State(_state): State<AppState>,
    Extension(account): Extension<Account>,
    Json(_request): Json<serde_json::Map<String, Value>>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
    let request_id = Uuid::new_v4().to_string();
    let metadata = ResponseMetadata {
        request_id,
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    if account.role != AccountRole::Superuser && account.role != AccountRole::Admin {
        return Ok(Json(ApiResponse::error(
            "PERMISSION_DENIED",
            "Only Owner or Admin can create databases".to_string(),
            metadata,
        )));
    }

    // v0.3: honest 501 — this endpoint was returning fake success.
    Ok(Json(ApiResponse::error(
        "NOT_IMPLEMENTED",
        "Database creation is not implemented in this version".to_string(),
        metadata,
    )))
}

pub async fn drop_database(
    State(_state): State<AppState>,
    Extension(account): Extension<Account>,
    Path(_name): Path<String>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
    let request_id = Uuid::new_v4().to_string();
    let metadata = ResponseMetadata {
        request_id,
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    if account.role != AccountRole::Superuser {
        return Ok(Json(ApiResponse::error(
            "PERMISSION_DENIED",
            "Only Superuser can drop databases".to_string(),
            metadata,
        )));
    }

    // v0.3: honest 501 — this endpoint was returning fake success.
    Ok(Json(ApiResponse::error(
        "NOT_IMPLEMENTED",
        "Database deletion is not implemented in this version".to_string(),
        metadata,
    )))
}

pub async fn list_tables(
    State(state): State<AppState>,
    Extension(account): Extension<Account>,
    Path(db): Path<String>,
) -> Result<Json<ApiResponse<Vec<Value>>>, StatusCode> {
    let start = Instant::now();
    let request_id = Uuid::new_v4().to_string();
    
    // Check if user has access to this database
    if !account.databases.iter().any(|d| d.database == db) {
        let metadata = ResponseMetadata {
            request_id,
            execution_time_ms: start.elapsed().as_millis(),
            rows_affected: None,
            instance_id: Some(account.instance_id.clone()),
            timestamp: Utc::now(),
        };
        return Ok(Json(ApiResponse::error(
            "FORBIDDEN",
            format!("No access to database: {db}"),
            metadata,
        )));
    }
    
    // Create query to list tables
    let query_req = QueryRequest {
        query: "SELECT schemaname as schema, tablename as name, 'table' as type, tableowner as owner FROM pg_tables WHERE schemaname NOT IN ('pg_catalog', 'information_schema') ORDER BY schemaname, tablename".to_string(),
        database: db,
        params: vec![],
        options: QueryOptions::default(),
    };
    
    // Execute query
    match query_handler(State(state), Extension(account.clone()), HeaderMap::new(), Json(query_req)).await {
        Ok(Json(response)) => {
            if response.success {
                let data = response.data.unwrap();
                let metadata = ResponseMetadata {
                    request_id,
                    execution_time_ms: start.elapsed().as_millis(),
                    rows_affected: Some(data.rows.len() as u64),
                    instance_id: Some(account.instance_id),
                    timestamp: Utc::now(),
                };
                Ok(Json(ApiResponse::success(data.rows, metadata)))
            } else {
                Ok(Json(ApiResponse { success: false, data: None, error: response.error, metadata: response.metadata }))
            }
        }
        Err(e) => Err(e),
    }
}

pub async fn get_schema(
    State(_state): State<AppState>,
    Extension(account): Extension<Account>,
    Path(_db): Path<String>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
    let request_id = Uuid::new_v4().to_string();
    let metadata = ResponseMetadata {
        request_id,
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    // v0.3: honest 501 — this endpoint was returning an empty object.
    Ok(Json(ApiResponse::error(
        "NOT_IMPLEMENTED",
        "Schema introspection is not implemented in this version".to_string(),
        metadata,
    )))
}

pub async fn get_account_info(
    State(state): State<AppState>,
    Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
    let metadata = ResponseMetadata {
        request_id: Uuid::new_v4().to_string(),
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    // Get current connection count
    let active_connections = state.active_connections
        .get(&account.id)
        .map(|entry| *entry.value())
        .unwrap_or(0);
    
    let info = json!({
        "id": account.id,
        "name": account.name,
        "role": account.role,
        "databases": account.databases.iter().map(|db| db.database.clone()).collect::<Vec<_>>(),
        "rate_limit": account.rate_limit,
        "max_connections": account.max_connections,
        "active_connections": active_connections,
    });
    
    Ok(Json(ApiResponse::success(info, metadata)))
}

pub async fn get_usage_stats(
    Extension(account): Extension<Account>,
) -> Result<Json<ApiResponse<Value>>, StatusCode> {
    let metadata = ResponseMetadata {
        request_id: Uuid::new_v4().to_string(),
        execution_time_ms: 0,
        rows_affected: None,
        instance_id: Some(account.instance_id.clone()),
        timestamp: Utc::now(),
    };
    
    // v0.3: honest 501 — this endpoint was returning fabricated zeros.
    Ok(Json(ApiResponse::error(
        "NOT_IMPLEMENTED",
        "Usage statistics are not implemented in this version".to_string(),
        metadata,
    )))
}

pub async fn serve_openapi() -> impl IntoResponse {
    let openapi_content = include_str!("../openapi.json");
    
    Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "application/json")
        .body(openapi_content.to_string())
        .unwrap()
}

pub async fn serve_docs() -> impl IntoResponse {
    Html(r#"<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>pg-api Documentation</title>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui.css">
    <style>
        body {
            margin: 0;
            padding: 0;
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
        }
        .header {
            background-color: #1a1a1a;
            color: white;
            padding: 1rem 2rem;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .header h1 {
            margin: 0;
            font-size: 1.5rem;
            font-weight: 500;
        }
        .header p {
            margin: 0.5rem 0 0 0;
            opacity: 0.8;
            font-size: 0.9rem;
        }
        #swagger-ui {
            margin-top: 0;
        }
        .swagger-ui .topbar {
            display: none;
        }
    </style>
</head>
<body>
    <div class="header">
        <h1>pg-api Documentation</h1>
        <p>PostgreSQL API Service - RESTful API for database operations</p>
    </div>
    <div id="swagger-ui"></div>
    
    <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui-bundle.js"></script>
    <script src="https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9/swagger-ui-standalone-preset.js"></script>
    <script>
        window.onload = function() {
            window.ui = SwaggerUIBundle({
                url: "/openapi.json",
                dom_id: '#swagger-ui',
                deepLinking: true,
                presets: [
                    SwaggerUIBundle.presets.apis,
                    SwaggerUIStandalonePreset
                ],
                plugins: [
                    SwaggerUIBundle.plugins.DownloadUrl
                ],
                layout: "StandaloneLayout",
                persistAuthorization: true,
                tryItOutEnabled: true,
                supportedSubmitMethods: ['get', 'post', 'put', 'delete', 'patch'],
                onComplete: function() {
                    console.log("Swagger UI loaded");
                }
            });
        };
    </script>
</body>
</html>"#)
}