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

use axum::http::StatusCode;
use platform_image::{
    DownloadedImage, GeneratedImages, PlatformImageError, PlatformImageStatusHint, ReferenceImage,
    VECTOR_ENGINE_PROVIDER, VectorEngineImageSettings, build_vector_engine_image_http_client,
    create_vector_engine_image_edit_with_references_and_model,
    create_vector_engine_image_generation, create_vector_engine_image_generation_with_model,
    create_vector_engine_nanobanana_generate_content,
};
#[cfg(test)]
use platform_image::{
    build_vector_engine_image_request_body, vector_engine_images_edit_url,
    vector_engine_images_generation_url,
};
use serde_json::{Value, json};
use std::time::Instant;
use time::OffsetDateTime;

use crate::{
    external_api_audit::{
        build_external_api_failure_draft_from_platform_image_audit, record_external_api_failure,
    },
    http_error::AppError,
    request_context::RequestContext,
    state::AppState,
    tracking::record_external_generation_run_after_success,
};

pub(crate) use platform_image::GPT_IMAGE_2_MODEL;
#[cfg(test)]
use platform_image::VECTOR_ENGINE_GPT_IMAGE_2_MODEL;

pub(crate) type OpenAiGeneratedImages = GeneratedImages;
pub(crate) type DownloadedOpenAiImage = DownloadedImage;
pub(crate) type OpenAiReferenceImage = ReferenceImage;

#[derive(Clone)]
pub(crate) struct OpenAiImageSettings {
    pub base_url: String,
    pub api_key: String,
    pub request_timeout_ms: u64,
    pub request_deadline: Option<Instant>,
    pub external_api_audit_state: Option<AppState>,
    pub external_api_audit_user_id: Option<String>,
    pub external_api_audit_profile_id: Option<String>,
    pub external_api_audit_request_id: Option<String>,
}

impl std::fmt::Debug for OpenAiImageSettings {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("OpenAiImageSettings")
            .field("base_url", &self.base_url)
            .field("api_key", &"<redacted>")
            .field("request_timeout_ms", &self.request_timeout_ms)
            .field("request_deadline_enabled", &self.request_deadline.is_some())
            .field(
                "external_api_audit_enabled",
                &self.external_api_audit_state.is_some(),
            )
            .field(
                "external_api_audit_user_id",
                &self.external_api_audit_user_id,
            )
            .field(
                "external_api_audit_profile_id",
                &self.external_api_audit_profile_id,
            )
            .field(
                "external_api_audit_request_id",
                &self.external_api_audit_request_id,
            )
            .finish()
    }
}

// 中文注释：api-server 只负责配置、审计和 HTTP envelope，Tiantoken 的 OpenAI-compatible
// 图片协议细节统一由 platform-image provider 承接。
pub(crate) fn require_openai_image_settings(
    state: &AppState,
) -> Result<OpenAiImageSettings, AppError> {
    let base_url = state.tiantoken_base_url().trim().trim_end_matches('/');
    if base_url.is_empty() {
        return Err(
            AppError::from_status(StatusCode::SERVICE_UNAVAILABLE).with_details(json!({
                "provider": "tiantoken",
                "reason": "TIANTOKEN_BASE_URL 未配置",
            })),
        );
    }

    let api_key = state
        .tiantoken_api_key()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| {
            AppError::from_status(StatusCode::SERVICE_UNAVAILABLE).with_details(json!({
                "provider": "tiantoken",
                "reason": "TIANTOKEN_API_KEY 未配置",
            }))
        })?;

    Ok(OpenAiImageSettings {
        base_url: base_url.to_string(),
        api_key: api_key.to_string(),
        request_timeout_ms: state.config.vector_engine_image_request_timeout_ms.max(1),
        request_deadline: None,
        external_api_audit_state: Some(state.clone()),
        external_api_audit_user_id: None,
        external_api_audit_profile_id: None,
        external_api_audit_request_id: None,
    })
}

pub(crate) fn build_openai_image_http_client(
    settings: &OpenAiImageSettings,
) -> Result<reqwest::Client, AppError> {
    build_vector_engine_image_http_client(&settings.provider_settings())
        .map_err(map_platform_image_error)
}

