zhealth/backend/src/handlers/sources.rs

//! Sources API handlers - file upload and management.

use axum::{
    extract::{Path, State},
    http::StatusCode,
    Json,
};
use axum_extra::extract::Multipart;
use chrono::Utc;
use sea_orm::{ActiveModelTrait, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, Set};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use tokio::fs;
use tokio::io::AsyncWriteExt;

use crate::models::bio::{biomarker_entry, source};

/// Response for a source.
#[derive(Serialize)]
pub struct SourceResponse {
    pub id: i32,
    pub user_id: i32,
    pub name: String,
    pub file_path: String,
    pub file_type: String,
    pub file_size: i64,
    pub status: String,
    pub biomarker_count: Option<i32>,
    pub ocr_data: Option<String>,
    pub description: Option<String>,
    pub uploaded_at: String,
}

/// State that includes config for upload path.
#[derive(Clone)]
pub struct SourcesState {
    pub db: DatabaseConnection,
    pub uploads_path: PathBuf,
}

/// GET /api/sources - List all sources for current user.
pub async fn list_sources(
    State(state): State<SourcesState>,
    // TODO: Get user_id from session
) -> Result<Json<Vec<SourceResponse>>, StatusCode> {
    let sources = source::Entity::find()
        .all(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let items: Vec<SourceResponse> = sources
        .into_iter()
        .map(|s| SourceResponse {
            id: s.id,
            user_id: s.user_id,
            name: s.name,
            file_path: s.file_path,
            file_type: s.file_type,
            file_size: s.file_size,
            status: s.status,
            biomarker_count: s.biomarker_count,
            ocr_data: s.ocr_data,
            description: s.description,
            uploaded_at: s.uploaded_at.to_string(),
        })
        .collect();

    Ok(Json(items))
}

/// GET /api/sources/:id - Get a source by ID.
pub async fn get_source(
    State(state): State<SourcesState>,
    Path(id): Path<i32>,
) -> Result<Json<SourceResponse>, StatusCode> {
    let s = source::Entity::find_by_id(id)
        .one(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    Ok(Json(SourceResponse {
        id: s.id,
        user_id: s.user_id,
        name: s.name,
        file_path: s.file_path,
        file_type: s.file_type,
        file_size: s.file_size,
        status: s.status,
        biomarker_count: s.biomarker_count,
        ocr_data: s.ocr_data,
        description: s.description,
        uploaded_at: s.uploaded_at.to_string(),
    }))
}

/// POST /api/sources - Upload a new source file.
pub async fn upload_source(
    State(state): State<SourcesState>,
    mut multipart: Multipart,
) -> Result<Json<SourceResponse>, StatusCode> {
    let mut file_name: Option<String> = None;
    let mut file_type: Option<String> = None;
    let mut file_data: Option<Vec<u8>> = None;
    let mut name: Option<String> = None;
    let mut description: Option<String> = None;
    let mut user_id: Option<i32> = None;

    while let Some(field) = multipart.next_field().await.map_err(|e| {
        tracing::error!("Multipart error: {:?}", e);
        StatusCode::BAD_REQUEST
    })? {
        let field_name = field.name().unwrap_or("").to_string();

        match field_name.as_str() {
            "file" => {
                file_name = field.file_name().map(|s| s.to_string());
                file_type = field.content_type().map(|s| s.to_string());
                file_data = Some(field.bytes().await.map_err(|e| {
                    tracing::error!("Failed to read file data: {:?}", e);
                    StatusCode::BAD_REQUEST
                })?.to_vec());
            }
            "name" => {
                name = Some(field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?);
            }
            "description" => {
                description = Some(field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?);
            }
            "user_id" => {
                let text = field.text().await.map_err(|_| StatusCode::BAD_REQUEST)?;
                user_id = Some(text.parse().map_err(|_| StatusCode::BAD_REQUEST)?);
            }
            _ => {}
        }
    }

    let file_data = file_data.ok_or(StatusCode::BAD_REQUEST)?;
    let user_id = user_id.ok_or(StatusCode::BAD_REQUEST)?;
    let original_name = file_name.unwrap_or_else(|| "upload".to_string());
    let display_name = name.unwrap_or_else(|| original_name.clone());
    let content_type = file_type.unwrap_or_else(|| "application/octet-stream".to_string());
    let file_size = file_data.len() as i64;

    // Generate unique filename
    let timestamp = Utc::now().timestamp_millis();
    let safe_name = original_name.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "_");
    let stored_name = format!("{}_{}", timestamp, safe_name);

    // Ensure uploads directory exists
    fs::create_dir_all(&state.uploads_path).await.map_err(|e| {
        tracing::error!("Failed to create uploads dir: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    // Write file
    let file_path = state.uploads_path.join(&stored_name);
    let mut file = fs::File::create(&file_path).await.map_err(|e| {
        tracing::error!("Failed to create file: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;
    file.write_all(&file_data).await.map_err(|e| {
        tracing::error!("Failed to write file: {:?}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let now = Utc::now().naive_utc();

    let new_source = source::ActiveModel {
        user_id: Set(user_id),
        name: Set(display_name.clone()),
        file_path: Set(file_path.to_string_lossy().to_string()),
        file_type: Set(content_type.clone()),
        file_size: Set(file_size),
        status: Set("pending".to_string()),
        biomarker_count: Set(None),
        ocr_data: Set(None),
        description: Set(description.clone()),
        uploaded_at: Set(now),
        ..Default::default()
    };

    let inserted = new_source
        .insert(&state.db)
        .await
        .map_err(|e| {
            tracing::error!("Failed to insert source: {:?}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })?;

    Ok(Json(SourceResponse {
        id: inserted.id,
        user_id: inserted.user_id,
        name: inserted.name,
        file_path: inserted.file_path,
        file_type: inserted.file_type,
        file_size: inserted.file_size,
        status: inserted.status,
        biomarker_count: inserted.biomarker_count,
        ocr_data: inserted.ocr_data,
        description: inserted.description,
        uploaded_at: inserted.uploaded_at.to_string(),
    }))
}

/// DELETE /api/sources/:id - Delete a source.
pub async fn delete_source(
    State(state): State<SourcesState>,
    Path(id): Path<i32>,
) -> Result<StatusCode, StatusCode> {
    // Get the source first to delete the file
    let s = source::Entity::find_by_id(id)
        .one(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    // Delete related biomarker entries first (cascade delete)
    biomarker_entry::Entity::delete_many()
        .filter(biomarker_entry::Column::SourceId.eq(id))
        .exec(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    // Delete file from disk
    if let Err(e) = fs::remove_file(&s.file_path).await {
        tracing::warn!("Failed to delete file {}: {:?}", s.file_path, e);
    }

    // Delete from database
    let result = source::Entity::delete_by_id(id)
        .exec(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    if result.rows_affected == 0 {
        return Err(StatusCode::NOT_FOUND);
    }

    Ok(StatusCode::NO_CONTENT)
}

/// Request to update OCR data for a source.
#[derive(Deserialize)]
pub struct UpdateOcrRequest {
    pub ocr_data: String,
}

/// PUT /api/sources/:id/ocr - Update OCR data for a source.
pub async fn update_ocr(
    State(state): State<SourcesState>,
    Path(id): Path<i32>,
    Json(req): Json<UpdateOcrRequest>,
) -> Result<Json<SourceResponse>, StatusCode> {
    let existing = source::Entity::find_by_id(id)
        .one(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
        .ok_or(StatusCode::NOT_FOUND)?;

    let mut active: source::ActiveModel = existing.into();
    active.ocr_data = Set(Some(req.ocr_data));

    let updated = active
        .update(&state.db)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    Ok(Json(SourceResponse {
        id: updated.id,
        user_id: updated.user_id,
        name: updated.name,
        file_path: updated.file_path,
        file_type: updated.file_type,
        file_size: updated.file_size,
        status: updated.status,
        biomarker_count: updated.biomarker_count,
        ocr_data: updated.ocr_data,
        description: updated.description,
        uploaded_at: updated.uploaded_at.to_string(),
    }))
}