GopherGate/src/dashboard/auth.rs

use axum::{extract::State, http::{HeaderMap, HeaderValue}, response::{Json, IntoResponse}};
use bcrypt;
use serde::Deserialize;
use sqlx::Row;
use tracing::warn;

use super::{ApiResponse, DashboardState};

// Authentication handlers
#[derive(Deserialize)]
pub(super) struct LoginRequest {
    pub(super) username: String,
    pub(super) password: String,
}

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

    let user_result = sqlx::query(
        "SELECT username, password_hash, display_name, role, must_change_password FROM users WHERE username = ?",
    )
    .bind(&payload.username)
    .fetch_optional(pool)
    .await;

    match user_result {
        Ok(Some(row)) => {
            let hash = row.get::<String, _>("password_hash");
            if bcrypt::verify(&payload.password, &hash).unwrap_or(false) {
                let username = row.get::<String, _>("username");
                let role = row.get::<String, _>("role");
                let display_name = row
                    .get::<Option<String>, _>("display_name")
                    .unwrap_or_else(|| username.clone());
                let must_change_password = row.get::<bool, _>("must_change_password");
                let token = state
                    .session_manager
                    .create_session(username.clone(), role.clone())
                    .await;
                Json(ApiResponse::success(serde_json::json!({
                    "token": token,
                    "must_change_password": must_change_password,
                    "user": {
                        "username": username,
                        "name": display_name,
                        "role": role
                    }
                })))
            } else {
                Json(ApiResponse::error("Invalid username or password".to_string()))
            }
        }
        Ok(None) => Json(ApiResponse::error("Invalid username or password".to_string())),
        Err(e) => {
            warn!("Database error during login: {}", e);
            Json(ApiResponse::error("Login failed due to system error".to_string()))
        }
    }
}

pub(super) async fn handle_auth_status(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
) -> impl IntoResponse {
    let token = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "));

    if let Some(token) = token
        && let Some((session, new_token)) = state.session_manager.validate_session_with_refresh(token).await
    {
        // Look up display_name from DB
        let display_name = sqlx::query_scalar::<_, Option<String>>(
            "SELECT display_name FROM users WHERE username = ?",
        )
        .bind(&session.username)
        .fetch_optional(&state.app_state.db_pool)
        .await
        .ok()
        .flatten()
        .flatten()
        .unwrap_or_else(|| session.username.clone());

        let mut headers = HeaderMap::new();
        if let Some(refreshed_token) = new_token {
            if let Ok(header_value) = HeaderValue::from_str(&refreshed_token) {
                headers.insert("X-Refreshed-Token", header_value);
            }
        }
        return (headers, Json(ApiResponse::success(serde_json::json!({
            "authenticated": true,
            "user": {
                "username": session.username,
                "name": display_name,
                "role": session.role
            }
        }))));
    }

    (HeaderMap::new(), Json(ApiResponse::error("Not authenticated".to_string())))
}

#[derive(Deserialize)]
pub(super) struct ChangePasswordRequest {
    pub(super) current_password: String,
    pub(super) new_password: String,
}

pub(super) async fn handle_change_password(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
    Json(payload): Json<ChangePasswordRequest>,
) -> impl IntoResponse {
    let pool = &state.app_state.db_pool;

    // Extract the authenticated user from the session token
    let token = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "));

    let (session, new_token) = match token {
        Some(t) => match state.session_manager.validate_session_with_refresh(t).await {
            Some((session, new_token)) => (Some(session), new_token),
            None => (None, None),
        },
        None => (None, None),
    };

    let mut response_headers = HeaderMap::new();
    if let Some(refreshed_token) = new_token {
        if let Ok(header_value) = HeaderValue::from_str(&refreshed_token) {
            response_headers.insert("X-Refreshed-Token", header_value);
        }
    }

    let username = match session {
        Some(s) => s.username,
        None => return (response_headers, Json(ApiResponse::error("Not authenticated".to_string()))),
    };

    let user_result = sqlx::query("SELECT password_hash FROM users WHERE username = ?")
        .bind(&username)
        .fetch_one(pool)
        .await;

    match user_result {
        Ok(row) => {
            let hash = row.get::<String, _>("password_hash");
            if bcrypt::verify(&payload.current_password, &hash).unwrap_or(false) {
                let new_hash = match bcrypt::hash(&payload.new_password, 12) {
                    Ok(h) => h,
                    Err(_) => return (response_headers, Json(ApiResponse::error("Failed to hash new password".to_string()))),
                };

                let update_result = sqlx::query(
                    "UPDATE users SET password_hash = ?, must_change_password = FALSE WHERE username = ?",
                )
                .bind(new_hash)
                .bind(&username)
                .execute(pool)
                .await;

                match update_result {
                    Ok(_) => (response_headers, Json(ApiResponse::success(
                        serde_json::json!({ "message": "Password updated successfully" }),
                    ))),
                    Err(e) => (response_headers, Json(ApiResponse::error(format!("Failed to update database: {}", e)))),
                }
            } else {
                (response_headers, Json(ApiResponse::error("Current password incorrect".to_string())))
            }
        }
        Err(e) => (response_headers, Json(ApiResponse::error(format!("User not found: {}", e)))),
    }
}

pub(super) async fn handle_logout(
    State(state): State<DashboardState>,
    headers: axum::http::HeaderMap,
) -> Json<ApiResponse<serde_json::Value>> {
    let token = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "));

    if let Some(token) = token {
        state.session_manager.revoke_session(token).await;
    }

    Json(ApiResponse::success(serde_json::json!({ "message": "Logged out" })))
}

/// Helper: Extract and validate a session from the Authorization header.
/// Returns the Session and optional new token if refreshed, or an error response.
pub(super) async fn extract_session(
    state: &DashboardState,
    headers: &axum::http::HeaderMap,
) -> Result<(super::sessions::Session, Option<String>), Json<ApiResponse<serde_json::Value>>> {
    let token = headers
        .get("Authorization")
        .and_then(|v| v.to_str().ok())
        .and_then(|v| v.strip_prefix("Bearer "));

    match token {
        Some(t) => match state.session_manager.validate_session_with_refresh(t).await {
            Some((session, new_token)) => Ok((session, new_token)),
            None => Err(Json(ApiResponse::error("Session expired or invalid".to_string()))),
        },
        None => Err(Json(ApiResponse::error("Not authenticated".to_string()))),
    }
}

/// Helper: Extract session and require admin role.
/// Returns session and optional new token if refreshed.
pub(super) async fn require_admin(
    state: &DashboardState,
    headers: &axum::http::HeaderMap,
) -> Result<(super::sessions::Session, Option<String>), Json<ApiResponse<serde_json::Value>>> {
    let (session, new_token) = extract_session(state, headers).await?;
    if session.role != "admin" {
        return Err(Json(ApiResponse::error("Admin access required".to_string())));
    }
    Ok((session, new_token))
}