ystp/src/worker/mod.rs

use crate::error::{AppError, ErrorCode};
use crate::services::billing;
use crate::services::compress;
use crate::services::quota;
use crate::services::storage;
use crate::state::AppState;

use redis::streams::{StreamClaimReply, StreamId, StreamPendingCountReply, StreamReadOptions};
use redis::AsyncCommands;
use sqlx::FromRow;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Instant;
use tokio::sync::Semaphore;
use tokio::task::JoinSet;
use uuid::Uuid;

const STREAM_KEY: &str = "stream:compress_jobs";
const GROUP_NAME: &str = "compress_workers";
const DEAD_STREAM_KEY: &str = "stream:compress_jobs:dead";
const MAX_DELIVERIES: usize = 3;
const STALE_MESSAGE_IDLE_MS: usize = 5 * 60 * 1000;

pub async fn run(state: AppState) -> Result<(), AppError> {
    tracing::info!("Worker started");

    crate::services::bootstrap::ensure_schema(&state).await?;

    let consumer = format!("worker_{}", Uuid::new_v4());
    ensure_group(&state, &consumer).await?;

    let mut last_maintenance = Instant::now();

    loop {
        if let Err(err) = poll_once(&state, &consumer).await {
            tracing::error!(error = ?err, "worker poll error");
            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        }

        if last_maintenance.elapsed().as_secs() >= 300 {
            if let Err(err) = maintenance(&state).await {
                tracing::error!(error = ?err, "maintenance failed");
            }
            last_maintenance = Instant::now();
        }
    }
}

async fn ensure_group(state: &AppState, _consumer: &str) -> Result<(), AppError> {
    let mut conn = state.redis.clone();

    let res: Result<redis::Value, redis::RedisError> = redis::cmd("XGROUP")
        .arg("CREATE")
        .arg(STREAM_KEY)
        .arg(GROUP_NAME)
        .arg("0")
        .arg("MKSTREAM")
        .query_async(&mut conn)
        .await;

    match res {
        Ok(_) => Ok(()),
        Err(err) => {
            let msg = err.to_string();
            if msg.contains("BUSYGROUP") {
                return Ok(());
            }
            Err(AppError::new(ErrorCode::Internal, "初始化队列失败").with_source(err))
        }
    }
}

async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> {
    let mut conn = state.redis.clone();

    if let Some(msg) = claim_stale_message(&mut conn, consumer).await? {
        return handle_message(state, &mut conn, msg).await;
    }

    // Retry messages already delivered to this consumer before taking new work.
    let pending_opts = StreamReadOptions::default()
        .group(GROUP_NAME, consumer)
        .count(1);
    let mut reply: redis::streams::StreamReadReply = conn
        .xread_options(&[STREAM_KEY], &["0"], &pending_opts)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?;

    if !reply.keys.iter().any(|key| !key.ids.is_empty()) {
        let opts = StreamReadOptions::default()
            .group(GROUP_NAME, consumer)
            .count(1)
            .block(5000);

        reply = conn
            .xread_options(&[STREAM_KEY], &[">"], &opts)
            .await
            .map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?;
    }

    if reply.keys.is_empty() {
        return Ok(());
    }

    for key in reply.keys {
        for msg in key.ids {
            handle_message(state, &mut conn, msg).await?;
        }
    }

    Ok(())
}

async fn handle_message(
    state: &AppState,
    conn: &mut redis::aio::ConnectionManager,
    msg: StreamId,
) -> Result<(), AppError> {
    let Some(task_id_str) = msg.get::<String>("task_id") else {
        ack_message(conn, &msg.id).await?;
        return Ok(());
    };
    let task_id = match Uuid::parse_str(&task_id_str) {
        Ok(value) => value,
        Err(_) => {
            ack_message(conn, &msg.id).await?;
            return Ok(());
        }
    };

    match process_task(state, task_id).await {
        Ok(()) => ack_message(conn, &msg.id).await,
        Err(err) => {
            let deliveries = pending_delivery_count(conn, &msg.id).await?;
            if should_dead_letter(deliveries) {
                write_dead_letter(conn, &msg.id, task_id, deliveries, &err).await?;
                mark_task_dead_letter(state, task_id, &err.message).await?;
                ack_message(conn, &msg.id).await?;
                tracing::error!(
                    task_id = %task_id,
                    deliveries,
                    error = %err,
                    "task moved to dead-letter stream"
                );
                return Ok(());
            }

            tracing::warn!(
                task_id = %task_id,
                deliveries,
                error = %err,
                "task processing failed; message left pending for retry"
            );
            Err(err)
        }
    }
}

