//! `/v1/chat/completions` 透传 — 替换 Authorization 头，把请求 body 原样 forward 到
//! 上游 LLM gateway，把响应 body 原样回吐给客户端。
//!
//! 一期只支持非 streaming（force `stream: false` 进 body），SSE 透传留给二期。

use std::sync::Arc;

use axum::{
    body::Bytes,
    extract::State,
    http::{HeaderMap, HeaderValue, StatusCode},
    response::{IntoResponse, Response},
};
use serde_json::Value;

#[derive(Clone, Debug)]
pub struct Config {
    pub upstream_url: String, // http://3.135.65.204:8848/v1/chat/completions
    pub upstream_token: String,
    pub proxy_auth_token: String,
}

impl Config {
    pub fn from_env() -> Self {
        let gateway = std::env::var("LLM_GATEWAY")
            .unwrap_or_else(|_| "http://3.135.65.204:8848/v1".to_string());
        let upstream_url = format!("{}/chat/completions", gateway.trim_end_matches('/'));
        Self {
            upstream_url,
            upstream_token: std::env::var("BACKEND_TOKEN").unwrap_or_default(),
            proxy_auth_token: std::env::var("PROXY_AUTH_TOKEN").unwrap_or_default(),
        }
    }
}

pub async fn handle(State(cfg): State<Arc<Config>>, body: Bytes) -> Response {
    // 1. parse body → 强制 stream=false（一期不支持流式）
    let body_bytes = match force_non_stream(&body) {
        Ok(b) => b,
        Err(e) => {
            return (StatusCode::BAD_REQUEST, format!("bad JSON body: {e}")).into_response();
        }
    };

    // 2. forward
    let client = reqwest::Client::new();
    let res = client
        .post(&cfg.upstream_url)
        .header("Authorization", format!("Bearer {}", cfg.upstream_token))
        .header("Content-Type", "application/json")
        .body(body_bytes)
        .send()
        .await;

    match res {
        Ok(r) => relay_response(r).await,
        Err(e) => {
            tracing::error!(error=%e, "upstream call failed");
            (StatusCode::BAD_GATEWAY, format!("upstream error: {e}")).into_response()
        }
    }
}

/// parse JSON、塞入 `stream: false`、重新 serialize。
/// 如果不是 JSON object 就保持原样（让上游自己报错）。
fn force_non_stream(body: &Bytes) -> Result<Vec<u8>, String> {
    if body.is_empty() {
        return Err("empty body".into());
    }
    let mut v: Value = serde_json::from_slice(body).map_err(|e| e.to_string())?;
    if let Some(obj) = v.as_object_mut() {
        obj.insert("stream".to_string(), Value::Bool(false));
    }
    serde_json::to_vec(&v).map_err(|e| e.to_string())
}

async fn relay_response(upstream: reqwest::Response) -> Response {
    let status = upstream.status();
    let ct = upstream
        .headers()
        .get(reqwest::header::CONTENT_TYPE)
        .cloned()
        .unwrap_or_else(|| HeaderValue::from_static("application/json"));
    let bytes = match upstream.bytes().await {
        Ok(b) => b,
        Err(e) => {
            tracing::error!(error=%e, "read upstream body");
            return (StatusCode::BAD_GATEWAY, "read upstream body failed").into_response();
        }
    };
    let mut headers = HeaderMap::new();
    headers.insert(axum::http::header::CONTENT_TYPE, ct);
    (
        StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY),
        headers,
        bytes,
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn force_non_stream_overrides_stream_true() {
        let input = Bytes::from(r#"{"model":"gemma","messages":[],"stream":true}"#);
        let out = force_non_stream(&input).unwrap();
        let v: Value = serde_json::from_slice(&out).unwrap();
        assert_eq!(v["stream"], Value::Bool(false));
    }

    #[test]
    fn force_non_stream_injects_when_absent() {
        let input = Bytes::from(r#"{"model":"gemma","messages":[]}"#);
        let out = force_non_stream(&input).unwrap();
        let v: Value = serde_json::from_slice(&out).unwrap();
        assert_eq!(v["stream"], Value::Bool(false));
    }

    #[test]
    fn force_non_stream_preserves_other_fields() {
        let input = Bytes::from(
            r#"{"model":"gemma-4-31b-it","temperature":0.7,"messages":[{"role":"user","content":"hi"}]}"#,
        );
        let out = force_non_stream(&input).unwrap();
        let v: Value = serde_json::from_slice(&out).unwrap();
        assert_eq!(v["model"], "gemma-4-31b-it");
        assert_eq!(v["temperature"], 0.7);
        assert_eq!(v["messages"][0]["role"], "user");
    }

    #[test]
    fn force_non_stream_rejects_empty() {
        assert!(force_non_stream(&Bytes::new()).is_err());
    }

    #[test]
    fn force_non_stream_rejects_invalid_json() {
        let input = Bytes::from(r#"not json"#);
        assert!(force_non_stream(&input).is_err());
    }

    #[test]
    fn config_from_env_builds_completions_url() {
        // Saved env keeps test isolation under cargo test (run in parallel)
        let prev_gateway = std::env::var("LLM_GATEWAY").ok();
        let prev_token = std::env::var("BACKEND_TOKEN").ok();
        let prev_proxy = std::env::var("PROXY_AUTH_TOKEN").ok();
        std::env::set_var("LLM_GATEWAY", "http://1.2.3.4:8848/v1/");
        std::env::set_var("BACKEND_TOKEN", "backend-xxx");
        std::env::set_var("PROXY_AUTH_TOKEN", "client-yyy");

        let cfg = Config::from_env();
        assert_eq!(cfg.upstream_url, "http://1.2.3.4:8848/v1/chat/completions");
        assert_eq!(cfg.upstream_token, "backend-xxx");
        assert_eq!(cfg.proxy_auth_token, "client-yyy");

        // restore
        match prev_gateway {
            Some(v) => std::env::set_var("LLM_GATEWAY", v),
            None => std::env::remove_var("LLM_GATEWAY"),
        }
        match prev_token {
            Some(v) => std::env::set_var("BACKEND_TOKEN", v),
            None => std::env::remove_var("BACKEND_TOKEN"),
        }
        match prev_proxy {
            Some(v) => std::env::set_var("PROXY_AUTH_TOKEN", v),
            None => std::env::remove_var("PROXY_AUTH_TOKEN"),
        }
    }
}
