//! Reverse-proxy primitives: small hand-written hyper-client handler for
//! `/api/...` HTTP, and an axum WebSocket pump for `/api/attach/:id` +
//! `/api/events`.
//!
//! Keeps dependencies minimal — reqwest for HTTP, tokio_tungstenite for the
//! upstream WS leg. ~150 LOC total.

use std::sync::Arc;

use axum::body::Body;
use axum::extract::ws::{Message as AxMessage, WebSocket};
use axum::extract::{Request, State};
use axum::http::{HeaderMap, HeaderValue, StatusCode, Uri};
use axum::response::{IntoResponse, Response};
use futures_util::{SinkExt, StreamExt};
use maestro_protocol::routes::{self, UrlError};
use maestro_protocol::ApiError;
use tokio_tungstenite::tungstenite::Message as WsMessage;

#[derive(Debug, Clone)]
pub struct UpstreamConfig {
    /// HTTP base of the daemon — `http://127.0.0.1:7070`.
    pub daemon_http_base: Arc<String>,
    /// WS base of the daemon — derived from `daemon_http_base` by replacing
    /// the scheme. Cached so we don't redo it per request.
    pub daemon_ws_base: Arc<String>,
}

impl UpstreamConfig {
    pub fn from_daemon_url(daemon_url: &str) -> Result<Self, UrlError> {
        let http_base = routes::normalize_base(daemon_url).to_string();
        let ws_base = routes::ws_base(&http_base)?;
        Ok(Self {
            daemon_http_base: Arc::new(http_base),
            daemon_ws_base: Arc::new(ws_base),
        })
    }
}

/// Shared state plumbed into the proxy handlers.
#[derive(Debug, Clone)]
pub struct ProxyState {
    pub upstream: UpstreamConfig,
    pub client: reqwest::Client,
}

impl ProxyState {
    pub fn new(upstream: UpstreamConfig) -> Self {
        Self {
            upstream,
            client: reqwest::Client::builder()
                .redirect(reqwest::redirect::Policy::none())
                .build()
                .expect("reqwest client"),
        }
    }
}

// ---------------------------------------------------------------------------
// HTTP proxy
// ---------------------------------------------------------------------------

/// Cap on the request body we'll buffer before forwarding upstream. The proxied
/// daemon endpoints (`/api/sessions` CRUD, etc.) all carry small JSON; 8 MiB is
/// far above anything legitimate. Without a cap, anyone reachable on the bind
/// address could POST a multi-GB body and OOM the process (NAP-257).
pub(crate) const MAX_REQUEST_BODY: usize = 8 * 1024 * 1024;

/// Forward an HTTP request to the upstream daemon and stream the response back.
pub async fn proxy_http(State(state): State<ProxyState>, req: Request) -> Response {
    let (parts, body) = req.into_parts();
    let upstream_uri = build_uri(&state.upstream.daemon_http_base, &parts.uri);

    // Bound the buffered request body. `to_bytes` returns `Err` once the body
    // exceeds the limit; treat that as 413 rather than buffering unboundedly.
    let body_bytes = match axum::body::to_bytes(body, MAX_REQUEST_BODY).await {
        Ok(b) => b,
        Err(_) => return payload_too_large(),
    };

    let mut upstream_req = state
        .client
        .request(parts.method.clone(), upstream_uri.as_str())
        .body(body_bytes.to_vec());

    for (name, value) in filter_request_headers(&parts.headers).iter() {
        upstream_req = upstream_req.header(name, value);
    }

    let upstream_resp = match upstream_req.send().await {
        Ok(r) => r,
        Err(e) => return bad_gateway(&format!("upstream send: {e}")),
    };

    let status = upstream_resp.status();
    let headers = filter_response_headers(upstream_resp.headers());

    // Stream the response body straight through instead of materializing it —
    // a large upstream reply never gets fully buffered in RAM (NAP-257).
    let mut resp = Response::builder().status(status);
    for (name, value) in headers.iter() {
        resp = resp.header(name, value);
    }
    resp.body(Body::from_stream(upstream_resp.bytes_stream()))
        .unwrap_or_else(|e| bad_gateway(&format!("build response: {e}")))
}