async fn claim_stale_message(
    conn: &mut redis::aio::ConnectionManager,
    consumer: &str,
) -> Result<Option<StreamId>, AppError> {
    let pending: StreamPendingCountReply = conn
        .xpending_count(STREAM_KEY, GROUP_NAME, "-", "+", 100)
        .await
        .map_err(|err| {
            AppError::new(ErrorCode::Internal, "查询停滞队列消息失败").with_source(err)
        })?;
    let Some(stale) = pending
        .ids
        .into_iter()
        .find(|item| item.consumer != consumer && item.last_delivered_ms >= STALE_MESSAGE_IDLE_MS)
    else {
        return Ok(None);
    };

    let claimed: StreamClaimReply = conn
        .xclaim(
            STREAM_KEY,
            GROUP_NAME,
            consumer,
            STALE_MESSAGE_IDLE_MS,
            &[stale.id],
        )
        .await
        .map_err(|err| {
            AppError::new(ErrorCode::Internal, "认领停滞队列消息失败").with_source(err)
        })?;
    Ok(claimed.ids.into_iter().next())
}

async fn pending_delivery_count(
    conn: &mut redis::aio::ConnectionManager,
    msg_id: &str,
) -> Result<usize, AppError> {
    let pending: StreamPendingCountReply = conn
        .xpending_count(STREAM_KEY, GROUP_NAME, msg_id, msg_id, 1)
        .await
        .map_err(|err| {
            AppError::new(ErrorCode::Internal, "查询队列重试次数失败").with_source(err)
        })?;
    Ok(pending
        .ids
        .first()
        .map(|item| item.times_delivered)
        .unwrap_or(1))
}

fn should_dead_letter(deliveries: usize) -> bool {
    deliveries >= MAX_DELIVERIES
}

async fn write_dead_letter(
    conn: &mut redis::aio::ConnectionManager,
    message_id: &str,
    task_id: Uuid,
    deliveries: usize,
    error: &AppError,
) -> Result<(), AppError> {
    redis::cmd("XADD")
        .arg(DEAD_STREAM_KEY)
        .arg("MAXLEN")
        .arg("~")
        .arg(10_000)
        .arg("*")
        .arg("task_id")
        .arg(task_id.to_string())
        .arg("source_message_id")
        .arg(message_id)
        .arg("deliveries")
        .arg(deliveries)
        .arg("error_code")
        .arg(error.code.as_str())
        .arg("error_message")
        .arg(&error.message)
        .arg("failed_at")
        .arg(chrono::Utc::now().to_rfc3339())
        .query_async::<_, redis::Value>(conn)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "写入死信队列失败").with_source(err))?;
    Ok(())
}

async fn mark_task_dead_letter(
    state: &AppState,
    task_id: Uuid,
    message: &str,
) -> Result<(), AppError> {
    let mut tx =
        state.db.begin().await.map_err(|err| {
            AppError::new(ErrorCode::Internal, "开启死信事务失败").with_source(err)
        })?;
    sqlx::query(
        r#"
        UPDATE task_files
        SET status = 'failed',
            error_message = $2,
            completed_at = NOW()
        WHERE task_id = $1 AND status IN ('pending', 'processing')
        "#,
    )
    .bind(task_id)
    .bind(message)
    .execute(&mut *tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "标记死信文件失败").with_source(err))?;
    sqlx::query(
        r#"
        UPDATE tasks
        SET status = 'failed',
            failed_files = GREATEST(total_files - completed_files, failed_files),
            error_message = $2,
            completed_at = NOW()
        WHERE id = $1 AND status NOT IN ('completed', 'failed', 'cancelled')
        "#,
    )
    .bind(task_id)
    .bind(message)
    .execute(&mut *tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "标记死信任务失败").with_source(err))?;
    tx.commit()
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "提交死信事务失败").with_source(err))?;
    Ok(())
}

