zerogravity/src/api/util.rs

//! Shared utilities for API handlers.

use axum::{
    http::StatusCode,
    response::{sse::Event, IntoResponse, Json},
};
use base64::Engine;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;

use super::types::{ErrorDetail, ErrorResponse};
use crate::proto::ImageData;

pub(crate) fn err_response(
    status: StatusCode,
    message: String,
    error_type: &str,
) -> axum::response::Response {
    let body = ErrorResponse {
        error: ErrorDetail {
            message,
            error_type: error_type.to_string(),
            code: Some(status.as_u16()),
            param: None,
        },
    };
    (status, Json(body)).into_response()
}

/// Convert a MITM-captured upstream error from Google into an HTTP response.
/// Forwards Google's exact error message and HTTP status code to the client.
pub(crate) fn upstream_err_response(
    err: &crate::mitm::store::UpstreamError,
) -> axum::response::Response {
    let status = StatusCode::from_u16(err.status).unwrap_or(StatusCode::BAD_GATEWAY);

    // Map Google error status to OpenAI-style error type
    let error_type = match err.error_status.as_deref() {
        Some("INVALID_ARGUMENT") => "invalid_request_error",
        Some("RESOURCE_EXHAUSTED") => "rate_limit_error",
        Some("PERMISSION_DENIED") | Some("UNAUTHENTICATED") => "authentication_error",
        Some("NOT_FOUND") => "not_found_error",
        Some("INTERNAL") | Some("UNAVAILABLE") => "server_error",
        _ => "upstream_error",
    };

    // Use Google's exact error message. Try parsed message first, then raw body.
    let message = if let Some(ref msg) = err.message {
        // Include Google's error status for context if available
        if let Some(ref gstatus) = err.error_status {
            format!("[{gstatus}] {msg}")
        } else {
            msg.clone()
        }
    } else if !err.body.is_empty() {
        // No parsed message — forward the raw body as-is so the client
        // sees exactly what Google returned (protobuf, HTML, etc.)
        err.body.clone()
    } else {
        format!("Google API error: HTTP {}", err.status)
    };

    // Extract param hint from Google's error details if available
    let param = serde_json::from_str::<serde_json::Value>(&err.body)
        .ok()
        .and_then(|v| {
            v["error"]["details"].as_array().and_then(|details| {
                details.iter().find_map(|d| {
                    d["fieldViolations"]
                        .as_array()
                        .and_then(|fv| fv.first())
                        .and_then(|v| v["field"].as_str().map(|s| s.to_string()))
                })
            })
        });

    let body = ErrorResponse {
        error: ErrorDetail {
            message,
            error_type: error_type.to_string(),
            code: Some(err.status),
            param,
        },
    };
    (status, Json(body)).into_response()
}

/// Extract the exact error message from a MITM-captured upstream error.
/// Preserves Google's original message verbatim. Used by streaming handlers.
pub(crate) fn upstream_error_message(err: &crate::mitm::store::UpstreamError) -> String {
    if let Some(ref msg) = err.message {
        if let Some(ref gstatus) = err.error_status {
            format!("[{gstatus}] {msg}")
        } else {
            msg.clone()
        }
    } else if !err.body.is_empty() {
        err.body.clone()
    } else {
        format!("Google API error: HTTP {}", err.status)
    }
}

/// Map Google's error status to OpenAI-compatible error type string.
pub(crate) fn upstream_error_type(err: &crate::mitm::store::UpstreamError) -> &'static str {
    match err.error_status.as_deref() {
        Some("INVALID_ARGUMENT") => "invalid_request_error",
        Some("RESOURCE_EXHAUSTED") => "rate_limit_error",
        Some("PERMISSION_DENIED") | Some("UNAUTHENTICATED") => "authentication_error",
        Some("NOT_FOUND") => "not_found_error",
        Some("INTERNAL") | Some("UNAVAILABLE") => "server_error",
        _ => "upstream_error",
    }
}

pub(crate) fn now_unix() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Default request timeout in seconds (used by serde defaults).
pub(crate) fn default_timeout() -> u64 {
    120
}

pub(crate) fn responses_sse_event(event_type: &str, data: serde_json::Value) -> Event {
    Event::default()
        .event(event_type)
        .data(serde_json::to_string(&data).unwrap_or_default())
}

// ─── Image extraction ────────────────────────────────────────────────────────

/// Parse a base64 data URI like `data:image/png;base64,iVBOR...` into ImageData.
/// Also accepts plain URLs (returns None — we only support inline base64).
pub(crate) fn parse_data_uri(url: &str) -> Option<ImageData> {
    // data:image/png;base64,<data>
    let rest = url.strip_prefix("data:")?;
    let (header, b64) = rest.split_once(";base64,")?;
    let mime_type = header.to_string();

    match base64::engine::general_purpose::STANDARD.decode(b64) {
        Ok(data) => {
            tracing::info!(mime = %mime_type, size = data.len(), "Decoded inline image");
            Some(ImageData { mime_type, data })
        }
        Err(e) => {
            warn!(error = %e, "Failed to decode base64 image data");
            None
        }
    }
}

/// Extract an image from an OpenAI content array item.
///
/// Supports:
/// - Chat Completions: `{"type": "image_url", "image_url": {"url": "data:..."}}`
/// - Responses API:    `{"type": "input_image", "image_url": "data:..."}` or
///   `{"type": "input_image", "url": "data:..."}`
pub(crate) fn extract_image_from_content(item: &serde_json::Value) -> Option<ImageData> {
    let item_type = item["type"].as_str().unwrap_or("");

    match item_type {
        // OpenAI Chat Completions format
        "image_url" => {
            let url = item["image_url"]["url"].as_str()?;
            parse_data_uri(url)
        }
        // OpenAI Responses API format
        "input_image" => {
            let url = item["image_url"]
                .as_str()
                .or_else(|| item["url"].as_str())?;
            parse_data_uri(url)
        }
        _ => None,
    }
}

/// Extract the first image from a content array (Value::Array of content parts).
pub(crate) fn extract_first_image(content: &serde_json::Value) -> Option<ImageData> {
    content
        .as_array()?
        .iter()
        .find_map(extract_image_from_content)
}