GopherGate/src/dashboard/clients.rs

use axum::{
    extract::{Path, State},
    response::Json,
};
use chrono;
use rand::Rng;
use serde::Deserialize;
use serde_json;
use sqlx::Row;
use tracing::warn;
use uuid;

use super::{ApiResponse, DashboardState};

/// Generate a random API token: sk-{48 hex chars}
fn generate_token() -> String {
    let mut rng = rand::rng();
    let bytes: Vec<u8> = (0..24).map(|_| rng.random::<u8>()).collect();
    format!("sk-{}", hex::encode(bytes))
}

#[derive(Deserialize)]
pub(super) struct CreateClientRequest {
    pub(super) name: String,
    pub(super) client_id: Option<String>,
}

#[derive(Deserialize)]
pub(super) struct UpdateClientPayload {
    pub(super) name: Option<String>,
    pub(super) description: Option<String>,
    pub(super) is_active: Option<bool>,
    pub(super) rate_limit_per_minute: Option<i64>,
}

pub(super) async fn handle_get_clients(State(state): State<DashboardState>) -> Json<ApiResponse<serde_json::Value>> {
    let pool = &state.app_state.db_pool;

    let result = sqlx::query(
        r#"
        SELECT 
            client_id as id,
            name,
            description,
            created_at,
            total_requests,
            total_tokens,
            total_cost,
            is_active,
            rate_limit_per_minute
        FROM clients
        ORDER BY created_at DESC
        "#,
    )
    .fetch_all(pool)
    .await;

    match result {
        Ok(rows) => {
            let clients: Vec<serde_json::Value> = rows
                .into_iter()
                .map(|row| {
                    serde_json::json!({
                        "id": row.get::<String, _>("id"),
                        "name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
                        "description": row.get::<Option<String>, _>("description"),
                        "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
                        "requests_count": row.get::<i64, _>("total_requests"),
                        "total_tokens": row.get::<i64, _>("total_tokens"),
                        "total_cost": row.get::<f64, _>("total_cost"),
                        "status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
                        "rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
                    })
                })
                .collect();

            Json(ApiResponse::success(serde_json::json!(clients)))
        }
        Err(e) => {
            warn!("Failed to fetch clients: {}", e);
            Json(ApiResponse::error("Failed to fetch clients".to_string()))
        }
    }
}

pub(super) async fn handle_create_client(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Json(payload): Json<CreateClientRequest>,
) -> Json<ApiResponse<serde_json::Value>> {
    let (_session, _) = match super::auth::require_admin(&state, &headers).await {
        Ok((session, new_token)) => (session, new_token),
        Err(e) => return e,
    };

    let pool = &state.app_state.db_pool;

    let client_id = payload
        .client_id
        .unwrap_or_else(|| format!("client-{}", &uuid::Uuid::new_v4().to_string()[..8]));

    let result = sqlx::query(
        r#"
        INSERT INTO clients (client_id, name, is_active)
        VALUES (?, ?, TRUE)
        RETURNING *
        "#,
    )
    .bind(&client_id)
    .bind(&payload.name)
    .fetch_one(pool)
    .await;

    match result {
        Ok(row) => {
            // Auto-generate a token for the new client
            let token = generate_token();
            let token_result = sqlx::query(
                "INSERT INTO client_tokens (client_id, token, name) VALUES (?, ?, 'default')",
            )
            .bind(&client_id)
            .bind(&token)
            .execute(pool)
            .await;

            if let Err(e) = token_result {
                warn!("Client created but failed to generate token: {}", e);
            }

            Json(ApiResponse::success(serde_json::json!({
                "id": row.get::<String, _>("client_id"),
                "name": row.get::<Option<String>, _>("name"),
                "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
                "status": "active",
                "token": token,
            })))
        }
        Err(e) => {
            warn!("Failed to create client: {}", e);
            Json(ApiResponse::error(format!("Failed to create client: {}", e)))
        }
    }
}

