normogen/backend/src/handlers/auth.rs

use axum::{
    extract::State,
    response::Json,
    http::StatusCode,
};
use serde_json::{json, Value};
use validator::Validate;
use uuid::Uuid;
use crate::config::AppState;
use crate::auth::PasswordService;
use crate::models::user::{User, RegisterUserRequest, LoginRequest, UserRepository};
use crate::models::refresh_token::RefreshToken;
use mongodb::bson::DateTime;
use serde::Deserialize;

#[derive(Deserialize)]
pub struct RefreshTokenRequest {
    pub refresh_token: String,
}

#[derive(Deserialize)]
pub struct LogoutRequest {
    pub refresh_token: String,
}

pub async fn register(
    State(state): State<AppState>,
    Json(payload): Json<RegisterUserRequest>,
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
    if let Err(errors) = payload.validate() {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(json!({ "error": "Validation failed", "details": errors.to_string() }))
        ));
    }
    
    let user_repo = UserRepository::new(state.db.collection("users"));
    if let Ok(Some(_)) = user_repo.find_by_email(&payload.email).await {
        return Err((
            StatusCode::CONFLICT,
            Json(json!({ "error": "Email already registered" }))
        ));
    }
    
    let user_id = Uuid::new_v4().to_string();
    let now = DateTime::now();
    
    let user = User {
        id: None,
        user_id: user_id.clone(),
        email: payload.email.clone(),
        password_hash: payload.password_hash,
        encrypted_recovery_phrase: payload.encrypted_recovery_phrase,
        recovery_phrase_iv: payload.recovery_phrase_iv,
        recovery_phrase_auth_tag: payload.recovery_phrase_auth_tag,
        token_version: 0,
        family_id: None,
        profile_ids: Vec::new(),
        created_at: now,
        updated_at: now,
    };
    
    if let Err(e) = user_repo.create(&user).await {
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({ "error": format!("Failed to create user: {}", e) }))
        ));
    }
    
    Ok(Json(json!({
        "message": "User registered successfully",
        "user_id": user_id,
        "email": user.email
    })))
}

pub async fn login(
    State(state): State<AppState>,
    Json(payload): Json<LoginRequest>,
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
    if let Err(errors) = payload.validate() {
        return Err((
            StatusCode::BAD_REQUEST,
            Json(json!({ "error": "Validation failed", "details": errors.to_string() }))
        ));
    }
    
    let user_repo = UserRepository::new(state.db.collection("users"));
    let user = match user_repo.find_by_email(&payload.email).await {
        Ok(Some(user)) => user,
        Ok(None) => {
            return Err((
                StatusCode::UNAUTHORIZED,
                Json(json!({ "error": "Invalid credentials" }))
            ));
        }
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": format!("Database error: {}", e) }))
            ));
        }
    };
    
    if user.password_hash != payload.password_hash {
        return Err((
            StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "Invalid credentials" }))
        ));
    }
    
    let access_token = match state.jwt_service.generate_access_token(
        &user.user_id,
        &user.email,
        user.family_id.as_deref(),
        vec!["read:own_data".to_string(), "write:own_data".to_string()],
    ) {
        Ok(token) => token,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": format!("Failed to generate token: {}", e) }))
            ));
        }
    };
    
    let refresh_token = match state.jwt_service.generate_refresh_token(&user.user_id) {
        Ok(token) => token,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": format!("Failed to generate refresh token: {}", e) }))
            ));
        }
    };
    
    let token_id = Uuid::new_v4().to_string();
    let now = DateTime::now();
    let expires_at = DateTime::now();
    // TODO: Set proper expiration (30 days from now)
    // For now, we'll need to update this when MongoDB provides proper datetime arithmetic
    
    let refresh_token_doc = RefreshToken {
        id: None,
        token_id,
        user_id: user.user_id.clone(),
        token_hash: match PasswordService::hash_password(&refresh_token) {
            Ok(hash) => hash,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    Json(json!({ "error": format!("Failed to hash token: {}", e) }))
                ));
            }
        },
        expires_at,
        created_at: now,
        revoked: false,
        revoked_at: None,
    };
    
    let refresh_token_collection = state.db.collection::<RefreshToken>("refresh_tokens");
    if let Err(e) = refresh_token_collection.insert_one(&refresh_token_doc, None).await {
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({ "error": format!("Failed to store refresh token: {}", e) }))
        ));
    }
    
    Ok(Json(json!({
        "access_token": access_token,
        "refresh_token": refresh_token,
        "user_id": user.user_id,
        "email": user.email,
        "family_id": user.family_id,
        "profile_ids": user.profile_ids
    })))
}

pub async fn refresh_token(
    State(state): State<AppState>,
    Json(payload): Json<RefreshTokenRequest>,
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
    let claims = match state.jwt_service.verify_refresh_token(&payload.refresh_token) {
        Ok(claims) => claims,
        Err(_) => {
            return Err((
                StatusCode::UNAUTHORIZED,
                Json(json!({ "error": "Invalid refresh token" }))
            ));
        }
    };
    
    let user_repo = UserRepository::new(state.db.collection("users"));
    let user = match user_repo.find_by_user_id(&claims.sub).await {
        Ok(Some(user)) => user,
        Ok(None) => {
            return Err((
                StatusCode::UNAUTHORIZED,
                Json(json!({ "error": "User not found" }))
            ));
        }
        Err(_) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": "Database error" }))
            ));
        }
    };
    
    let new_access_token = match state.jwt_service.generate_access_token(
        &user.user_id,
        &user.email,
        user.family_id.as_deref(),
        vec!["read:own_data".to_string(), "write:own_data".to_string()],
    ) {
        Ok(token) => token,
        Err(_) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": "Failed to generate token" }))
            ));
        }
    };
    
    let new_refresh_token = match state.jwt_service.generate_refresh_token(&user.user_id) {
        Ok(token) => token,
        Err(_) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(json!({ "error": "Failed to generate refresh token" }))
            ));
        }
    };
    
    Ok(Json(json!({
        "access_token": new_access_token,
        "refresh_token": new_refresh_token
    })))
}

pub async fn logout(
    State(state): State<AppState>,
    Json(payload): Json<LogoutRequest>,
) -> Result<Json<Value>, (StatusCode, Json<Value>)> {
    let _claims = match state.jwt_service.verify_refresh_token(&payload.refresh_token) {
        Ok(claims) => claims,
        Err(_) => {
            return Err((
                StatusCode::UNAUTHORIZED,
                Json(json!({ "error": "Invalid refresh token" }))
            ));
        }
    };
    
    // TODO: Mark token as revoked in database
    
    Ok(Json(json!({ "message": "Logged out successfully" })))
}