async fn ack_message(
    conn: &mut redis::aio::ConnectionManager,
    msg_id: &str,
) -> Result<(), AppError> {
    let _: i64 = redis::cmd("XACK")
        .arg(STREAM_KEY)
        .arg(GROUP_NAME)
        .arg(msg_id)
        .query_async(conn)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "确认队列消息失败").with_source(err))?;
    Ok(())
}

#[derive(Debug, FromRow)]
struct TaskProcRow {
    status: String,
    compression_level: String,
    compression_rate: Option<i16>,
    max_width: Option<i32>,
    max_height: Option<i32>,
    preserve_metadata: bool,
    user_id: Option<Uuid>,
    session_id: Option<String>,
    api_key_id: Option<Uuid>,
    source: String,
    client_ip: Option<String>,
    retention_hours: i32,
    anonymous_units_reserved: i32,
}

#[derive(Debug, FromRow)]
struct TaskFileProcRow {
    id: Uuid,
    input_path: Option<String>,
    original_format: String,
    output_format: String,
    status: String,
}

#[derive(Debug, FromRow)]
struct CleanupFileRow {
    storage_backend: String,
    storage_endpoint_id: Option<Uuid>,
    storage_key: Option<String>,
    input_path: Option<String>,
}

#[derive(Debug, FromRow)]
struct CleanupZipRow {
    zip_storage_backend: Option<String>,
    zip_storage_endpoint_id: Option<Uuid>,
    zip_storage_key: Option<String>,
}

#[derive(Clone)]
struct TaskContext {
    api_key_id: Option<Uuid>,
    source: String,
    preserve_metadata: bool,
    session_id: Option<String>,
    anon_ip: Option<IpAddr>,
    is_anonymous: bool,
    retention_hours: i32,
    anonymous_quota_reserved: bool,
}