fn build_uri(base: &str, original: &Uri) -> String {
    let mut out = base.to_string();
    out.push_str(original.path());
    if let Some(q) = original.query() {
        out.push('?');
        out.push_str(q);
    }
    out
}

/// Headers we strip before forwarding upstream / downstream. `host` must be
/// dropped (reqwest will set it correctly for the upstream); `content-length`
/// is regenerated by the client; hop-by-hop headers per RFC 7230 §6.1 should
/// not be relayed.
const HOP_BY_HOP: &[&str] = &[
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
];

fn filter_request_headers(headers: &HeaderMap) -> HeaderMap {
    let mut out = HeaderMap::new();
    for (name, value) in headers {
        let n = name.as_str();
        if n.eq_ignore_ascii_case("host") || n.eq_ignore_ascii_case("content-length") {
            continue;
        }
        if HOP_BY_HOP.iter().any(|h| n.eq_ignore_ascii_case(h)) {
            continue;
        }
        out.append(name.clone(), value.clone());
    }
    out
}

fn filter_response_headers(headers: &reqwest::header::HeaderMap) -> HeaderMap {
    let mut out = HeaderMap::new();
    for (name, value) in headers {
        let n = name.as_str();
        if HOP_BY_HOP.iter().any(|h| n.eq_ignore_ascii_case(h)) {
            continue;
        }
        if let (Ok(n), Ok(v)) = (
            axum::http::HeaderName::from_bytes(name.as_str().as_bytes()),
            HeaderValue::from_bytes(value.as_bytes()),
        ) {
            out.append(n, v);
        }
    }
    out
}

fn bad_gateway(msg: &str) -> Response {
    tracing::warn!(error = msg, "proxy: bad gateway");
    api_error(StatusCode::BAD_GATEWAY, "proxy_failed", msg)
}

fn payload_too_large() -> Response {
    tracing::warn!(limit = MAX_REQUEST_BODY, "proxy: request body exceeds cap");
    api_error(
        StatusCode::PAYLOAD_TOO_LARGE,
        "payload_too_large",
        &format!("request body exceeds {MAX_REQUEST_BODY} bytes"),
    )
}

/// Answer in the same envelope the daemon uses, so a client cannot tell whether
/// a failure came from the daemon or from this proxy.
fn api_error(status: StatusCode, code: &str, message: &str) -> Response {
    (status, axum::Json(ApiError::new(code, message))).into_response()
}

// ---------------------------------------------------------------------------
// WebSocket proxy
// ---------------------------------------------------------------------------

