litecloud/api/src/routes/files.rs

use axum::{
    body::StreamBody,
    extract::{Extension, Multipart, Path},
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
    routing::{delete, get, post},
    Json, Router,
};
use sqlx::PgPool;
use std::{path::PathBuf, sync::Arc};
use tokio::fs::File;
use tokio_util::io::ReaderStream;
use uuid::Uuid;

use crate::config::Config;
use crate::error::AppError;
use crate::models::file::{CreateDirectoryDto, CreateFileDto, File as FileModel, FileType};
use crate::services::auth::AuthUser;
use crate::services::encryption::EncryptionService;
use crate::services::storage::StorageService;

pub fn routes() -> Router {
    Router::new()
        .route("/files", get(list_files))
        .route("/files/:id", get(get_file))
        .route("/files/:id/download", get(download_file))
        .route("/files/upload", post(upload_file))
        .route("/files/directory", post(create_directory))
        .route("/files/:id", delete(delete_file))
}

async fn list_files(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Path(parent_id): Option<Path<Uuid>>,
) -> Result<Json<Vec<FileModel>>, AppError> {
    let files = match parent_id {
        Some(parent_id) => FileModel::list_by_parent(&pool, parent_id, auth_user.id).await?,
        None => FileModel::list_root_directory(&pool, auth_user.id).await?,
    };

    Ok(Json(files))
}

async fn get_file(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Path(id): Path<Uuid>,
) -> Result<Json<FileModel>, AppError> {
    let file = FileModel::find_by_id(&pool, id).await?;

    // Check if user has access to this file
    if file.owner_id != auth_user.id {
        return Err(AppError::AccessDenied("You don't have access to this file".to_string()));
    }

    Ok(Json(file))
}

async fn upload_file(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Extension(config): Extension<Arc<Config>>,
    Extension(encryption_service): Extension<EncryptionService>,
    mut multipart: Multipart,
) -> Result<Json<FileModel>, AppError> {
    // Extract file data from multipart form
    let mut file_name = None;
    let mut file_data = None;
    let mut parent_id = None;

    while let Some(field) = multipart.next_field().await.map_err(|e| AppError::InvalidInput(e.to_string()))? {
        let name = field.name().unwrap_or("").to_string();

        if name == "file" {
            file_name = field.file_name().map(|s| s.to_string());
            file_data = Some(field.bytes().await.map_err(|e| AppError::InvalidInput(e.to_string()))?);
        } else if name == "parent_id" {
            let parent_id_str = field.text().await.map_err(|e| AppError::InvalidInput(e.to_string()))?;
            if !parent_id_str.is_empty() {
                parent_id = Some(Uuid::parse_str(&parent_id_str).map_err(|e| AppError::InvalidInput(e.to_string()))?);
            }
        }
    }

    let file_name = file_name.ok_or_else(|| AppError::InvalidInput("File name is required".to_string()))?;
    let file_data = file_data.ok_or_else(|| AppError::InvalidInput("File data is required".to_string()))?;

    // Check user storage quota
    let user = crate::models::user::User::find_by_id(&pool, auth_user.id).await?;
    if user.storage_used + file_data.len() as i64 > user.storage_quota {
        return Err(AppError::StorageQuotaExceeded("Storage quota exceeded".to_string()));
    }

    let storage_service = StorageService::new(&config.storage_path);

    // Create file record in database
    let create_file_dto = CreateFileDto {
        name: file_name,
        parent_id,
        size: file_data.len() as i64,
        mime_type: mime_guess::from_path(&file_name).first_or_octet_stream().to_string(),
    };

    let file = FileModel::create(&pool, auth_user.id, create_file_dto).await?;

    // Encrypt and save file to disk
    let encrypted_data = encryption_service.encrypt(&file_data)?;
    storage_service.save_file(auth_user.id, file.id, &encrypted_data).await?;

    // Update user storage used
    crate::models::user::User::update_storage_used(&pool, auth_user.id, file_data.len() as i64).await?;

    Ok(Json(file))
}

async fn download_file(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Extension(config): Extension<Arc<Config>>,
    Extension(encryption_service): Extension<EncryptionService>,
    Path(id): Path<Uuid>,
) -> Result<impl IntoResponse, AppError> {
    let file = FileModel::find_by_id(&pool, id).await?;

    if file.owner_id != auth_user.id {
        return Err(AppError::AccessDenied("You don't have access to this file".to_string()));
    }

    if file.file_type != FileType::File {
        return Err(AppError::InvalidInput("Cannot download a directory".to_string()));
    }

    let storage_service = StorageService::new(&config.storage_path);

    let encrypted_data = storage_service.read_file(auth_user.id, file.id).await?;

    let decrypted_data = encryption_service.decrypt(&encrypted_data)?;

    let mut headers = HeaderMap::new();
    headers.insert(
        axum::http::header::CONTENT_TYPE,
        file.mime_type.parse().unwrap_or_else(|_| "application/octet-stream".parse().unwrap()),
    );
    headers.insert(
        axum::http::header::CONTENT_DISPOSITION,
        format!("attachment; filename=\"{}\"", file.name).parse().unwrap(),
    );

    let stream = tokio_util::io::ReaderStream::new(std::io::Cursor::new(decrypted_data));
    let body = StreamBody::new(stream);

    Ok((StatusCode::OK, headers, body))
}

async fn create_directory(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Json(create_dir_dto): Json<CreateDirectoryDto>,
) -> Result<Json<FileModel>, AppError> {
    let directory = FileModel::create_directory(&pool, auth_user.id, create_dir_dto).await?;
    Ok(Json(directory))
}

async fn delete_file(
    auth_user: AuthUser,
    Extension(pool): Extension<PgPool>,
    Extension(config): Extension<Arc<Config>>,
    Path(id): Path<Uuid>,
) -> Result<StatusCode, AppError> {
    let file = FileModel::find_by_id(&pool, id).await?;
    if file.owner_id != auth_user.id {
        return Err(AppError::AccessDenied("You don't have access to this file".to_string()));
    }
    let storage_service = StorageService::new(&config.storage_path);
    if file.file_type == FileType::Directory {
        FileModel::delete_directory_recursive(&pool, id, auth_user.id, &storage_service).await?;
    } else {
        storage_service.delete_file(auth_user.id, file.id).await?;
        FileModel::delete(&pool, id).await?;
        crate::models::user::User::update_storage_used(&pool, auth_user.id, -file.size).await?;
    }

    Ok(StatusCode::NO_CONTENT)
}