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::{Value, json};
use shared_contracts::external_generation::{
    ExternalGenerationJobStatus, ExternalGenerationJobStatusRecord,
    ExternalGenerationJobStatusResponse, ExternalGenerationQueueOverview,
    ExternalGenerationQueueOverviewResponse, ExternalGenerationTaskAcknowledgeRequest,
    ExternalGenerationTaskAcknowledgeResponse, ExternalGenerationTaskListResponse,
    ExternalGenerationTaskRecord,
};
use spacetime_client::{
    ExternalGenerationJobAcknowledgeRecordInput, ExternalGenerationJobGetRecordInput,
    ExternalGenerationJobListRecord, ExternalGenerationJobListRecordInput,
    ExternalGenerationJobRecord, 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_jobs(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_jobs(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_jobs(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(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(job),
        },
    ))
}

fn map_external_generation_queue_overview(
    list: &ExternalGenerationJobListRecord,
) -> 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: ExternalGenerationJobRecord,
) -> ExternalGenerationJobStatusRecord {
    let (status, phase_detail, progress) = match job.status.as_str() {
        "completed" => (ExternalGenerationJobStatus::Completed, "生成已完成。", 100),
        "running" => (ExternalGenerationJobStatus::Running, "正在生成。", 35),
        "failed" => (ExternalGenerationJobStatus::Failed, "生成失败。", 0),
        _ => (ExternalGenerationJobStatus::Queued, "排队中。", 8),
    };

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

fn map_external_generation_task_record(
    job: ExternalGenerationJobRecord,
) -> ExternalGenerationTaskRecord {
    let status_record = map_external_generation_job_status(job.clone());
    let request_prompt = extract_external_generation_request_prompt(&job.request_payload_json);
    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,
        status: status_record.status,
        phase_label: status_record.phase_label,
        phase_detail: status_record.phase_detail,
        progress: status_record.progress,
        error: status_record.error,
        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 extract_external_generation_request_prompt(request_payload_json: &str) -> Option<String> {
    let payload: Value = serde_json::from_str(request_payload_json).ok()?;
    for key in ["prompt", "promptText", "spritesheetLabel"] {
        if let Some(prompt) = payload
            .get(key)
            .and_then(Value::as_str)
            .and_then(normalize_external_generation_prompt_text)
        {
            return Some(prompt);
        }
    }
    if let Some(prompt) = payload
        .get("iconDescriptions")
        .and_then(Value::as_array)
        .map(|items| {
            items
                .iter()
                .filter_map(Value::as_str)
                .filter_map(normalize_external_generation_prompt_text)
                .collect::<Vec<_>>()
                .join("、")
        })
        .and_then(|value| normalize_external_generation_prompt_text(&value))
    {
        return Some(prompt);
    }
    payload
        .get("generationInputs")
        .and_then(|value| value.get("fields"))
        .and_then(Value::as_array)
        .and_then(|fields| {
            fields.iter().find_map(|field| {
                let title = field.get("title").and_then(Value::as_str)?.trim();
                if !matches!(title, "prompt" | "gpt_description_prompt") {
                    return None;
                }
                field
                    .get("value")
                    .and_then(Value::as_str)
                    .and_then(normalize_external_generation_prompt_text)
            })
        })
}

fn normalize_external_generation_prompt_text(value: &str) -> Option<String> {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        None
    } else {
        Some(trimmed.to_string())
    }
}

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 extracts_external_generation_request_prompt_from_payload() {
        assert_eq!(
            extract_external_generation_request_prompt(
                r#"{"prompt":"  发光主视觉  ","sourceImageSrc":"data:image/png;base64,secret"}"#,
            )
            .as_deref(),
            Some("发光主视觉"),
        );
        assert_eq!(
            extract_external_generation_request_prompt(
                r#"{"iconDescriptions":["返回按钮"," 设置按钮 "],"referenceImageSrc":"data:image/png;base64,secret"}"#,
            )
            .as_deref(),
            Some("返回按钮、设置按钮"),
        );
    }
}