pub(crate) async fn create_openai_image_generation(
    http_client: &reqwest::Client,
    settings: &OpenAiImageSettings,
    prompt: &str,
    negative_prompt: Option<&str>,
    size: &str,
    candidate_count: u32,
    reference_images: &[String],
    failure_context: &str,
) -> Result<OpenAiGeneratedImages, AppError> {
    let started_at_micros = current_utc_micros();
    let request_payload = json!({
        "size": size,
        "candidateCount": candidate_count,
        "promptChars": prompt.chars().count(),
        "negativePromptChars": negative_prompt.map(str::chars).map(Iterator::count),
        "referenceImageCount": reference_images.len(),
    });
    let result = create_vector_engine_image_generation(
        http_client,
        &settings.provider_settings(),
        prompt,
        negative_prompt,
        size,
        candidate_count,
        reference_images,
        failure_context,
    )
    .await;
    map_platform_image_result(
        settings,
        result,
        "image_generation",
        failure_context,
        request_payload,
        started_at_micros,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn create_openai_image_generation_with_model(
    http_client: &reqwest::Client,
    settings: &OpenAiImageSettings,
    model: &str,
    prompt: &str,
    negative_prompt: Option<&str>,
    size: &str,
    candidate_count: u32,
    reference_images: &[String],
    failure_context: &str,
) -> Result<OpenAiGeneratedImages, AppError> {
    let started_at_micros = current_utc_micros();
    let request_payload = json!({
        "model": model,
        "size": size,
        "candidateCount": candidate_count,
        "promptChars": prompt.chars().count(),
        "negativePromptChars": negative_prompt.map(str::chars).map(Iterator::count),
        "referenceImageCount": reference_images.len(),
    });
    let result = create_vector_engine_image_generation_with_model(
        http_client,
        &settings.provider_settings(),
        model,
        prompt,
        negative_prompt,
        size,
        candidate_count,
        reference_images,
        failure_context,
    )
    .await;
    map_platform_image_result(
        settings,
        result,
        "image_generation",
        failure_context,
        request_payload,
        started_at_micros,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn create_openai_nanobanana_generate_content(
    http_client: &reqwest::Client,
    settings: &OpenAiImageSettings,
    model: &str,
    prompt: &str,
    negative_prompt: Option<&str>,
    aspect_ratio: &str,
    image_size: &str,
    reference_images: &[OpenAiReferenceImage],
    failure_context: &str,
) -> Result<OpenAiGeneratedImages, AppError> {
    let started_at_micros = current_utc_micros();
    let request_payload = json!({
        "model": model,
        "aspectRatio": aspect_ratio,
        "imageSize": image_size,
        "promptChars": prompt.chars().count(),
        "negativePromptChars": negative_prompt.map(str::chars).map(Iterator::count),
        "referenceImageCount": reference_images.len(),
    });
    let result = create_vector_engine_nanobanana_generate_content(
        http_client,
        &settings.provider_settings(),
        model,
        prompt,
        negative_prompt,
        aspect_ratio,
        image_size,
        reference_images,
        failure_context,
    )
    .await;
    map_platform_image_result(
        settings,
        result,
        "nanobanana_generate_content",
        failure_context,
        request_payload,
        started_at_micros,
    )
    .await
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn create_openai_image_edit_with_references_and_model(
    http_client: &reqwest::Client,
    settings: &OpenAiImageSettings,
    model: &str,
    prompt: &str,
    negative_prompt: Option<&str>,
    size: &str,
    candidate_count: u32,
    reference_images: &[OpenAiReferenceImage],
    failure_context: &str,
) -> Result<OpenAiGeneratedImages, AppError> {
    let started_at_micros = current_utc_micros();
    let request_payload = json!({
        "model": model,
        "size": size,
        "promptChars": prompt.chars().count(),
        "negativePromptChars": negative_prompt.map(str::chars).map(Iterator::count),
        "referenceImageCount": reference_images.len(),
    });
    let result = create_vector_engine_image_edit_with_references_and_model(
        http_client,
        &settings.provider_settings(),
        model,
        prompt,
        negative_prompt,
        size,
        candidate_count,
        reference_images,
        failure_context,
    )
    .await;
    map_platform_image_result(
        settings,
        result,
        "image_edit_with_references",
        failure_context,
        request_payload,
        started_at_micros,
    )
    .await
}

#[cfg(test)]
pub(crate) fn build_openai_image_request_body(
    prompt: &str,
    negative_prompt: Option<&str>,
    size: &str,
    candidate_count: u32,
    reference_images: &[String],
) -> Value {
    build_vector_engine_image_request_body(
        prompt,
        negative_prompt,
        size,
        candidate_count,
        reference_images,
    )
}

impl OpenAiImageSettings {
    pub(crate) fn with_external_api_audit_context(
        mut self,
        request_context: &RequestContext,
        user_id: Option<String>,
        profile_id: Option<String>,
    ) -> Self {
        self.external_api_audit_user_id = user_id;
        self.external_api_audit_profile_id = profile_id;
        self.external_api_audit_request_id = Some(request_context.request_id().to_string());
        self.request_deadline = request_context.external_call_deadline();
        self
    }

    pub(crate) fn provider_settings(&self) -> VectorEngineImageSettings {
        VectorEngineImageSettings {
            base_url: self.base_url.clone(),
            api_key: self.api_key.clone(),
            request_timeout_ms: self.request_timeout_ms.max(1),
            request_deadline: self.request_deadline,
        }
    }
}

async fn map_platform_image_result(
    settings: &OpenAiImageSettings,
    result: Result<OpenAiGeneratedImages, PlatformImageError>,
    operation: &'static str,
    failure_context: &str,
    request_payload: Value,
    started_at_micros: i64,
) -> Result<OpenAiGeneratedImages, AppError> {
    match result {
        Ok(value) => {
            for audit in &value.recovered_failure_audits {
                record_openai_image_failure_audit_if_configured(settings, audit).await;
            }
            if let Some(state) = settings.external_api_audit_state.as_ref() {
                record_external_generation_run_after_success(
                    state,
                    VECTOR_ENGINE_PROVIDER,
                    operation,
                    failure_context,
                    request_payload,
                    started_at_micros,
                    true,
                    None,
                    Some(value.task_id.clone()),
                    Some(json!({
                        "imageCount": value.images.len(),
                        "actualPromptChars": value.actual_prompt.as_ref().map(|prompt| prompt.chars().count()),
                        "recoveredFailureCount": value.recovered_failure_audits.len(),
                    })),
                )
                .await;
            }
            Ok(value)
        }
        Err(error) => {
            for audit in error.recovered_failure_audits() {
                record_openai_image_failure_audit_if_configured(settings, audit).await;
            }
            if let Some(state) = settings.external_api_audit_state.as_ref() {
                record_external_generation_run_after_success(
                    state,
                    VECTOR_ENGINE_PROVIDER,
                    operation,
                    failure_context,
                    request_payload,
                    started_at_micros,
                    false,
                    Some(error.message().to_string()),
                    None,
                    None,
                )
                .await;
            }
            record_openai_image_failure_if_configured(settings, &error).await;
            Err(map_platform_image_error(error))
        }
    }
}
pub(crate) async fn record_openai_image_failure_if_configured(
    settings: &OpenAiImageSettings,
    error: &PlatformImageError,
) {
    let Some(audit) = error.audit() else {
        return;
    };
    record_openai_image_failure_audit_if_configured(settings, audit).await;
}

async fn record_openai_image_failure_audit_if_configured(
    settings: &OpenAiImageSettings,
    audit: &platform_image::PlatformImageFailureAudit,
) {
    let Some(state) = settings.external_api_audit_state.as_ref() else {
        return;
    };
    let draft = build_external_api_failure_draft_from_platform_image_audit(audit)
        .with_user_id(settings.external_api_audit_user_id.clone())
        .with_profile_id(settings.external_api_audit_profile_id.clone())
        .with_request_id(settings.external_api_audit_request_id.clone());
    record_external_api_failure(state, draft).await;
}

pub(crate) fn map_platform_image_error(error: PlatformImageError) -> AppError {
    let error = error.into_final_error();
    let status = match error.status_hint() {
        PlatformImageStatusHint::BadRequest => StatusCode::BAD_REQUEST,
        PlatformImageStatusHint::ServiceUnavailable => StatusCode::SERVICE_UNAVAILABLE,
        PlatformImageStatusHint::BadGateway => StatusCode::BAD_GATEWAY,
        PlatformImageStatusHint::GatewayTimeout => StatusCode::GATEWAY_TIMEOUT,
    };

    let mut details = json!({
        "provider": error.provider(),
        "message": error.message(),
    });

    match &error {
        PlatformImageError::InvalidConfig { .. } | PlatformImageError::InvalidRequest { .. } => {}
        PlatformImageError::Request {
            endpoint,
            timeout,
            connect,
            request,
            body,
            status_code,
            source,
            ..
        } => {
            details["endpoint"] = json!(endpoint);
            details["timeout"] = json!(timeout);
            details["connect"] = json!(connect);
            details["request"] = json!(request);
            details["body"] = json!(body);
            details["status"] = json!(status_code);
            details["source"] = json!(source);
        }
        PlatformImageError::Upstream {
            upstream_status,
            raw_excerpt,
            ..
        } => {
            details["upstreamStatus"] = json!(upstream_status);
            details["rawExcerpt"] = json!(raw_excerpt);
        }
        PlatformImageError::ResponseParse { raw_excerpt, .. } => {
            details["rawExcerpt"] = json!(raw_excerpt);
        }
        PlatformImageError::MissingImage { .. } => {}
        PlatformImageError::FallbackFailed { .. } => {
            unreachable!("fallback wrapper should be removed before HTTP error mapping")
        }
    }

    if let Some(audit) = error.audit() {
        details["endpoint"] = json!(audit.endpoint);
        details["failureStage"] = json!(audit.failure_stage);
        details["statusClass"] = json!(audit.status_class);
        details["retryable"] = json!(audit.retryable);
        details["timeout"] = json!(audit.timeout);
        details["latencyMs"] = json!(audit.latency_ms);
        details["promptChars"] = json!(audit.prompt_chars);
        details["referenceImageCount"] = json!(audit.reference_image_count);
        details["imageModel"] = json!(audit.image_model);
        details["rawExcerpt"] = json!(audit.raw_excerpt);
        details["errorSource"] = json!(audit.error_source);
    }

    AppError::from_status(status).with_details(details)
}

#[cfg(test)]
fn vector_engine_images_generation_url_for_test(settings: &OpenAiImageSettings) -> String {
    vector_engine_images_generation_url(&settings.provider_settings())
}

#[cfg(test)]
fn vector_engine_images_edit_url_for_test(settings: &OpenAiImageSettings) -> String {
    vector_engine_images_edit_url(&settings.provider_settings())
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};

    #[test]
    fn external_api_audit_context_forwards_external_call_deadline() {
        let deadline = Instant::now() + std::time::Duration::from_secs(30);
        let request_context = RequestContext::new(
            "request-1".to_string(),
            "external-generation-worker editor_image_edit".to_string(),
            std::time::Duration::ZERO,
            false,
        )
        .with_external_call_deadline(deadline);
        let settings = OpenAiImageSettings {
            base_url: "https://vector.example".to_string(),
            api_key: "test-key".to_string(),
            request_timeout_ms: 1_000_000,
            request_deadline: None,
            external_api_audit_state: None,
            external_api_audit_user_id: None,
            external_api_audit_profile_id: None,
            external_api_audit_request_id: None,
        }
        .with_external_api_audit_context(&request_context, None, None);

        assert_eq!(settings.request_deadline, Some(deadline));
        assert_eq!(
            settings.provider_settings().request_deadline,
            Some(deadline)
        );
    }

    #[test]
    fn gpt_image_2_generation_request_uses_create_model_without_reference_images() {
        let body = build_openai_image_request_body(
            "雾海神殿",
            Some("文字，水印"),
            "16:9",
            2,
            &["data:image/png;base64,abcd".to_string()],
        );

        assert_eq!(body["model"], GPT_IMAGE_2_MODEL);
        assert_eq!(body["size"], "1536x1024");
        assert_eq!(body["n"], 2);
        assert!(body.get("official_fallback").is_none());
        assert!(body.get("image").is_none());
        assert!(body["prompt"].as_str().unwrap_or_default().contains("避免"));
    }

    #[test]
    fn vector_engine_generation_url_normalizes_base_url() {
        let root_settings = OpenAiImageSettings {
            base_url: "https://vector.example".to_string(),
            api_key: "test-key".to_string(),
            request_timeout_ms: 1_000_000,
            request_deadline: None,
            external_api_audit_state: None,
            external_api_audit_user_id: None,
            external_api_audit_profile_id: None,
            external_api_audit_request_id: None,
        };
        let v1_settings = OpenAiImageSettings {
            base_url: "https://vector.example/v1".to_string(),
            api_key: "test-key".to_string(),
            request_timeout_ms: 1_000_000,
            request_deadline: None,
            external_api_audit_state: None,
            external_api_audit_user_id: None,
            external_api_audit_profile_id: None,
            external_api_audit_request_id: None,
        };

        assert_eq!(
            vector_engine_images_generation_url_for_test(&root_settings),
            "https://vector.example/v1/images/generations"
        );
        assert_eq!(
            vector_engine_images_generation_url_for_test(&v1_settings),
            "https://vector.example/v1/images/generations"
        );
    }

    #[test]
    fn vector_engine_edit_url_normalizes_base_url() {
        let root_settings = OpenAiImageSettings {
            base_url: "https://vector.example".to_string(),
            api_key: "test-key".to_string(),
            request_timeout_ms: 1_000_000,
            request_deadline: None,
            external_api_audit_state: None,
            external_api_audit_user_id: None,
            external_api_audit_profile_id: None,
            external_api_audit_request_id: None,
        };
        let v1_settings = OpenAiImageSettings {
            base_url: "https://vector.example/v1".to_string(),
            api_key: "test-key".to_string(),
            request_timeout_ms: 1_000_000,
            request_deadline: None,
            external_api_audit_state: None,
            external_api_audit_user_id: None,
            external_api_audit_profile_id: None,
            external_api_audit_request_id: None,
        };

        assert_eq!(
            vector_engine_images_edit_url_for_test(&root_settings),
            "https://vector.example/v1/images/edits"
        );
        assert_eq!(
            vector_engine_images_edit_url_for_test(&v1_settings),
            "https://vector.example/v1/images/edits"
        );
    }

    #[test]
    fn reference_data_url_stays_provider_owned() {
        let source = format!(
            "data:image/png;base64,{}",
            BASE64_STANDARD.encode(b"pngbytes")
        );
        let body = build_openai_image_request_body("提示词", None, "1:1", 1, &[source]);

        assert!(body.get("image").is_none());
    }

    #[test]
    fn vector_engine_upstream_failure_builds_tracking_ready_audit_event() {
        let audit = platform_image::PlatformImageFailureAudit {
            provider: VECTOR_ENGINE_PROVIDER,
            endpoint: "https://vector.example/v1/images/generations".to_string(),
            operation: "拼图 UI 背景图生成失败".to_string(),
            failure_stage: "upstream_status",
            status_code: Some(429),
            status_class: None,
            timeout: false,
            retryable: true,
            error_message: "上游限流".to_string(),
            error_source: None,
            raw_excerpt: Some("{\"error\":\"rate limited\"}".to_string()),
            latency_ms: Some(321),
            prompt_chars: Some(42),
            reference_image_count: Some(1),
            image_model: Some(VECTOR_ENGINE_GPT_IMAGE_2_MODEL),
        };
        let tracking = crate::external_api_audit::build_external_api_failure_tracking_draft(
            &build_external_api_failure_draft_from_platform_image_audit(&audit),
        );

        assert_eq!(
            tracking.event_key,
            crate::external_api_audit::EXTERNAL_API_FAILURE_EVENT_KEY
        );
        assert_eq!(tracking.scope_id, VECTOR_ENGINE_PROVIDER);
        assert_eq!(tracking.metadata["provider"], VECTOR_ENGINE_PROVIDER);
        assert_eq!(tracking.metadata["statusCode"], 429);
        assert_eq!(tracking.metadata["statusClass"], "4xx");
        assert_eq!(tracking.metadata["failureStage"], "upstream_status");
        assert_eq!(tracking.metadata["retryable"], true);
        assert_eq!(tracking.metadata["promptChars"], 42);
        assert_eq!(tracking.metadata["referenceImageCount"], 1);
        assert_eq!(
            tracking.metadata["imageModel"],
            VECTOR_ENGINE_GPT_IMAGE_2_MODEL
        );
    }
}

fn current_utc_micros() -> i64 {
    (OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000) as i64
}