Genarrative/server-rs/crates/api-server/src/external_generation.rs

use axum::{
    Json,
    extract::{Extension, Path, Query, State},
    http::StatusCode,
    response::Response,
};
use serde::Deserialize;
use serde_json::json;
use shared_contracts::external_generation::{
    ExternalGenerationJobStatus, ExternalGenerationJobStatusDetailRecord,
    ExternalGenerationJobStatusRecord, ExternalGenerationJobStatusResponse,
    ExternalGenerationQueueOverview, ExternalGenerationQueueOverviewResponse,
    ExternalGenerationTaskAcknowledgeRequest, ExternalGenerationTaskAcknowledgeResponse,
    ExternalGenerationTaskListResponse, ExternalGenerationTaskRecord,
};
use spacetime_client::{
    ExternalGenerationJobAcknowledgeRecordInput, ExternalGenerationJobGetRecordInput,
    ExternalGenerationJobListRecordInput, ExternalGenerationJobSummaryListRecord,
    ExternalGenerationJobSummaryRecord, SpacetimeClientError,
};

use crate::{
    api_response::json_success_body, auth::AuthenticatedAccessToken, http_error::AppError,
    request_context::RequestContext, state::AppState,
};

const EXTERNAL_GENERATION_PROVIDER: &str = "external_generation";

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExternalGenerationTaskListQuery {
    limit: Option<u32>,
    include_acknowledged_terminal: Option<bool>,
    statuses: Option<String>,
}

pub async fn get_external_generation_queue_overview(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    Extension(authenticated): Extension<AuthenticatedAccessToken>,
) -> Result<Json<serde_json::Value>, Response> {
    let owner_user_id = authenticated.claims().user_id().to_string();
    let list = state
        .spacetime_client()
        .list_external_generation_job_summaries(ExternalGenerationJobListRecordInput {
            owner_user_id,
            limit: 1,
            include_acknowledged_terminal: false,
            statuses: Vec::new(),
        })
        .await
        .map_err(|error| external_generation_error_response(&request_context, error))?;

    Ok(json_success_body(
        Some(&request_context),
        ExternalGenerationQueueOverviewResponse {
            overview: map_external_generation_queue_overview(&list),
        },
    ))
}

pub async fn list_external_generation_tasks(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    Extension(authenticated): Extension<AuthenticatedAccessToken>,
    Query(query): Query<ExternalGenerationTaskListQuery>,
) -> Result<Json<serde_json::Value>, Response> {
    let owner_user_id = authenticated.claims().user_id().to_string();
    let requested_limit = query.limit.unwrap_or(50).clamp(1, 100);
    let status_filter = parse_external_generation_status_filter(query.statuses.as_deref());
    let statuses = external_generation_status_filter_input(&status_filter);
    let list = state
        .spacetime_client()
        .list_external_generation_job_summaries(ExternalGenerationJobListRecordInput {
            owner_user_id,
            limit: requested_limit,
            include_acknowledged_terminal: query.include_acknowledged_terminal.unwrap_or(false),
            statuses,
        })
        .await
        .map_err(|error| external_generation_error_response(&request_context, error))?;
    let overview = map_external_generation_queue_overview(&list);
    let mut tasks: Vec<_> = list
        .jobs
        .into_iter()
        .map(map_external_generation_task_record)
        .filter(|task| {
            status_filter.is_empty() || status_filter.iter().any(|status| status == &task.status)
        })
        .collect();
    tasks.truncate(requested_limit as usize);

    Ok(json_success_body(
        Some(&request_context),
        ExternalGenerationTaskListResponse { overview, tasks },
    ))
}

pub async fn acknowledge_external_generation_tasks(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    Extension(authenticated): Extension<AuthenticatedAccessToken>,
    Json(payload): Json<ExternalGenerationTaskAcknowledgeRequest>,
) -> Result<Json<serde_json::Value>, Response> {
    let owner_user_id = authenticated.claims().user_id().to_string();
    let acknowledged = state
        .spacetime_client()
        .acknowledge_external_generation_job_summaries(
            ExternalGenerationJobAcknowledgeRecordInput {
                owner_user_id,
                job_ids: payload.job_ids,
                acknowledged_at_micros: current_utc_micros(),
            },
        )
        .await
        .map_err(|error| external_generation_error_response(&request_context, error))?;

    Ok(json_success_body(
        Some(&request_context),
        ExternalGenerationTaskAcknowledgeResponse {
            acknowledged_tasks: acknowledged
                .jobs
                .into_iter()
                .map(map_external_generation_task_record)
                .collect(),
        },
    ))
}

pub async fn get_external_generation_job_status(
    State(state): State<AppState>,
    Extension(request_context): Extension<RequestContext>,
    Extension(authenticated): Extension<AuthenticatedAccessToken>,
    Path(job_id): Path<String>,
) -> Result<Json<serde_json::Value>, Response> {
    let owner_user_id = authenticated.claims().user_id().to_string();
    let job = state
        .spacetime_client()
        .get_external_generation_job_summary(ExternalGenerationJobGetRecordInput {
            job_id,
            owner_user_id,
        })
        .await
        .map_err(|error| external_generation_error_response(&request_context, error))?;

    Ok(json_success_body(
        Some(&request_context),
        ExternalGenerationJobStatusResponse {
            job: map_external_generation_job_status_detail(job),
        },
    ))
}

fn map_external_generation_queue_overview(
    list: &ExternalGenerationJobSummaryListRecord,
) -> ExternalGenerationQueueOverview {
    ExternalGenerationQueueOverview {
        pending_count: list.pending_count,
        running_count: list.running_count,
        unacknowledged_terminal_count: list.unacknowledged_terminal_count,
        updated_at_micros: list.now_micros,
    }
}