async fn process_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
    let Some(mut task): Option<TaskProcRow> = sqlx::query_as(
        r#"
        SELECT
          status::text AS status,
          compression_level::text AS compression_level,
          compression_rate,
          max_width,
          max_height,
          preserve_metadata,
          user_id,
          session_id,
          api_key_id,
          source::text AS source,
          host(client_ip) AS client_ip,
          retention_hours,
          anonymous_units_reserved
        FROM tasks
        WHERE id = $1
        "#,
    )
    .bind(task_id)
    .fetch_optional(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?
    else {
        tracing::info!(task_id = %task_id, "task was deleted before processing; acknowledging message");
        return Ok(());
    };

    if matches!(task.status.as_str(), "completed" | "failed" | "cancelled") {
        return Ok(());
    }

    let is_retry = task.status == "processing";

    let updated = sqlx::query(
        r#"
        UPDATE tasks
        SET status = 'processing', started_at = NOW()
        WHERE id = $1 AND status = 'pending'
        "#,
    )
    .bind(task_id)
    .execute(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;

    if updated.rows_affected() == 0 && task.status == "pending" {
        // Another worker might have taken it.
        return Ok(());
    }

    // Refresh task row after status change
    task.status = "processing".to_string();

    if is_retry {
        sqlx::query(
            "UPDATE task_files SET status = 'pending' WHERE task_id = $1 AND status = 'processing'",
        )
        .bind(task_id)
        .execute(&state.db)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?;
    }

    let compression_rate = task.compression_rate.and_then(|v| u8::try_from(v).ok());
    let level = compression_rate
        .map(compress::rate_to_level)
        .unwrap_or(compress::parse_level(&task.compression_level)?);
    let max_width = task.max_width.and_then(|v| u32::try_from(v).ok());
    let max_height = task.max_height.and_then(|v| u32::try_from(v).ok());

    let mut files: Vec<TaskFileProcRow> = sqlx::query_as(
        r#"
        SELECT
          id,
          input_path,
          original_format,
          output_format,
          status::text AS status
        FROM task_files
        WHERE task_id = $1
        ORDER BY created_at ASC
        "#,
    )
    .bind(task_id)
    .fetch_all(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务文件失败").with_source(err))?;

    let billing_ctx = if let Some(user_id) = task.user_id {
        Some(billing::get_user_billing(state, user_id).await?)
    } else {
        None
    };

    let anon_ip: Option<IpAddr> = task
        .client_ip
        .as_deref()
        .and_then(|s| s.parse::<IpAddr>().ok());

    let ctx = TaskContext {
        api_key_id: task.api_key_id,
        source: task.source.clone(),
        preserve_metadata: task.preserve_metadata,
        session_id: task.session_id.clone(),
        anon_ip,
        is_anonymous: task.user_id.is_none(),
        retention_hours: task.retention_hours,
        anonymous_quota_reserved: task.anonymous_units_reserved > 0,
    };

    let concurrency = state.config.worker_concurrency.max(1) as usize;
    let semaphore = Arc::new(Semaphore::new(concurrency));
    let mut join_set = JoinSet::new();

    for file in files.drain(..) {
        if file.status != "pending" {
            continue;
        }

        let permit = semaphore.clone().acquire_owned().await.unwrap();
        let state = state.clone();
        let ctx = ctx.clone();
        let billing_ctx = billing_ctx.clone();
        let file_id = file.id;

        join_set.spawn(async move {
            let _permit = permit;
            let result = process_task_file(
                state,
                task_id,
                file,
                level,
                compression_rate,
                max_width,
                max_height,
                ctx,
                billing_ctx,
            )
            .await;
            if let Err(err) = &result {
                tracing::error!(task_id = %task_id, file_id = %file_id, error = %err, "file processing failed");
            }
            result
        });
    }

    let mut first_error = None;
    while let Some(result) = join_set.join_next().await {
        match result {
            Ok(Ok(())) => {}
            Ok(Err(err)) => {
                if first_error.is_none() {
                    first_error = Some(err);
                }
            }
            Err(err) => {
                return Err(
                    AppError::new(ErrorCode::Internal, "文件处理线程异常退出").with_source(err)
                );
            }
        }
    }

    if let Some(err) = first_error {
        return Err(err);
    }

    finalize_task_status(state, task_id).await?;
    Ok(())
}

fn parse_image_fmt(value: &str) -> Result<compress::ImageFmt, AppError> {
    match value.trim().to_ascii_lowercase().as_str() {
        "png" => Ok(compress::ImageFmt::Png),
        "jpeg" | "jpg" => Ok(compress::ImageFmt::Jpeg),
        "webp" => Ok(compress::ImageFmt::Webp),
        "avif" => Ok(compress::ImageFmt::Avif),
        "gif" => Ok(compress::ImageFmt::Gif),
        "bmp" => Ok(compress::ImageFmt::Bmp),
        "tif" | "tiff" => Ok(compress::ImageFmt::Tiff),
        "ico" => Ok(compress::ImageFmt::Ico),
        _ => Err(AppError::new(ErrorCode::InvalidRequest, "未知图片格式")),
    }
}

async fn is_task_cancelled(state: &AppState, task_id: Uuid) -> Result<bool, AppError> {
    let status: Option<String> = sqlx::query_scalar("SELECT status::text FROM tasks WHERE id = $1")
        .bind(task_id)
        .fetch_optional(&state.db)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务状态失败").with_source(err))?;
    Ok(matches!(status.as_deref(), Some("cancelled")))
}

#[allow(clippy::too_many_arguments)]
async fn process_task_file(
    state: AppState,
    task_id: Uuid,
    file: TaskFileProcRow,
    level: compress::CompressionLevel,
    compression_rate: Option<u8>,
    max_width: Option<u32>,
    max_height: Option<u32>,
    ctx: TaskContext,
    billing_ctx: Option<billing::BillingContext>,
) -> Result<(), AppError> {
    if is_task_cancelled(&state, task_id).await? {
        return Ok(());
    }

    let updated = sqlx::query(
        "UPDATE task_files SET status = 'processing' WHERE id = $1 AND status = 'pending'",
    )
    .bind(file.id)
    .execute(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件处理状态失败").with_source(err))?;
    if updated.rows_affected() == 0 {
        return Ok(());
    }

    if is_task_cancelled(&state, task_id).await? {
        mark_file_failed(&state, task_id, file.id, "已取消").await?;
        return Ok(());
    }

    let Some(input_path) = file.input_path.clone() else {
        mark_file_failed(&state, task_id, file.id, "原文件不存在").await?;
        return Ok(());
    };

    let input_bytes = match tokio::fs::read(&input_path).await {
        Ok(v) => v,
        Err(_) => {
            mark_file_failed(&state, task_id, file.id, "读取原文件失败").await?;
            return Ok(());
        }
    };

    let format_in = match parse_image_fmt(&file.original_format) {
        Ok(format) => format,
        Err(err) => {
            mark_file_failed(&state, task_id, file.id, &err.message).await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        }
    };
    let format_out = match parse_image_fmt(&file.output_format) {
        Ok(format) => format,
        Err(err) => {
            mark_file_failed(&state, task_id, file.id, &err.message).await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        }
    };

    let original_size = input_bytes.len() as u64;
    let compressed = match compress::compress_image_bytes(
        &state,
        input_bytes,
        format_in,
        format_out,
        level,
        compression_rate,
        None, // target_size_bytes: worker 批量任务不支持精确大小
        max_width,
        max_height,
        ctx.preserve_metadata,
    )
    .await
    {
        Ok(v) => v,
        Err(err) => {
            mark_file_failed(&state, task_id, file.id, &err.message).await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        }
    };

    if is_task_cancelled(&state, task_id).await? {
        mark_file_failed(&state, task_id, file.id, "已取消").await?;
        let _ = tokio::fs::remove_file(&input_path).await;
        return Ok(());
    }

    let compressed_size = compressed.len() as u64;
    let saved_percent = if original_size == 0 {
        0.0
    } else {
        (original_size.saturating_sub(compressed_size) as f64) * 100.0 / (original_size as f64)
    };
    let skip_charge = compression_rate == Some(100)
        && format_in == format_out
        && max_width.is_none()
        && max_height.is_none();
    let charge_units = !skip_charge && compressed_size < original_size;

    let object_key = storage::result_key(
        ctx.retention_hours as i64,
        task_id,
        file.id,
        format_out.extension(),
    );
    let stored = match storage::store_bytes(
        &state,
        &object_key,
        compressed,
        format_out.content_type(),
    )
    .await
    {
        Ok(value) => value,
        Err(err) => {
            reset_file_for_retry(&state, file.id, "对象存储暂时不可用").await?;
            return Err(err);
        }
    };

    if ctx.is_anonymous && charge_units && !ctx.anonymous_quota_reserved {
        let Some(session_id) = ctx.session_id.as_deref() else {
            let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
            mark_file_failed(&state, task_id, file.id, "匿名任务缺少 session_id").await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        };
        let Some(ip) = ctx.anon_ip else {
            let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
            mark_file_failed(&state, task_id, file.id, "匿名任务缺少 client_ip").await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        };
        if let Err(err) = quota::consume_anonymous_units(&state, session_id, ip, 1).await {
            let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
            mark_file_failed(&state, task_id, file.id, &err.message).await?;
            let _ = tokio::fs::remove_file(&input_path).await;
            return Ok(());
        }
    }

    if is_task_cancelled(&state, task_id).await? {
        let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
        mark_file_failed(&state, task_id, file.id, "已取消").await?;
        let _ = tokio::fs::remove_file(&input_path).await;
        return Ok(());
    }

    if let Err(err) = finalize_file(
        &state,
        &billing_ctx,
        ctx.api_key_id,
        &ctx.source,
        task_id,
        file.id,
        &stored,
        original_size as i64,
        compressed_size as i64,
        saved_percent,
        format_in,
        format_out,
        charge_units,
    )
    .await
    {
        let _ = storage::delete_object(&state, &stored_locator(&stored)).await;
        mark_file_failed(&state, task_id, file.id, &err.message).await?;
        let _ = tokio::fs::remove_file(&input_path).await;
        return Ok(());
    }

    let _ = tokio::fs::remove_file(&input_path).await;
    Ok(())
}

#[allow(clippy::too_many_arguments)]
async fn finalize_file(
    state: &AppState,
    billing_ctx: &Option<billing::BillingContext>,
    api_key_id: Option<Uuid>,
    source: &str,
    task_id: Uuid,
    task_file_id: Uuid,
    stored: &storage::StoredObject,
    bytes_in: i64,
    bytes_out: i64,
    saved_percent: f64,
    format_in: compress::ImageFmt,
    format_out: compress::ImageFmt,
    charge_units: bool,
) -> Result<(), AppError> {
    let mut tx = state
        .db
        .begin()
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;

    // Paid users: charge before marking file completed (atomic w/ status update).
    if charge_units {
        if let Some(billing) = billing_ctx {
            charge_one_unit(
                &mut tx,
                billing,
                api_key_id,
                source,
                task_id,
                task_file_id,
                format_in,
                format_out,
                bytes_in as u64,
                bytes_out as u64,
            )
            .await?;
        }
    }

    sqlx::query(
        r#"
        UPDATE task_files
        SET storage_path = $2,
            storage_backend = $3,
            storage_endpoint_id = $4,
            storage_key = $5,
            storage_etag = $6,
            input_path = NULL,
            compressed_size = $7,
            saved_percent = $8,
            status = 'completed',
            completed_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(task_file_id)
    .bind(if stored.backend == "local" {
        Some(stored.key.as_str())
    } else {
        None
    })
    .bind(&stored.backend)
    .bind(stored.endpoint_id)
    .bind(&stored.key)
    .bind(&stored.etag)
    .bind(bytes_out)
    .bind(saved_percent)
    .execute(&mut *tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;

    sqlx::query(
        r#"
        UPDATE tasks
        SET completed_files = completed_files + 1,
            total_compressed_size = total_compressed_size + $2
        WHERE id = $1
        "#,
    )
    .bind(task_id)
    .bind(bytes_out)
    .execute(&mut *tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;

    tx.commit()
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;

    Ok(())
}

async fn mark_file_failed(
    state: &AppState,
    task_id: Uuid,
    task_file_id: Uuid,
    message: &str,
) -> Result<(), AppError> {
    let mut tx = state
        .db
        .begin()
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "开启事务失败").with_source(err))?;

    let updated = sqlx::query(
        r#"
        UPDATE task_files
        SET status = 'failed',
            error_message = $2,
            storage_path = NULL,
            input_path = NULL,
            completed_at = NOW()
        WHERE id = $1
          AND status NOT IN ('completed', 'failed')
        "#,
    )
    .bind(task_file_id)
    .bind(message)
    .execute(&mut *tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "更新文件失败").with_source(err))?;

    if updated.rows_affected() > 0 {
        sqlx::query(
            r#"
            UPDATE tasks
            SET failed_files = failed_files + 1
            WHERE id = $1
            "#,
        )
        .bind(task_id)
        .execute(&mut *tx)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务统计失败").with_source(err))?;
    }

    tx.commit()
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "提交事务失败").with_source(err))?;

    Ok(())
}

async fn reset_file_for_retry(
    state: &AppState,
    task_file_id: Uuid,
    message: &str,
) -> Result<(), AppError> {
    sqlx::query(
        "UPDATE task_files SET status = 'pending', error_message = $2 WHERE id = $1 AND status = 'processing'",
    )
    .bind(task_file_id)
    .bind(message)
    .execute(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "恢复待重试文件失败").with_source(err))?;
    Ok(())
}

