yomo 2.1.5

use std::sync::Arc;

use anyhow::anyhow;
use axum::{
    Router,
    body::Bytes,
    extract::{Path, State},
    http::{self, HeaderMap, HeaderName, HeaderValue, StatusCode},
    response::{IntoResponse, Sse, sse::Event},
};
use futures_util::stream;
use log::{debug, error};
use tokio::io::{AsyncWriteExt, ReadHalf, SimplexStream};

use crate::{
    connector::{Connector, MemoryConnector},
    io::{receive_bytes, receive_frame, send_bytes, send_frame},
    types::{BodyFormat, RequestHeaders, ResponseHeaders},
};

/// Custom error with HTTP status code
pub struct CustomError {
    status_code: StatusCode,
    msg: String,
}

impl IntoResponse for CustomError {
    fn into_response(self) -> axum::response::Response {
        (self.status_code, self.msg).into_response()
    }
}

impl<E> From<E> for CustomError
where
    E: Into<anyhow::Error>,
{
    fn from(err: E) -> Self {
        Self {
            status_code: StatusCode::INTERNAL_SERVER_ERROR,
            msg: err.into().to_string(),
        }
    }
}

/// Parse HTTP header value
fn parse_http_headers(http_headers: &HeaderMap, key: &str) -> String {
    match http_headers.get(key) {
        Some(value) => value.to_str().unwrap_or_default(),
        None => "",
    }
    .to_string()
}

/// Return application-specific headers while excluding standard HTTP headers.
pub(crate) fn custom_http_headers(http_headers: &HeaderMap) -> HeaderMap {
    const STANDARD_HEADERS: &[&str] = &[
        "accept",
        "accept-charset",
        "accept-encoding",
        "accept-language",
        "authorization",
        "cache-control",
        "connection",
        "content-encoding",
        "content-length",
        "content-type",
        "cookie",
        "date",
        "expect",
        "forwarded",
        "host",
        "if-match",
        "if-modified-since",
        "if-none-match",
        "if-range",
        "if-unmodified-since",
        "origin",
        "pragma",
        "referer",
        "te",
        "trailer",
        "transfer-encoding",
        "user-agent",
        "upgrade",
        "via",
    ];

    http_headers
        .iter()
        .filter(|(name, _)| {
            !STANDARD_HEADERS
                .iter()
                .any(|standard| name == &HeaderName::from_static(standard))
        })
        .map(|(name, value)| (name.clone(), value.clone()))
        .collect()
}

/// Return all headers with credential values redacted.
pub(crate) fn redacted_http_headers(http_headers: &HeaderMap) -> HeaderMap {
    const SENSITIVE_HEADERS: &[&str] = &[
        "authorization",
        "proxy-authorization",
        "cookie",
        "x-api-key",
        "api-key",
        "x-goog-api-key",
    ];

    http_headers
        .iter()
        .map(|(name, value)| {
            let value = if SENSITIVE_HEADERS.contains(&name.as_str()) {
                HeaderValue::from_static("[redacted]")
            } else {
                value.clone()
            };
            (name.clone(), value)
        })
        .collect()
}

/// Create request headers from HTTP headers
fn new_request_headers(tool_name: &str, http_headers: &HeaderMap) -> RequestHeaders {
    RequestHeaders {
        name: tool_name.to_owned(),
        body_format: BodyFormat::Bytes,
        trace_id: parse_http_headers(http_headers, "X-Trace-Id"),
        span_id: parse_http_headers(http_headers, "X-Span-Id"),
        extension: parse_http_headers(http_headers, "X-Extension"),
    }
}

/// Custom response supporting both regular bytes body and SSE streaming
pub struct CustomResponse {
    body: Option<Vec<u8>>,
    reader: Option<ReadHalf<SimplexStream>>,
}

impl IntoResponse for CustomResponse {
    fn into_response(self) -> axum::response::Response {
        if let Some(body) = self.body {
            debug!("recv body: {}", String::from_utf8_lossy(&body));
            (StatusCode::OK, body).into_response()
        } else if let Some(reader) = self.reader {
            let stream = stream::unfold(reader, move |mut r| async move {
                match receive_bytes(&mut r).await {
                    Ok(Some(chunk)) => {
                        let data = String::from_utf8_lossy(&chunk);
                        debug!("recv chunk: {}", data);
                        Some((Ok(Event::default().data(data)), r))
                    }
                    Ok(None) => {
                        debug!("recv chunk done");
                        None
                    }
                    Err(e) => {
                        error!("receiving frame error: {}", e);
                        Some((Err(anyhow!("receiving frame error: {}", e)), r))
                    }
                }
            });
            Sse::new(stream).into_response()
        } else {
            (StatusCode::OK, "".to_string()).into_response()
        }
    }
}

/// HTTP stream handler: forward request to corresponding QUIC tool with SSE response
#[axum::debug_handler]
pub async fn tool_invoke_handler(
    http_headers: HeaderMap,
    Path(tool_name): Path<String>,
    State(connector): State<Arc<MemoryConnector>>,
    body: Bytes,
) -> Result<CustomResponse, CustomError> {
    let request_headers = new_request_headers(&tool_name, &http_headers);

    debug!(
        "request custom HTTP headers: {:?}",
        custom_http_headers(&http_headers)
    );
    debug!(
        "[{}|{}] new request to [{}]: {}",
        request_headers.trace_id,
        request_headers.span_id,
        request_headers.name,
        String::from_utf8_lossy(&body)
    );

    let (mut reader, mut writer) = connector.open_new_stream().await?;

    // send request headers
    send_frame(&mut writer, &request_headers).await?;

    // send request body
    send_bytes(&mut writer, &body.to_vec()).await?;
    writer.shutdown().await?;

    let response_headers: ResponseHeaders = receive_frame(&mut reader)
        .await?
        .ok_or(anyhow!("Failed to receive response headers"))?;

    if response_headers.status_code != http::StatusCode::OK {
        return Err(CustomError {
            status_code: StatusCode::from_u16(response_headers.status_code)?,
            msg: response_headers.error_msg,
        });
    }

    match response_headers.body_format {
        BodyFormat::Null => Ok(CustomResponse {
            body: None,
            reader: None,
        }),
        BodyFormat::Bytes => {
            let body = receive_bytes(&mut reader)
                .await?
                .ok_or(anyhow!("Failed to receive response"))?;

            Ok(CustomResponse {
                body: Some(body),
                reader: None,
            })
        }
        BodyFormat::Chunk => {
            // Stream response using SSE
            Ok(CustomResponse {
                body: None,
                reader: Some(reader),
            })
        }
    }
}

/// Tool API server: listen and receive external requests for tool invocation
pub async fn build_tool_api(connector: MemoryConnector) -> anyhow::Result<Router> {
    let state = Arc::new(connector);

    let app = axum::Router::new()
        .route("/{tool_name}", axum::routing::post(tool_invoke_handler))
        .with_state(state);

    Ok(app)
}