fn map_external_generation_job_status(
    job: ExternalGenerationJobSummaryRecord,
) -> ExternalGenerationJobStatusRecord {
    let (status, phase_detail, progress, error) = match job.status.as_str() {
        "completed" => (
            ExternalGenerationJobStatus::Completed,
            "生成已完成。",
            100,
            None,
        ),
        "running" => (
            ExternalGenerationJobStatus::Running,
            "正在生成。",
            35,
            None,
        ),
        "failed" => (
            ExternalGenerationJobStatus::Failed,
            "生成失败。",
            0,
            job.last_error_message.clone(),
        ),
        _ => (
            ExternalGenerationJobStatus::Queued,
            "排队中。",
            8,
            None,
        ),
    };

    ExternalGenerationJobStatusRecord {
        operation_id: job.job_id.clone(),
        status,
        phase_label: job.request_label.clone(),
        phase_detail: phase_detail.to_string(),
        progress,
        error,
        updated_at_micros: job.updated_at_micros,
    }
}

fn map_external_generation_job_status_detail(
    job: ExternalGenerationJobSummaryRecord,
) -> ExternalGenerationJobStatusDetailRecord {
    let warning = job.warning_message.clone();
    ExternalGenerationJobStatusDetailRecord {
        status: map_external_generation_job_status(job),
        warning,
    }
}

fn map_external_generation_task_record(
    job: ExternalGenerationJobSummaryRecord,
) -> ExternalGenerationTaskRecord {
    let status_record = map_external_generation_job_status(job.clone());
    ExternalGenerationTaskRecord {
        job_id: job.job_id,
        job_kind: job.job_kind,
        source_module: job.source_module,
        source_entity_id: job.source_entity_id,
        request_label: job.request_label,
        request_prompt: job.request_prompt,
        status: status_record.status,
        phase_label: status_record.phase_label,
        phase_detail: status_record.phase_detail,
        progress: status_record.progress,
        error: status_record.error,
        warning: job.warning_message,
        price_mud_points: job.price_mud_points,
        refund_ledger_id: job.refund_ledger_id,
        notification_acknowledged_at: job.notification_acknowledged_at,
        created_at: job.created_at,
        started_at: job.started_at,
        completed_at: job.completed_at,
        updated_at: job.updated_at,
        updated_at_micros: job.updated_at_micros,
    }
}

fn parse_external_generation_status_filter(
    statuses: Option<&str>,
) -> Vec<ExternalGenerationJobStatus> {
    statuses
        .unwrap_or_default()
        .split(',')
        .filter_map(|status| match status.trim() {
            "queued" => Some(ExternalGenerationJobStatus::Queued),
            "running" => Some(ExternalGenerationJobStatus::Running),
            "completed" => Some(ExternalGenerationJobStatus::Completed),
            "failed" => Some(ExternalGenerationJobStatus::Failed),
            _ => None,
        })
        .collect()
}

fn external_generation_status_filter_input(
    statuses: &[ExternalGenerationJobStatus],
) -> Vec<String> {
    statuses
        .iter()
        .map(|status| match status {
            ExternalGenerationJobStatus::Queued => "queued",
            ExternalGenerationJobStatus::Running => "running",
            ExternalGenerationJobStatus::Completed => "completed",
            ExternalGenerationJobStatus::Failed => "failed",
        })
        .map(str::to_string)
        .collect()
}

fn current_utc_micros() -> i64 {
    shared_kernel::offset_datetime_to_unix_micros(time::OffsetDateTime::now_utc())
}

fn external_generation_error_response(
    request_context: &RequestContext,
    error: SpacetimeClientError,
) -> Response {
    AppError::from_status(StatusCode::BAD_GATEWAY)
        .with_details(json!({
            "provider": EXTERNAL_GENERATION_PROVIDER,
            "message": error.to_string(),
        }))
        .into_response_with_context(Some(request_context))
}

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

    #[test]
    fn parses_external_generation_status_filter() {
        assert_eq!(
            parse_external_generation_status_filter(Some("running,queued,unknown")),
            vec![
                ExternalGenerationJobStatus::Running,
                ExternalGenerationJobStatus::Queued
            ]
        );
    }

    #[test]
    fn maps_task_from_payload_free_summary_projection() {
        let summary = ExternalGenerationJobSummaryRecord {
            job_id: "task-1".to_string(),
            job_kind: "editor_image_generation".to_string(),
            owner_user_id: "user-1".to_string(),
            source_module: "editor-canvas".to_string(),
            source_entity_id: "project-1".to_string(),
            request_label: "图片生成".to_string(),
            request_prompt: Some("发光主视觉".to_string()),
            status: "completed".to_string(),
            last_error_message: Some("上一次尝试失败".to_string()),
            created_at: "2026-07-10T08:00:00Z".to_string(),
            started_at: Some("2026-07-10T08:00:01Z".to_string()),
            completed_at: Some("2026-07-10T08:00:10Z".to_string()),
            updated_at: "2026-07-10T08:00:10Z".to_string(),
            updated_at_micros: 1_000,
            price_mud_points: 4,
            refund_ledger_id: None,
            notification_acknowledged_at: None,
            notification_acknowledged_at_micros: None,
            warning_message: Some("连通域数量不足".to_string()),
        };
        let status = map_external_generation_job_status_detail(summary.clone());
        let task = map_external_generation_task_record(summary);

        assert_eq!(status.warning.as_deref(), Some("连通域数量不足"));
        assert!(status.status.error.is_none());
        assert_eq!(task.request_prompt.as_deref(), Some("发光主视觉"));
        assert_eq!(task.status, ExternalGenerationJobStatus::Completed);
        assert_eq!(task.warning.as_deref(), Some("连通域数量不足"));
        assert!(task.error.is_none());
    }
}