pub(super) async fn handle_get_client(
    State(state): State<DashboardState>,
    Path(id): Path<String>,
) -> Json<ApiResponse<serde_json::Value>> {
    let pool = &state.app_state.db_pool;

    let result = sqlx::query(
        r#"
        SELECT 
            c.client_id as id,
            c.name,
            c.description,
            c.is_active,
            c.rate_limit_per_minute,
            c.created_at,
            COALESCE(c.total_tokens, 0) as total_tokens,
            COALESCE(c.total_cost, 0.0) as total_cost,
            COUNT(r.id) as total_requests,
            MAX(r.timestamp) as last_request
        FROM clients c
        LEFT JOIN llm_requests r ON c.client_id = r.client_id
        WHERE c.client_id = ?
        GROUP BY c.client_id
        "#,
    )
    .bind(&id)
    .fetch_optional(pool)
    .await;

    match result {
        Ok(Some(row)) => Json(ApiResponse::success(serde_json::json!({
            "id": row.get::<String, _>("id"),
            "name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
            "description": row.get::<Option<String>, _>("description"),
            "is_active": row.get::<bool, _>("is_active"),
            "rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
            "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
            "total_tokens": row.get::<i64, _>("total_tokens"),
            "total_cost": row.get::<f64, _>("total_cost"),
            "total_requests": row.get::<i64, _>("total_requests"),
            "last_request": row.get::<Option<chrono::DateTime<chrono::Utc>>, _>("last_request"),
            "status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
        }))),
        Ok(None) => Json(ApiResponse::error(format!("Client '{}' not found", id))),
        Err(e) => {
            warn!("Failed to fetch client: {}", e);
            Json(ApiResponse::error(format!("Failed to fetch client: {}", e)))
        }
    }
}

pub(super) async fn handle_update_client(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Path(id): Path<String>,
    Json(payload): Json<UpdateClientPayload>,
) -> Json<ApiResponse<serde_json::Value>> {
    let (_session, _) = match super::auth::require_admin(&state, &headers).await {
        Ok((session, new_token)) => (session, new_token),
        Err(e) => return e,
    };

    let pool = &state.app_state.db_pool;

    // Build dynamic UPDATE query from provided fields
    let mut sets = Vec::new();
    let mut binds: Vec<String> = Vec::new();

    if let Some(ref name) = payload.name {
        sets.push("name = ?");
        binds.push(name.clone());
    }
    if let Some(ref desc) = payload.description {
        sets.push("description = ?");
        binds.push(desc.clone());
    }
    if payload.is_active.is_some() {
        sets.push("is_active = ?");
    }
    if payload.rate_limit_per_minute.is_some() {
        sets.push("rate_limit_per_minute = ?");
    }

    if sets.is_empty() {
        return Json(ApiResponse::error("No fields to update".to_string()));
    }

    // Always update updated_at
    sets.push("updated_at = CURRENT_TIMESTAMP");

    let sql = format!("UPDATE clients SET {} WHERE client_id = ?", sets.join(", "));
    let mut query = sqlx::query(&sql);

    // Bind in the same order as sets
    for b in &binds {
        query = query.bind(b);
    }
    if let Some(active) = payload.is_active {
        query = query.bind(active);
    }
    if let Some(rate) = payload.rate_limit_per_minute {
        query = query.bind(rate);
    }
    query = query.bind(&id);

    match query.execute(pool).await {
        Ok(result) => {
            if result.rows_affected() == 0 {
                return Json(ApiResponse::error(format!("Client '{}' not found", id)));
            }
            // Return the updated client
            let row = sqlx::query(
                r#"
                SELECT client_id as id, name, description, is_active, rate_limit_per_minute,
                       created_at, total_requests, total_tokens, total_cost
                FROM clients WHERE client_id = ?
                "#,
            )
            .bind(&id)
            .fetch_one(pool)
            .await;

            match row {
                Ok(row) => Json(ApiResponse::success(serde_json::json!({
                    "id": row.get::<String, _>("id"),
                    "name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "Unnamed".to_string()),
                    "description": row.get::<Option<String>, _>("description"),
                    "is_active": row.get::<bool, _>("is_active"),
                    "rate_limit_per_minute": row.get::<Option<i64>, _>("rate_limit_per_minute"),
                    "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
                    "total_requests": row.get::<i64, _>("total_requests"),
                    "total_tokens": row.get::<i64, _>("total_tokens"),
                    "total_cost": row.get::<f64, _>("total_cost"),
                    "status": if row.get::<bool, _>("is_active") { "active" } else { "inactive" },
                }))),
                Err(e) => {
                    warn!("Failed to fetch updated client: {}", e);
                    // Update succeeded but fetch failed — still report success
                    Json(ApiResponse::success(serde_json::json!({ "message": "Client updated" })))
                }
            }
        }
        Err(e) => {
            warn!("Failed to update client: {}", e);
            Json(ApiResponse::error(format!("Failed to update client: {}", e)))
        }
    }
}

