RustBox/src/routes/attachments.rs

use std::sync::Arc;

use axum::extract::Multipart;
use axum::extract::Path as AxumPath;
use axum::response::{IntoResponse, Response};
use axum::{Extension, Json};
use serde::Serialize;
use sha2::{Digest, Sha256};

use crate::auth::session::SessionState;
use crate::config::AppConfig;
use crate::db;
use crate::error::AppError;
use crate::storage::{BlobStorage, attachment_key};

#[derive(Debug, Serialize)]
struct UploadResponse {
    id: String,
    filename: String,
    content_type: String,
    size: i64,
}

#[derive(Debug, Serialize)]
struct DeleteResponse {
    status: String,
}

pub async fn upload_attachment(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Extension(storage): Extension<Arc<dyn BlobStorage>>,
    AxumPath(draft_id): AxumPath<String>,
    mut multipart: Multipart,
) -> Result<Response, AppError> {
    db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?;
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;

    let existing = db::drafts::get_draft(&conn, &draft_id).map_err(AppError::InternalError)?;
    if existing.is_none() {
        db::drafts::upsert_draft(
            &conn, &draft_id, "", "", "", "", "", None, None, None, false, "normal", "html",
            "none", false, false, false,
        )
        .map_err(AppError::InternalError)?;
    }

    let mut uploaded = Vec::new();
    let max_attachment_size = config.max_attachment_size_bytes();

    while let Some(mut field) = multipart
        .next_field()
        .await
        .map_err(|e| AppError::BadRequest(format!("Invalid multipart data: {e}")))?
    {
        let filename = field.file_name().unwrap_or("unnamed").to_string();
        let content_type = field
            .content_type()
            .unwrap_or("application/octet-stream")
            .to_string();
        let att_id = uuid::Uuid::new_v4().to_string();
        let key = attachment_key(&session.user_hash, &draft_id, &att_id)
            .map_err(AppError::InternalError)?;
        let mut writer = storage
            .begin_upload(&key)
            .await
            .map_err(|e| AppError::ServiceUnavailable(format!("Storage unavailable: {e}")))?;
        let mut size: usize = 0;
        let mut checksum = Sha256::new();

        loop {
            let chunk = match field.chunk().await {
                Ok(chunk) => chunk,
                Err(e) => {
                    let _ = writer.abort().await;
                    return Err(AppError::BadRequest(format!(
                        "Failed to read field data: {e}"
                    )));
                }
            };
            let Some(chunk) = chunk else { break };
            size = size.saturating_add(chunk.len());
            if size > max_attachment_size {
                let _ = writer.abort().await;
                return Err(AppError::BadRequest(format!(
                    "File '{}' exceeds maximum size of {} MiB",
                    filename, config.max_attachment_size_mb
                )));
            }
            checksum.update(&chunk);
            if let Err(e) = writer.write(chunk).await {
                let _ = writer.abort().await;
                return Err(AppError::ServiceUnavailable(format!(
                    "Failed to upload attachment: {e}"
                )));
            }
        }

        let stat = writer.finish().await.map_err(|e| {
            AppError::ServiceUnavailable(format!("Failed to finalize attachment upload: {e}"))
        })?;
        if stat.size != size as u64 {
            let _ = storage.delete(&key).await;
            return Err(AppError::InternalError(format!(
                "Attachment size verification failed for '{filename}'"
            )));
        }
        let checksum = format!("{:x}", checksum.finalize());
        let size = size as i64;
        let attachment = db::drafts::NewDraftAttachment {
            id: &att_id,
            draft_id: &draft_id,
            filename: &filename,
            content_type: &content_type,
            size,
            storage_key: &key,
            checksum_sha256: Some(&checksum),
        };
        if let Err(e) = db::drafts::add_draft_attachment(&conn, &attachment) {
            let _ = storage.delete(&key).await;
            return Err(AppError::InternalError(e));
        }

        uploaded.push(UploadResponse {
            id: att_id,
            filename,
            content_type,
            size,
        });
    }

    if uploaded.is_empty() {
        return Err(AppError::BadRequest("No files uploaded".to_string()));
    }

    Ok(Json(serde_json::json!({
        "attachments": uploaded,
    }))
    .into_response())
}

pub async fn get_attachment_content(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Extension(storage): Extension<Arc<dyn BlobStorage>>,
    AxumPath((draft_id, attachment_id)): AxumPath<(String, String)>,
) -> Result<Response, AppError> {
    db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?;
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;

    let attachments =
        db::drafts::get_draft_attachments(&conn, &draft_id).map_err(AppError::InternalError)?;
    let attachment = attachments
        .iter()
        .find(|a| a.id == attachment_id)
        .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?;

    let stream = storage
        .get(&attachment.storage_key)
        .await
        .map_err(|e| AppError::ServiceUnavailable(format!("Storage read failed: {e}")))?;

    Ok(Response::builder()
        .header("content-type", &attachment.content_type)
        .header("content-length", attachment.size.max(0).to_string())
        .header("cache-control", "private, max-age=3600")
        .body(axum::body::Body::from_stream(stream))
        .unwrap())
}

pub async fn delete_attachment(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Extension(storage): Extension<Arc<dyn BlobStorage>>,
    AxumPath((draft_id, attachment_id)): AxumPath<(String, String)>,
) -> Result<Response, AppError> {
    db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?;
    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;

    let attachments =
        db::drafts::get_draft_attachments(&conn, &draft_id).map_err(AppError::InternalError)?;
    let attachment = attachments
        .iter()
        .find(|a| a.id == attachment_id)
        .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?;

    let storage_key = attachment.storage_key.clone();
    let deleted = db::drafts::delete_draft_attachment(&conn, &attachment_id)
        .map_err(AppError::InternalError)?;
    if !deleted {
        return Err(AppError::NotFound("Attachment not found".to_string()));
    }
    if let Err(error) = storage.delete(&storage_key).await {
        // SQL is authoritative. A failed physical delete leaves an orphaned blob,
        // which is safer than leaving a DB row pointing at a missing attachment.
        tracing::warn!(error = %error, storage_key = %storage_key, "Failed to remove deleted attachment blob");
    }

    Ok(Json(DeleteResponse {
        status: "deleted".to_string(),
    })
    .into_response())
}