ystp/src/services/idempotency.rs

use crate::error::{AppError, ErrorCode};
use crate::state::AppState;

use chrono::{DateTime, Duration, Utc};
use serde_json::Value as JsonValue;
use sha2::{Digest, Sha256};
use sqlx::FromRow;
use uuid::Uuid;

#[derive(Debug, Clone, Copy)]
pub enum Scope {
    User(Uuid),
    ApiKey(Uuid),
}

#[derive(Debug)]
pub enum BeginResult {
    Acquired,
    Replay { response_body: JsonValue },
    InProgress,
}

#[derive(Debug, FromRow)]
struct IdemRow {
    request_hash: String,
    response_status: i32,
    response_body: Option<JsonValue>,
}

pub fn sha256_hex(parts: &[&[u8]]) -> String {
    let mut hasher = Sha256::new();
    for p in parts {
        hasher.update(p);
        hasher.update([0u8]); // separator
    }
    hex::encode(hasher.finalize())
}

pub async fn begin(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    request_hash: &str,
    ttl_hours: i64,
) -> Result<BeginResult, AppError> {
    if idempotency_key.trim().is_empty() {
        return Err(AppError::new(
            ErrorCode::InvalidRequest,
            "Idempotency-Key 不能为空",
        ));
    }
    if idempotency_key.len() > 128 {
        return Err(AppError::new(
            ErrorCode::InvalidRequest,
            "Idempotency-Key 过长",
        ));
    }
    if request_hash.len() != 64 {
        return Err(AppError::new(
            ErrorCode::InvalidRequest,
            "request_hash 不合法",
        ));
    }

    let now = Utc::now();
    let expires_at = now + Duration::hours(ttl_hours.max(1));

    cleanup_expired_for_key(state, scope, idempotency_key, now).await?;

    let inserted = match scope {
        Scope::User(user_id) => {
            sqlx::query(
                r#"
                INSERT INTO idempotency_keys (
                  user_id, idempotency_key, request_hash,
                  response_status, response_body,
                  expires_at
                ) VALUES (
                  $1, $2, $3,
                  0, NULL,
                  $4
                )
                ON CONFLICT DO NOTHING
                "#,
            )
            .bind(user_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .bind(expires_at)
            .execute(&state.db)
            .await
        }
        Scope::ApiKey(api_key_id) => {
            sqlx::query(
                r#"
                INSERT INTO idempotency_keys (
                  api_key_id, idempotency_key, request_hash,
                  response_status, response_body,
                  expires_at
                ) VALUES (
                  $1, $2, $3,
                  0, NULL,
                  $4
                )
                ON CONFLICT DO NOTHING
                "#,
            )
            .bind(api_key_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .bind(expires_at)
            .execute(&state.db)
            .await
        }
    }
    .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等记录失败").with_source(err))?;

    if inserted.rows_affected() > 0 {
        return Ok(BeginResult::Acquired);
    }

    let row = get_row(state, scope, idempotency_key, now).await?;
    let Some(row) = row else {
        return Ok(BeginResult::Acquired);
    };

    if row.request_hash != request_hash {
        return Err(AppError::new(
            ErrorCode::IdempotencyConflict,
            "同一个 Idempotency-Key 的请求参数不一致",
        ));
    }

    if row.response_status == 0 || row.response_body.is_none() {
        return Ok(BeginResult::InProgress);
    }

    Ok(BeginResult::Replay {
        response_body: row.response_body.unwrap_or(JsonValue::Null),
    })
}

pub async fn wait_for_replay(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    request_hash: &str,
    max_wait_ms: u64,
) -> Result<Option<(i32, JsonValue)>, AppError> {
    let started = tokio::time::Instant::now();

    loop {
        let row = get_row(state, scope, idempotency_key, Utc::now()).await?;
        let Some(row) = row else { return Ok(None) };

        if row.request_hash != request_hash {
            return Err(AppError::new(
                ErrorCode::IdempotencyConflict,
                "同一个 Idempotency-Key 的请求参数不一致",
            ));
        }

        if row.response_status != 0 {
            return Ok(Some((
                row.response_status,
                row.response_body.unwrap_or(JsonValue::Null),
            )));
        }

        if started.elapsed().as_millis() as u64 >= max_wait_ms {
            return Ok(None);
        }

        tokio::time::sleep(std::time::Duration::from_millis(200)).await;
    }
}

pub async fn complete(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    request_hash: &str,
    response_status: i32,
    response_body: JsonValue,
) -> Result<(), AppError> {
    let updated = match scope {
        Scope::User(user_id) => {
            sqlx::query(
                r#"
                UPDATE idempotency_keys
                SET response_status = $4,
                    response_body = $5
                WHERE user_id = $1
                  AND idempotency_key = $2
                  AND request_hash = $3
                  AND response_status = 0
                "#,
            )
            .bind(user_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .bind(response_status)
            .bind(response_body)
            .execute(&state.db)
            .await
        }
        Scope::ApiKey(api_key_id) => {
            sqlx::query(
                r#"
                UPDATE idempotency_keys
                SET response_status = $4,
                    response_body = $5
                WHERE api_key_id = $1
                  AND idempotency_key = $2
                  AND request_hash = $3
                  AND response_status = 0
                "#,
            )
            .bind(api_key_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .bind(response_status)
            .bind(response_body)
            .execute(&state.db)
            .await
        }
    }
    .map_err(|err| AppError::new(ErrorCode::Internal, "写入幂等结果失败").with_source(err))?;

    if updated.rows_affected() == 0 {
        tracing::warn!("idempotency record not updated (maybe already completed?)");
    }

    Ok(())
}

pub async fn abort(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    request_hash: &str,
) -> Result<(), AppError> {
    match scope {
        Scope::User(user_id) => {
            let _ = sqlx::query(
                "DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
            )
            .bind(user_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .execute(&state.db)
            .await;
        }
        Scope::ApiKey(api_key_id) => {
            let _ = sqlx::query(
                "DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND request_hash = $3 AND response_status = 0",
            )
            .bind(api_key_id)
            .bind(idempotency_key)
            .bind(request_hash)
            .execute(&state.db)
            .await;
        }
    }
    Ok(())
}

async fn cleanup_expired_for_key(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    now: DateTime<Utc>,
) -> Result<(), AppError> {
    match scope {
        Scope::User(user_id) => {
            let _ = sqlx::query(
                "DELETE FROM idempotency_keys WHERE user_id = $1 AND idempotency_key = $2 AND expires_at < $3",
            )
            .bind(user_id)
            .bind(idempotency_key)
            .bind(now)
            .execute(&state.db)
            .await;
        }
        Scope::ApiKey(api_key_id) => {
            let _ = sqlx::query(
                "DELETE FROM idempotency_keys WHERE api_key_id = $1 AND idempotency_key = $2 AND expires_at < $3",
            )
            .bind(api_key_id)
            .bind(idempotency_key)
            .bind(now)
            .execute(&state.db)
            .await;
        }
    }
    Ok(())
}

async fn get_row(
    state: &AppState,
    scope: Scope,
    idempotency_key: &str,
    now: DateTime<Utc>,
) -> Result<Option<IdemRow>, AppError> {
    let row = match scope {
        Scope::User(user_id) => {
            sqlx::query_as::<_, IdemRow>(
                r#"
                SELECT request_hash, response_status, response_body
                FROM idempotency_keys
                WHERE user_id = $1
                  AND idempotency_key = $2
                  AND expires_at > $3
                ORDER BY created_at DESC
                LIMIT 1
                "#,
            )
            .bind(user_id)
            .bind(idempotency_key)
            .bind(now)
            .fetch_optional(&state.db)
            .await
        }
        Scope::ApiKey(api_key_id) => {
            sqlx::query_as::<_, IdemRow>(
                r#"
                SELECT request_hash, response_status, response_body
                FROM idempotency_keys
                WHERE api_key_id = $1
                  AND idempotency_key = $2
                  AND expires_at > $3
                ORDER BY created_at DESC
                LIMIT 1
                "#,
            )
            .bind(api_key_id)
            .bind(idempotency_key)
            .bind(now)
            .fetch_optional(&state.db)
            .await
        }
    }
    .map_err(|err| AppError::new(ErrorCode::Internal, "查询幂等记录失败").with_source(err))?;

    Ok(row)
}