async fn finalize_task_status(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
    let row: Option<(i32, i32, i32, String)> = sqlx::query_as(
        "SELECT total_files, completed_files, failed_files, status::text AS status FROM tasks WHERE id = $1",
    )
    .bind(task_id)
    .fetch_optional(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询任务失败").with_source(err))?;

    let Some((total, completed, failed, status)) = row else {
        return Ok(());
    };
    if status == "cancelled" {
        let paths: Vec<Option<String>> = sqlx::query_scalar(
            "SELECT input_path FROM task_files WHERE task_id = $1 AND status IN ('pending','processing')",
        )
        .bind(task_id)
        .fetch_all(&state.db)
        .await
        .unwrap_or_default();
        for p in paths.into_iter().flatten() {
            let _ = tokio::fs::remove_file(p).await;
        }

        let _ = sqlx::query(
            "UPDATE task_files SET status = 'failed', error_message = '已取消', storage_path = NULL, input_path = NULL, completed_at = NOW() WHERE task_id = $1 AND status IN ('pending','processing')",
        )
        .bind(task_id)
        .execute(&state.db)
        .await;

        let _ = sqlx::query(
            "UPDATE tasks SET failed_files = GREATEST(total_files - completed_files, 0), completed_at = NOW() WHERE id = $1 AND completed_at IS NULL",
        )
        .bind(task_id)
        .execute(&state.db)
        .await;
        return Ok(());
    }

    if completed + failed >= total && total > 0 {
        let final_status = if completed == 0 && failed == total {
            "failed"
        } else {
            "completed"
        };
        sqlx::query(
            "UPDATE tasks SET status = $2::task_status, completed_at = NOW() WHERE id = $1",
        )
        .bind(task_id)
        .bind(final_status)
        .execute(&state.db)
        .await
        .map_err(|err| AppError::new(ErrorCode::Internal, "更新任务状态失败").with_source(err))?;
    }

    Ok(())
}

#[allow(clippy::too_many_arguments)]
async fn charge_one_unit(
    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
    billing: &billing::BillingContext,
    api_key_id: Option<Uuid>,
    source: &str,
    task_id: Uuid,
    task_file_id: Uuid,
    format_in: compress::ImageFmt,
    format_out: compress::ImageFmt,
    bytes_in: u64,
    bytes_out: u64,
) -> Result<(), AppError> {
    quota::consume_user_unit(tx, billing, bytes_in, bytes_out).await?;

    sqlx::query(
        r#"
        INSERT INTO usage_events (
          user_id, api_key_id, source,
          task_id, task_file_id,
          units, bytes_in, bytes_out, format_in, format_out
        ) VALUES (
          $1, $2, $3::task_source,
          $4, $5,
          1, $6, $7, $8, $9
        )
        "#,
    )
    .bind(billing.user_id)
    .bind(api_key_id)
    .bind(source)
    .bind(task_id)
    .bind(task_file_id)
    .bind(bytes_in as i64)
    .bind(bytes_out as i64)
    .bind(format_in.as_str())
    .bind(format_out.as_str())
    .execute(&mut **tx)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "写入用量明细失败").with_source(err))?;

    Ok(())
}

