Genarrative/server-rs/crates/platform-image/src/generated_asset_sheets/persist.rs

use std::{collections::BTreeMap, time::Duration};

use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
use platform_oss::{LegacyAssetPrefix, OssClient, OssObjectAccess, OssPutObjectRequest};

use super::error::GeneratedAssetSheetError;

const GENERATED_ASSET_SHEET_OSS_PUT_TIMEOUT_MS: u64 = 3 * 60_000;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GeneratedAssetSheetUpload {
    pub src: String,
    pub object_key: String,
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct GeneratedAssetSheetPersistPrompt {
    pub sheet_prompt: Option<String>,
    pub item_name_prompt: Option<String>,
    pub special_prompt: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct GeneratedAssetSheetPersistInput {
    pub prefix: LegacyAssetPrefix,
    pub owner_user_id: String,
    pub session_id: String,
    pub profile_id: String,
    pub path_segments: Vec<String>,
    pub file_name: String,
    pub content_type: String,
    pub bytes: Vec<u8>,
    pub asset_kind: String,
    pub source_job_id: Option<String>,
    pub generated_at_micros: i64,
    pub grid_size: usize,
    pub row_index: usize,
    pub view_index: usize,
    pub prompt: GeneratedAssetSheetPersistPrompt,
}

pub fn prepare_generated_asset_sheet_put_request(
    input: GeneratedAssetSheetPersistInput,
) -> Result<OssPutObjectRequest, GeneratedAssetSheetError> {
    if input.grid_size == 0 {
        return Err(GeneratedAssetSheetError::invalid_request(
            "系列素材图集的 n 必须大于 0。",
        ));
    }
    if input.row_index == 0
        || input.view_index == 0
        || input.row_index > input.grid_size
        || input.view_index > input.grid_size
    {
        return Err(GeneratedAssetSheetError::invalid_request(format!(
            "系列素材图集持久化的行列索引必须落在 n*n 范围内。gridSize={}, rowIndex={}, viewIndex={}",
            input.grid_size, input.row_index, input.view_index
        )));
    }

    let mut metadata = BTreeMap::new();
    metadata.insert(
        "x-oss-meta-asset-kind".to_string(),
        input.asset_kind.clone(),
    );
    metadata.insert(
        "x-oss-meta-owner-user-id".to_string(),
        input.owner_user_id.clone(),
    );
    metadata.insert(
        "x-oss-meta-profile-id".to_string(),
        input.profile_id.clone(),
    );
    metadata.insert(
        "x-oss-meta-generated-asset-sheet-grid-size".to_string(),
        input.grid_size.to_string(),
    );
    metadata.insert(
        "x-oss-meta-generated-asset-sheet-row-index".to_string(),
        input.row_index.to_string(),
    );
    metadata.insert(
        "x-oss-meta-generated-asset-sheet-view-index".to_string(),
        input.view_index.to_string(),
    );
    metadata.insert(
        "x-oss-meta-generated-at-micros".to_string(),
        input.generated_at_micros.to_string(),
    );
    if let Some(source_job_id) = input
        .source_job_id
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        metadata.insert(
            "x-oss-meta-source-job-id".to_string(),
            source_job_id.to_string(),
        );
    }
    insert_generated_asset_sheet_prompt_metadata(
        &mut metadata,
        "generated-asset-sheet-prompt-b64",
        input.prompt.sheet_prompt.as_deref(),
    );
    insert_generated_asset_sheet_prompt_metadata(
        &mut metadata,
        "generated-asset-sheet-item-name-prompt-b64",
        input.prompt.item_name_prompt.as_deref(),
    );
    insert_generated_asset_sheet_prompt_metadata(
        &mut metadata,
        "generated-asset-sheet-special-prompt-b64",
        input.prompt.special_prompt.as_deref(),
    );
    if input.prompt.sheet_prompt.is_some()
        || input.prompt.item_name_prompt.is_some()
        || input.prompt.special_prompt.is_some()
    {
        metadata.insert(
            "x-oss-meta-generated-asset-sheet-prompt-encoding".to_string(),
            "utf8-base64".to_string(),
        );
    }

    Ok(OssPutObjectRequest {
        prefix: input.prefix,
        path_segments: std::iter::once(input.session_id.as_str())
            .chain(std::iter::once(input.profile_id.as_str()))
            .chain(input.path_segments.iter().map(String::as_str))
            .map(|segment| sanitize_generated_asset_sheet_path_segment(segment, "asset"))
            .collect(),
        file_name: input.file_name,
        content_type: Some(input.content_type),
        access: OssObjectAccess::Private,
        metadata,
        body: input.bytes,
    })
}

pub async fn persist_generated_asset_sheet_bytes(
    oss_client: &OssClient,
    input: GeneratedAssetSheetPersistInput,
) -> Result<GeneratedAssetSheetUpload, GeneratedAssetSheetError> {
    let put_request = prepare_generated_asset_sheet_put_request(input)?;
    let oss_http_client = reqwest::Client::builder()
        .timeout(Duration::from_millis(
            GENERATED_ASSET_SHEET_OSS_PUT_TIMEOUT_MS,
        ))
        .build()
        .map_err(|error| {
            GeneratedAssetSheetError::build_http_client(format!(
                "构造系列素材图集 OSS 上传客户端失败：{error}"
            ))
        })?;
    let put_result = oss_client
        .put_object(&oss_http_client, put_request)
        .await
        .map_err(GeneratedAssetSheetError::Oss)?;

    Ok(GeneratedAssetSheetUpload {
        src: put_result.legacy_public_path,
        object_key: put_result.object_key,
    })
}

fn insert_generated_asset_sheet_prompt_metadata(
    metadata: &mut BTreeMap<String, String>,
    key: &str,
    value: Option<&str>,
) {
    let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
        return;
    };
    metadata.insert(
        format!("x-oss-meta-{key}"),
        BASE64_STANDARD.encode(value.as_bytes()),
    );
}

fn sanitize_generated_asset_sheet_path_segment(raw: &str, fallback: &str) -> String {
    let normalized = raw
        .trim()
        .chars()
        .map(|ch| {
            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
                ch.to_ascii_lowercase()
            } else {
                '-'
            }
        })
        .collect::<String>();
    let collapsed = normalized
        .split('-')
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>()
        .join("-");
    if collapsed.is_empty() {
        fallback.to_string()
    } else {
        collapsed.chars().take(64).collect()
    }
}