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 bytes::Bytes;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::auth::session::SessionState;
use crate::config::AppConfig;
use crate::db;
use crate::error::AppError;
use crate::imap::client::{ImapClient, ImapCredentials};
use crate::storage::{self, BlobStorage, attachment_key};

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

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

#[derive(Debug, Deserialize)]
pub(crate) struct AttachMessageSourceRequest {
    folder: String,
    uid: u32,
}

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) => {
                    if writer.abort().await.is_err()
                        && let Err(cleanup_error) =
                            storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key)
                                .await
                    {
                        tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue aborted attachment cleanup");
                    }
                    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 {
                if writer.abort().await.is_err()
                    && let Err(cleanup_error) =
                        storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
                {
                    tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue oversized attachment cleanup");
                }
                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 {
                if writer.abort().await.is_err()
                    && let Err(cleanup_error) =
                        storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
                {
                    tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue interrupted attachment cleanup");
                }
                return Err(AppError::ServiceUnavailable(format!(
                    "Failed to upload attachment: {e}"
                )));
            }
        }

        let stat = match writer.finish().await {
            Ok(stat) => stat,
            Err(e) => {
                if let Err(cleanup_error) =
                    storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
                {
                    tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue failed-finalize attachment cleanup");
                }
                return Err(AppError::ServiceUnavailable(format!(
                    "Failed to finalize attachment upload: {e}"
                )));
            }
        };
        if stat.size != size as u64 {
            if let Err(cleanup_error) =
                storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
            {
                tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue invalid attachment cleanup");
            }
            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) {
            if let Err(cleanup_error) =
                storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
            {
                tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue unattached blob cleanup");
            }
            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 attach_message_source(
    Extension(session): Extension<SessionState>,
    Extension(config): Extension<Arc<AppConfig>>,
    Extension(storage): Extension<Arc<dyn BlobStorage>>,
    Extension(imap_client): Extension<Arc<dyn ImapClient>>,
    AxumPath(draft_id): AxumPath<String>,
    Json(req): Json<AttachMessageSourceRequest>,
) -> Result<Response, AppError> {
    db::drafts::validate_draft_id(&draft_id).map_err(AppError::BadRequest)?;
    if req.folder.trim().is_empty() || req.uid == 0 {
        return Err(AppError::BadRequest(
            "Message folder and UID are required".to_string(),
        ));
    }

    let creds = ImapCredentials {
        host: session.imap_host.clone(),
        port: session.imap_port,
        tls: session.imap_tls,
        tls_verify: session.imap_tls_verify,
        email: session.email.clone(),
        password: session.password.clone(),
    };
    let raw = imap_client
        .fetch_raw_message(&creds, &req.folder, req.uid)
        .await
        .map_err(|e| AppError::ServiceUnavailable(format!("IMAP error: {e}")))?;
    if raw.is_empty() {
        return Err(AppError::BadRequest("Message source is empty".to_string()));
    }
    if raw.len() > config.max_attachment_size_bytes() {
        return Err(AppError::BadRequest(format!(
            "Message source exceeds maximum attachment size of {} MiB",
            config.max_attachment_size_mb
        )));
    }

    let conn = db::pool::open_user_db(&config.data_dir, &session.user_hash)
        .map_err(|e| AppError::InternalError(format!("Failed to open database: {e}")))?;
    if db::drafts::get_draft(&conn, &draft_id)
        .map_err(AppError::InternalError)?
        .is_none()
    {
        db::drafts::upsert_draft(
            &conn, &draft_id, "", "", "", "", "", None, None, None, false, "normal", "html",
            "none", false, false, false,
        )
        .map_err(AppError::InternalError)?;
    }

    let filename = format!("message-{}.eml", req.uid);
    let content_type = "message/rfc822".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 size = raw.len() as i64;
    let checksum = format!("{:x}", Sha256::digest(&raw));
    let mut writer = storage
        .begin_upload(&key)
        .await
        .map_err(|e| AppError::ServiceUnavailable(format!("Storage unavailable: {e}")))?;

    if let Err(error) = writer.write(Bytes::from(raw)).await {
        if writer.abort().await.is_err()
            && let Err(cleanup_error) =
                storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
        {
            tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue generated attachment cleanup");
        }
        return Err(AppError::ServiceUnavailable(format!(
            "Failed to store message attachment: {error}"
        )));
    }

    let stat = match writer.finish().await {
        Ok(stat) => stat,
        Err(error) => {
            if let Err(cleanup_error) =
                storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
            {
                tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue failed generated attachment cleanup");
            }
            return Err(AppError::ServiceUnavailable(format!(
                "Failed to finalize message attachment: {error}"
            )));
        }
    };
    if stat.size != size as u64 {
        if let Err(cleanup_error) =
            storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
        {
            tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue invalid generated attachment cleanup");
        }
        return Err(AppError::InternalError(
            "Message attachment size verification failed".to_string(),
        ));
    }

    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(error) = db::drafts::add_draft_attachment(&conn, &attachment) {
        if let Err(cleanup_error) =
            storage::enqueue_unreferenced_delete(&conn, storage.as_ref(), &key).await
        {
            tracing::error!(error = %cleanup_error, storage_key = %key, "Failed to queue unattached generated blob cleanup");
        }
        return Err(AppError::InternalError(error));
    }

    Ok(Json(serde_json::json!({
        "attachments": [UploadResponse {
            id: att_id,
            filename,
            content_type,
            size,
        }],
    }))
    .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 transaction = conn.unchecked_transaction().map_err(|e| {
        AppError::InternalError(format!(
            "Failed to start attachment delete transaction: {e}"
        ))
    })?;
    let attachment =
        db::drafts::get_draft_attachment_for_update(&transaction, &draft_id, &attachment_id)
            .map_err(AppError::InternalError)?
            .ok_or_else(|| AppError::NotFound("Attachment not found".to_string()))?;
    let storage_key = attachment.storage_key.clone();
    let cleanup_event = db::blob_delete_outbox::enqueue(&transaction, &storage_key)
        .map_err(AppError::InternalError)?;
    let deleted = db::drafts::delete_draft_attachment(&transaction, &attachment_id)
        .map_err(AppError::InternalError)?;
    if !deleted {
        return Err(AppError::NotFound("Attachment not found".to_string()));
    }
    transaction.commit().map_err(|e| {
        AppError::InternalError(format!("Failed to commit attachment deletion: {e}"))
    })?;
    if let Err(error) =
        storage::process_queued_delete(&conn, storage.as_ref(), &cleanup_event).await
    {
        tracing::warn!(error = %error, storage_key = %storage_key, "Failed to process queued attachment blob deletion");
    }

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