async fn maintenance(state: &AppState) -> Result<(), AppError> {
    cleanup_expired_tasks(state).await?;
    cleanup_stale_zip_temp(state).await?;
    cleanup_expired_records(state).await?;
    Ok(())
}

async fn cleanup_stale_zip_temp(state: &AppState) -> Result<(), AppError> {
    let root = std::path::Path::new(&state.config.storage_path).join("tmp/zips");
    let mut entries = match tokio::fs::read_dir(&root).await {
        Ok(entries) => entries,
        Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
        Err(err) => {
            return Err(
                AppError::new(ErrorCode::StorageUnavailable, "读取 ZIP 临时目录失败")
                    .with_source(err),
            )
        }
    };
    let cutoff = std::time::SystemTime::now()
        .checked_sub(std::time::Duration::from_secs(6 * 60 * 60))
        .unwrap_or(std::time::UNIX_EPOCH);

    while let Some(entry) = entries.next_entry().await.map_err(|err| {
        AppError::new(ErrorCode::StorageUnavailable, "遍历 ZIP 临时目录失败").with_source(err)
    })? {
        let metadata = match entry.metadata().await {
            Ok(metadata) => metadata,
            Err(_) => continue,
        };
        if !metadata.is_dir()
            || metadata
                .modified()
                .map(|time| time >= cutoff)
                .unwrap_or(true)
        {
            continue;
        }
        let _ = tokio::fs::remove_dir_all(entry.path()).await;
    }
    Ok(())
}