pub(super) async fn handle_delete_client(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Path(id): Path<String>,
) -> Json<ApiResponse<serde_json::Value>> {
    let (_session, _) = match super::auth::require_admin(&state, &headers).await {
        Ok((session, new_token)) => (session, new_token),
        Err(e) => return e,
    };

    let pool = &state.app_state.db_pool;

    // Don't allow deleting the default client
    if id == "default" {
        return Json(ApiResponse::error("Cannot delete default client".to_string()));
    }

    let result = sqlx::query("DELETE FROM clients WHERE client_id = ?")
        .bind(id)
        .execute(pool)
        .await;

    match result {
        Ok(_) => Json(ApiResponse::success(serde_json::json!({ "message": "Client deleted" }))),
        Err(e) => Json(ApiResponse::error(format!("Failed to delete client: {}", e))),
    }
}

pub(super) async fn handle_client_usage(
    State(state): State<DashboardState>,
    Path(id): Path<String>,
) -> Json<ApiResponse<serde_json::Value>> {
    let pool = &state.app_state.db_pool;

    // Get per-model breakdown for this client
    let result = sqlx::query(
        r#"
        SELECT 
            model,
            provider,
            COUNT(*) as request_count,
            SUM(prompt_tokens) as prompt_tokens,
            SUM(completion_tokens) as completion_tokens,
            SUM(total_tokens) as total_tokens,
            SUM(cost) as total_cost,
            AVG(duration_ms) as avg_duration_ms
        FROM llm_requests
        WHERE client_id = ?
        GROUP BY model, provider
        ORDER BY total_cost DESC
        "#,
    )
    .bind(&id)
    .fetch_all(pool)
    .await;

    match result {
        Ok(rows) => {
            let breakdown: Vec<serde_json::Value> = rows
                .into_iter()
                .map(|row| {
                    serde_json::json!({
                        "model": row.get::<String, _>("model"),
                        "provider": row.get::<String, _>("provider"),
                        "request_count": row.get::<i64, _>("request_count"),
                        "prompt_tokens": row.get::<i64, _>("prompt_tokens"),
                        "completion_tokens": row.get::<i64, _>("completion_tokens"),
                        "total_tokens": row.get::<i64, _>("total_tokens"),
                        "total_cost": row.get::<f64, _>("total_cost"),
                        "avg_duration_ms": row.get::<f64, _>("avg_duration_ms"),
                    })
                })
                .collect();

            Json(ApiResponse::success(serde_json::json!({
                "client_id": id,
                "breakdown": breakdown,
            })))
        }
        Err(e) => {
            warn!("Failed to fetch client usage: {}", e);
            Json(ApiResponse::error(format!("Failed to fetch client usage: {}", e)))
        }
    }
}

// ── Token management endpoints ──────────────────────────────────────