/// Open an upstream WS connection at `daemon_ws_base + path?query` and pump
/// frames in both directions until either side closes. Used for
/// `/api/attach/:id` and `/api/events`.
pub async fn pump_ws(mut client: WebSocket, upstream_url: String) {
    let (upstream, _) = match tokio_tungstenite::connect_async(&upstream_url).await {
        Ok(pair) => pair,
        Err(e) => {
            tracing::warn!(url = %upstream_url, error = %e, "proxy: upstream WS connect failed");
            // The daemon rejects an attach it can't serve (unknown session)
            // with a structured HTTP error *before* the upgrade. We've already
            // upgraded the browser socket, so a bare drop reaches it as an
            // opaque 1006 close — the attach pane just flashes and vanishes with
            // no reason. Relay the daemon's reason as an application close frame
            // the UI can show (#39).
            let (code, reason) = upstream_close_reason(&e);
            let _ = client
                .send(AxMessage::Close(Some(axum::extract::ws::CloseFrame {
                    code,
                    reason: reason.into(),
                })))
                .await;
            return;
        }
    };
    let (mut client_sink, mut client_stream) = client.split();
    let (mut up_sink, mut up_stream) = upstream.split();

    let client_to_up = async {
        while let Some(msg) = client_stream.next().await {
            let Ok(msg) = msg else { break };
            let frame = match msg {
                AxMessage::Text(t) => WsMessage::Text(t),
                AxMessage::Binary(b) => WsMessage::Binary(b),
                AxMessage::Ping(p) => WsMessage::Ping(p),
                AxMessage::Pong(p) => WsMessage::Pong(p),
                AxMessage::Close(c) => {
                    let _ = up_sink
                        .send(WsMessage::Close(c.map(|cf| {
                            tokio_tungstenite::tungstenite::protocol::CloseFrame {
                                code: cf.code.into(),
                                reason: cf.reason,
                            }
                        })))
                        .await;
                    break;
                }
            };
            if up_sink.send(frame).await.is_err() {
                break;
            }
        }
    };

    let up_to_client = async {
        while let Some(msg) = up_stream.next().await {
            let Ok(msg) = msg else { break };
            let frame = match msg {
                WsMessage::Text(t) => AxMessage::Text(t),
                WsMessage::Binary(b) => AxMessage::Binary(b),
                WsMessage::Ping(p) => AxMessage::Ping(p),
                WsMessage::Pong(p) => AxMessage::Pong(p),
                WsMessage::Close(c) => {
                    let _ = client_sink
                        .send(AxMessage::Close(c.map(|cf| {
                            axum::extract::ws::CloseFrame {
                                code: cf.code.into(),
                                reason: cf.reason.to_string().into(),
                            }
                        })))
                        .await;
                    break;
                }
                WsMessage::Frame(_) => continue,
            };
            if client_sink.send(frame).await.is_err() {
                break;
            }
        }
    };

    // Stop as soon as either side hangs up.
    tokio::select! {
        _ = client_to_up => {}
        _ = up_to_client => {}
    }
}

/// Largest application close code we emit; the browser surfaces it as
/// `CloseEvent.code`. Codes in 4000-4999 are reserved for application use.
const CLOSE_UPSTREAM_REJECTED: u16 = 4502;

/// Turn a failed upstream WS handshake into a browser-facing close code +
/// reason. When the daemon rejected the upgrade with an [`ApiError`] body,
/// surface that message; otherwise fall back to a generic gateway reason. The
/// reason rides a WS control frame, so it's truncated to the 123-byte limit.
fn upstream_close_reason(e: &tokio_tungstenite::tungstenite::Error) -> (u16, String) {
    use tokio_tungstenite::tungstenite::Error as WsError;
    let msg = match e {
        WsError::Http(resp) => resp
            .body()
            .as_ref()
            .and_then(|b| serde_json::from_slice::<ApiError>(b).ok())
            .map(|e| e.error.message)
            .unwrap_or_else(|| "the session's daemon rejected the attach".to_owned()),
        _ => "could not reach the session's daemon".to_owned(),
    };
    (CLOSE_UPSTREAM_REJECTED, truncate_reason(&msg))
}

/// Clamp a close reason to the WS control-frame budget (125 bytes total, 2 for
/// the code → 123 for the reason), respecting UTF-8 char boundaries.
fn truncate_reason(s: &str) -> String {
    const MAX: usize = 123;
    if s.len() <= MAX {
        return s.to_owned();
    }
    let mut end = MAX;
    while end > 0 && !s.is_char_boundary(end) {
        end -= 1;
    }
    s[..end].to_owned()
}