async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
    let _ = sqlx::query("DELETE FROM idempotency_keys WHERE expires_at < NOW()")
        .execute(&state.db)
        .await;

    let _ = sqlx::query(
        "DELETE FROM email_verifications WHERE expires_at < NOW() AND verified_at IS NULL",
    )
    .execute(&state.db)
    .await;

    let _ = sqlx::query("DELETE FROM password_resets WHERE expires_at < NOW() - INTERVAL '7 days'")
        .execute(&state.db)
        .await;

    let _ =
        sqlx::query("DELETE FROM webhook_events WHERE received_at < NOW() - INTERVAL '90 days'")
            .execute(&state.db)
            .await;

    let _ = sqlx::query(
        r#"
        DELETE FROM storage_endpoints e
        WHERE e.deleted_at < NOW() - INTERVAL '30 days'
          AND NOT EXISTS (SELECT 1 FROM task_files f WHERE f.storage_endpoint_id = e.id)
          AND NOT EXISTS (SELECT 1 FROM tasks t WHERE t.zip_storage_endpoint_id = e.id)
        "#,
    )
    .execute(&state.db)
    .await;

    Ok(())
}

async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
    let task_ids: Vec<Uuid> =
        sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200")
            .fetch_all(&state.db)
            .await
            .unwrap_or_default();

    if task_ids.is_empty() {
        return Ok(());
    }

    for task_id in task_ids {
        if let Err(err) = cleanup_expired_task(state, task_id).await {
            tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred");
        }
    }

    Ok(())
}