pub(super) async fn handle_get_client_tokens(
    State(state): State<DashboardState>,
    Path(id): Path<String>,
) -> Json<ApiResponse<serde_json::Value>> {
    let pool = &state.app_state.db_pool;

    let result = sqlx::query(
        r#"
        SELECT id, token, name, is_active, created_at, last_used_at
        FROM client_tokens
        WHERE client_id = ?
        ORDER BY created_at DESC
        "#,
    )
    .bind(&id)
    .fetch_all(pool)
    .await;

    match result {
        Ok(rows) => {
            let tokens: Vec<serde_json::Value> = rows
                .into_iter()
                .map(|row| {
                    let token: String = row.get("token");
                    // Mask all but last 8 chars: sk-••••abcd1234
                    let masked = if token.len() > 8 {
                        format!("{}••••{}", &token[..3], &token[token.len() - 8..])
                    } else {
                        "••••".to_string()
                    };
                    serde_json::json!({
                        "id": row.get::<i64, _>("id"),
                        "token_masked": masked,
                        "name": row.get::<Option<String>, _>("name").unwrap_or_else(|| "default".to_string()),
                        "is_active": row.get::<bool, _>("is_active"),
                        "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
                        "last_used_at": row.get::<Option<chrono::DateTime<chrono::Utc>>, _>("last_used_at"),
                    })
                })
                .collect();

            Json(ApiResponse::success(serde_json::json!(tokens)))
        }
        Err(e) => {
            warn!("Failed to fetch client tokens: {}", e);
            Json(ApiResponse::error(format!("Failed to fetch client tokens: {}", e)))
        }
    }
}

#[derive(Deserialize)]
pub(super) struct CreateTokenRequest {
    pub(super) name: Option<String>,
}

pub(super) async fn handle_create_client_token(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Path(id): Path<String>,
    Json(payload): Json<CreateTokenRequest>,
) -> Json<ApiResponse<serde_json::Value>> {
    let (_session, _) = match super::auth::require_admin(&state, &headers).await {
        Ok((session, new_token)) => (session, new_token),
        Err(e) => return e,
    };

    let pool = &state.app_state.db_pool;

    // Verify client exists
    let exists: Option<(i64,)> = sqlx::query_as("SELECT 1 as x FROM clients WHERE client_id = ?")
        .bind(&id)
        .fetch_optional(pool)
        .await
        .unwrap_or(None);

    if exists.is_none() {
        return Json(ApiResponse::error(format!("Client '{}' not found", id)));
    }

    let token = generate_token();
    let token_name = payload.name.unwrap_or_else(|| "default".to_string());

    let result = sqlx::query(
        "INSERT INTO client_tokens (client_id, token, name) VALUES (?, ?, ?) RETURNING id, created_at",
    )
    .bind(&id)
    .bind(&token)
    .bind(&token_name)
    .fetch_one(pool)
    .await;

    match result {
        Ok(row) => Json(ApiResponse::success(serde_json::json!({
            "id": row.get::<i64, _>("id"),
            "token": token,
            "name": token_name,
            "created_at": row.get::<chrono::DateTime<chrono::Utc>, _>("created_at"),
        }))),
        Err(e) => {
            warn!("Failed to create client token: {}", e);
            Json(ApiResponse::error(format!("Failed to create token: {}", e)))
        }
    }
}

pub(super) async fn handle_delete_client_token(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Path((client_id, token_id)): Path<(String, i64)>,
) -> Json<ApiResponse<serde_json::Value>> {
    let (_session, _) = match super::auth::require_admin(&state, &headers).await {
        Ok((session, new_token)) => (session, new_token),
        Err(e) => return e,
    };

    let pool = &state.app_state.db_pool;

    let result = sqlx::query("DELETE FROM client_tokens WHERE id = ? AND client_id = ?")
        .bind(token_id)
        .bind(&client_id)
        .execute(pool)
        .await;

    match result {
        Ok(r) => {
            if r.rows_affected() == 0 {
                Json(ApiResponse::error("Token not found".to_string()))
            } else {
                Json(ApiResponse::success(serde_json::json!({ "message": "Token revoked" })))
            }
        }
        Err(e) => {
            warn!("Failed to delete client token: {}", e);
            Json(ApiResponse::error(format!("Failed to revoke token: {}", e)))
        }
    }
}