/// Compose the upstream WS URL from the base + path + query of the inbound
/// request.
pub fn upstream_ws_url(state: &ProxyState, uri: &Uri) -> String {
    let mut out = state.upstream.daemon_ws_base.as_str().to_string();
    out.push_str(uri.path());
    if let Some(q) = uri.query() {
        out.push('?');
        out.push_str(q);
    }
    out
}

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

    #[test]
    fn ws_base_swaps_http_scheme() {
        let u = UpstreamConfig::from_daemon_url("http://127.0.0.1:7070/").unwrap();
        assert_eq!(&*u.daemon_http_base, "http://127.0.0.1:7070");
        assert_eq!(&*u.daemon_ws_base, "ws://127.0.0.1:7070");
    }

    #[test]
    fn ws_base_swaps_https_scheme() {
        let u = UpstreamConfig::from_daemon_url("https://daemon.example/").unwrap();
        assert_eq!(&*u.daemon_ws_base, "wss://daemon.example");
    }

    /// NAP-257: a request body over the cap is rejected with 413 before we ever
    /// dial upstream — so an oversized POST can't OOM the proxy. The upstream is
    /// a bogus unreachable address; a 413 (not a 502) proves we bailed early.
    #[tokio::test]
    async fn oversized_request_body_is_rejected_with_413() {
        let state = ProxyState::new(UpstreamConfig::from_daemon_url("http://127.0.0.1:9").unwrap());
        let big = vec![0u8; MAX_REQUEST_BODY + 1];
        let req = Request::builder()
            .method("POST")
            .uri("/api/sessions")
            .body(Body::from(big))
            .unwrap();

        let resp = proxy_http(State(state), req).await;
        assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE);
    }

    /// A body at the cap still proceeds to the upstream send (which fails here
    /// with 502 since nothing is listening) — i.e. the cap is exclusive of the
    /// limit, not off-by-one against legitimate traffic.
    #[tokio::test]
    async fn body_at_cap_is_not_rejected_by_the_size_check() {
        let state = ProxyState::new(UpstreamConfig::from_daemon_url("http://127.0.0.1:9").unwrap());
        let body = vec![0u8; MAX_REQUEST_BODY];
        let req = Request::builder()
            .method("POST")
            .uri("/api/sessions")
            .body(Body::from(body))
            .unwrap();

        let resp = proxy_http(State(state), req).await;
        assert_eq!(resp.status(), StatusCode::BAD_GATEWAY);
    }

    /// #39: a daemon rejection carrying the structured error body surfaces its
    /// `message` as the close reason, under an application close code, so the
    /// browser can show *why* the attach failed instead of an opaque 1006.
    #[test]
    fn upstream_close_reason_extracts_daemon_message() {
        use tokio_tungstenite::tungstenite::{http, Error as WsError};
        let body =
            br#"{"error":{"code":"unknown_session","message":"no such session on this daemon"}}"#
                .to_vec();
        let resp = http::Response::builder()
            .status(404)
            .body(Some(body))
            .unwrap();
        let (code, reason) = upstream_close_reason(&WsError::Http(resp));
        assert_eq!(code, CLOSE_UPSTREAM_REJECTED);
        assert_eq!(reason, "no such session on this daemon");
    }

    /// A non-HTTP failure (e.g. TCP refused — peer down) still yields a usable
    /// generic reason rather than a panic or empty string.
    #[test]
    fn upstream_close_reason_falls_back_for_non_http_errors() {
        use tokio_tungstenite::tungstenite::Error as WsError;
        let e = WsError::Io(std::io::Error::new(
            std::io::ErrorKind::ConnectionRefused,
            "nope",
        ));
        let (code, reason) = upstream_close_reason(&e);
        assert_eq!(code, CLOSE_UPSTREAM_REJECTED);
        assert!(!reason.is_empty());
    }

    /// Close reasons ride a control frame capped at 123 bytes; an over-long
    /// message is truncated on a char boundary (no panic on multi-byte chars).
    #[test]
    fn truncate_reason_respects_control_frame_budget() {
        let ascii = "x".repeat(200);
        assert_eq!(truncate_reason(&ascii).len(), 123);

        let multibyte = "é".repeat(100); // 2 bytes each
        let out = truncate_reason(&multibyte);
        assert!(out.len() <= 123);
        assert!(out.is_char_boundary(out.len()));
    }
}