async fn cleanup_expired_task(state: &AppState, task_id: Uuid) -> Result<(), AppError> {
    let files: Vec<CleanupFileRow> = sqlx::query_as(
        r#"
        SELECT storage_backend, storage_endpoint_id,
               COALESCE(storage_key, storage_path) AS storage_key,
               input_path
        FROM task_files
        WHERE task_id = $1
        "#,
    )
    .bind(task_id)
    .fetch_all(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务文件失败").with_source(err))?;

    for file in files {
        if let Some(key) = file.storage_key {
            storage::delete_object(
                state,
                &storage::ObjectLocator {
                    backend: file.storage_backend,
                    endpoint_id: file.storage_endpoint_id,
                    key,
                },
            )
            .await?;
        }
        if let Some(input_path) = file.input_path {
            let _ = tokio::fs::remove_file(input_path).await;
        }
    }

    let zip: Option<CleanupZipRow> = sqlx::query_as(
        "SELECT zip_storage_backend, zip_storage_endpoint_id, zip_storage_key FROM tasks WHERE id = $1",
    )
    .bind(task_id)
    .fetch_optional(&state.db)
    .await
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询过期 ZIP 失败").with_source(err))?;
    if let Some(zip) = zip {
        if let (Some(backend), Some(key)) = (zip.zip_storage_backend, zip.zip_storage_key) {
            storage::delete_object(
                state,
                &storage::ObjectLocator {
                    backend,
                    endpoint_id: zip.zip_storage_endpoint_id,
                    key,
                },
            )
            .await?;
        }
    }

    let legacy_zip_path = format!("{}/zips/{task_id}.zip", state.config.storage_path);
    let _ = tokio::fs::remove_file(legacy_zip_path).await;
    let orig_dir = format!("{}/orig/{task_id}", state.config.storage_path);
    let _ = tokio::fs::remove_dir_all(orig_dir).await;

    sqlx::query("DELETE FROM tasks WHERE id = $1 AND expires_at < NOW()")
        .bind(task_id)
        .execute(&state.db)
        .await
        .map_err(|err| {
            AppError::new(ErrorCode::Internal, "删除过期任务记录失败").with_source(err)
        })?;
    Ok(())
}

fn stored_locator(stored: &storage::StoredObject) -> storage::ObjectLocator {
    storage::ObjectLocator {
        backend: stored.backend.clone(),
        endpoint_id: stored.endpoint_id,
        key: stored.key.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn third_delivery_moves_message_to_dead_letter() {
        assert!(!should_dead_letter(1));
        assert!(!should_dead_letter(2));
        assert!(should_dead_letter(3));
        assert!(should_dead_letter(10));
    }
}