unigateway 1.3.0

use std::fmt::Display;
use std::sync::Arc;
use std::time::Instant;

use axum::{
    http::{HeaderMap, StatusCode},
    response::Response,
};
use serde_json::Value;
use tracing::info;
use unigateway_host::host::HostContext;

use crate::middleware::{GatewayAuth, error_json, extract_openai_api_key, extract_x_api_key};
use crate::routing::target_provider_hint;
use crate::types::GatewayRequestState;

pub(super) struct PreparedGatewayRequest<'a> {
    pub start: Instant,
    pub host: HostContext<'a>,
    pub token: String,
    pub hint: Option<String>,
    pub auth: Option<GatewayAuth>,
}

pub(super) async fn prepare_and_parse_openai_request<'a, Request, Parse, ParseError>(
    state: &'a Arc<GatewayRequestState>,
    headers: &HeaderMap,
    payload: &Value,
    parse_request: Parse,
) -> Result<(PreparedGatewayRequest<'a>, Request), Response>
where
    Parse: FnOnce(&Value, &str) -> Result<Request, ParseError>,
    ParseError: Display,
{
    let prepared = prepare_openai_request(state, headers, payload).await?;
    parse_prepared_request(prepared, payload, parse_request, |_| {
        state.provider_model(unigateway_host::host::HostEnvProvider::OpenAi)
    })
    .await
}

pub(super) async fn prepare_and_parse_anthropic_request<'a, Request, Parse, ParseError>(
    state: &'a Arc<GatewayRequestState>,
    headers: &HeaderMap,
    payload: &Value,
    parse_request: Parse,
) -> Result<(PreparedGatewayRequest<'a>, Request), Response>
where
    Parse: FnOnce(&Value, &str) -> Result<Request, ParseError>,
    ParseError: Display,
{
    let prepared = prepare_anthropic_request(state, headers, payload).await?;
    parse_prepared_request(prepared, payload, parse_request, |_| {
        state.provider_model(unigateway_host::host::HostEnvProvider::Anthropic)
    })
    .await
}

async fn prepare_openai_request<'a>(
    state: &'a Arc<GatewayRequestState>,
    headers: &HeaderMap,
    payload: &Value,
) -> Result<PreparedGatewayRequest<'a>, Response> {
    prepare_gateway_request(state, headers, payload, || {
        extract_openai_api_key(headers, "")
    })
    .await
}

async fn prepare_anthropic_request<'a>(
    state: &'a Arc<GatewayRequestState>,
    headers: &HeaderMap,
    payload: &Value,
) -> Result<PreparedGatewayRequest<'a>, Response> {
    let prepared =
        prepare_gateway_request(state, headers, payload, || extract_x_api_key(headers, "")).await?;

    log_prepared_anthropic_request(headers, payload, &prepared);

    Ok(prepared)
}

async fn prepare_gateway_request<'a, ExtractToken>(
    state: &'a Arc<GatewayRequestState>,
    headers: &HeaderMap,
    payload: &Value,
    extract_token: ExtractToken,
) -> Result<PreparedGatewayRequest<'a>, Response>
where
    ExtractToken: FnOnce() -> String,
{
    let start = Instant::now();
    let token = extract_token();
    let host = HostContext::from_parts(state.as_ref(), state.as_ref());
    let hint = target_provider_hint(headers, payload);
    let auth = GatewayAuth::try_authenticate(state, &token).await?;

    Ok(PreparedGatewayRequest {
        start,
        host,
        token,
        hint,
        auth,
    })
}

async fn parse_prepared_request<'a, Request, Parse, ParseError, Model>(
    prepared: PreparedGatewayRequest<'a>,
    payload: &Value,
    parse_request: Parse,
    model: Model,
) -> Result<(PreparedGatewayRequest<'a>, Request), Response>
where
    Parse: FnOnce(&Value, &str) -> Result<Request, ParseError>,
    ParseError: Display,
    Model: FnOnce(&PreparedGatewayRequest<'a>) -> &'a str,
{
    let request = parse_request(payload, model(&prepared)).map_err(invalid_request_response)?;

    Ok((prepared, request))
}

fn log_prepared_anthropic_request(
    headers: &HeaderMap,
    payload: &Value,
    prepared: &PreparedGatewayRequest<'_>,
) {
    let has_x_api_key = headers.get("x-api-key").is_some();
    let has_bearer = headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .map(|value| value.starts_with("Bearer "))
        .unwrap_or(false);

    info!(
        endpoint = "/v1/messages",
        has_x_api_key,
        has_bearer,
        token_present = !prepared.token.is_empty(),
        model = payload
            .get("model")
            .and_then(|value| value.as_str())
            .unwrap_or(""),
        "received anthropic request"
    );
}

fn invalid_request_response(error: impl Display) -> Response {
    error_json(
        StatusCode::BAD_REQUEST,
        &format!("invalid request: {error}"),
    )
}