//! POST pipeline for the streamable HTTP transport.
//!
//! One envelope per POST: send with signing, retry transient failures
//! with backoff, capture the response headers for the budget gate
//! (WOR-1287), and dispatch the JSON or SSE-framed body into the
//! inbound channel. Split out of `mod.rs` so the transport module
//! stays within the module-size guard; the connection lifecycle, the
//! SSE GET reader, and request signing remain in the parent.

use std::time::Duration;

use bytes::Bytes;
use futures_util::StreamExt;
use reqwest::header::{ACCEPT, CONTENT_TYPE};
use reqwest::{Client as HttpClient, StatusCode};
use serde_json::Value;
use tokio::sync::mpsc;

use reqwest::Method;
use tokio::time::sleep;

use super::{
    apply_signature, auth_hint, SharedState, TransportMode, BACKOFF_STEPS_MS,
    HTTP_TRANSPORT_ERROR_CODE, METHOD_HEADER, NAME_HEADER, SESSION_HEADER,
};
use crate::protocol::{JsonRpcErrorBody, JsonRpcMessage};
use crate::transport::sse_parser::SseParser;

/// POST one envelope with exponential backoff on transient failures.
pub(super) async fn post_with_retry(
    http: &HttpClient,
    url: &str,
    shared: &SharedState,
    signer: Option<&crate::auth::RequestSigner>,
    msg: &JsonRpcMessage,
    inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>,
) -> Result<(), ()> {
    let mut attempt = 0usize;
    loop {
        match do_post(http, url, shared, signer, msg, inbound_tx).await {
            Ok(()) => return Ok(()),
            Err(err) => {
                if attempt >= BACKOFF_STEPS_MS.len() || !should_retry(&err) {
                    return Err(());
                }
                sleep(retry_wait(attempt, err.retry_after)).await;
                attempt += 1;
            }
        }
    }
}

/// Compute the wait before the next retry. Honors `Retry-After` if the
/// server supplied one; otherwise uses the next step in [`BACKOFF_STEPS_MS`].
pub(super) fn retry_wait(attempt: usize, retry_after: Option<u64>) -> Duration {
    if let Some(secs) = retry_after {
        return Duration::from_secs(secs);
    }
    let ms = BACKOFF_STEPS_MS
        .get(attempt)
        .copied()
        .unwrap_or(BACKOFF_STEPS_MS[BACKOFF_STEPS_MS.len() - 1]);
    Duration::from_millis(ms)
}

/// Local POST-error type so retry logic can read the transience without
/// wrapping every error path in a custom enum. The writer discards the
/// concrete cause (it only decides retry-or-poison), so we keep just the
/// retry signal and the optional `Retry-After` window.
pub(super) struct PostError {
    pub(super) transient: bool,
    pub(super) retry_after: Option<u64>,
}

impl PostError {
    fn from_reqwest(inner: reqwest::Error, transient: bool, retry_after: Option<u64>) -> Self {
        tracing::debug!(
            target: "mcptest_core::transport::streamable_http",
            error = %inner,
            transient,
            "post attempt failed"
        );
        Self {
            transient,
            retry_after,
        }
    }
}

fn should_retry(err: &PostError) -> bool {
    err.transient
}

/// One POST attempt. Reads the response, parses JSON or SSE depending on
/// `Content-Type`, and pushes results into `inbound_tx`.
async fn do_post(
    http: &HttpClient,
    url: &str,
    shared: &SharedState,
    signer: Option<&crate::auth::RequestSigner>,
    msg: &JsonRpcMessage,
    inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>,
) -> Result<(), PostError> {
    let resp = send_post_request(http, url, shared, signer, msg).await?;

    let status = resp.status();
    let hint = auth_hint(status, &resp);
    if !status.is_success() {
        tracing::warn!(
            target: "mcptest_core::transport::streamable_http",
            status = %status,
            "non-2xx HTTP response{hint}"
        );
    }
    let retry_after = parse_retry_after(&resp);

    // Capture session ID off any response. Skipped in stateless mode:
    // SEP-2575 strips the per-session table out, so a server that
    // returns `Mcp-Session-Id` against a stateless client is either a
    // mode mismatch or a bug, and capturing it would cause the next
    // request to look like a session re-entry.
    if matches!(shared.mode, TransportMode::Session) {
        if let Some(sid) = resp
            .headers()
            .get(SESSION_HEADER)
            .and_then(|v| v.to_str().ok())
        {
            shared.set_session(sid.to_owned());
        }
    }

    if is_retryable_status(status) {
        return Err(synthesize_retry_error(resp, retry_after));
    }

    // Past the retry gate: this response is final for the request, so its
    // headers are the ones the `response_headers` budget gate should see.
    // Recorded before `dispatch_response` consumes the body, which also
    // covers SSE-framed POST responses (WOR-1287).
    record_response_headers(shared, msg, &resp);

    // A non-retryable client error (401/403/404/400, etc.) is a transport-level
    // rejection with no JSON-RPC body to dispatch. For a request, synthesize an
    // error response keyed to its id and push it through `inbound_tx` so the
    // caller resolves immediately with the HTTP status, rather than waiting out
    // the request timeout and surfacing a misleading "timed out" on an auth-gated
    // server. (Dropping the writer's channel clone would not wake the caller: the
    // reader task holds its own clone, so the channel stays open.) A notification
    // has no caller to resolve, so we just log it. Either way the POST is handled,
    // not retried, and we keep the writer alive for later requests.
    if status.is_client_error() {
        if let Some(id) = msg.id.clone() {
            let _ = inbound_tx.send(JsonRpcMessage::error_response(
                id,
                JsonRpcErrorBody {
                    code: HTTP_TRANSPORT_ERROR_CODE,
                    message: format!("HTTP {status}{hint}"),
                    data: None,
                },
            ));
        }
        return Ok(());
    }

    dispatch_response(resp, status, retry_after, inbound_tx).await
}

/// Build and send the POST, attaching the JSON body, the standard MCP headers,
/// the session id when one is known, and any per-request signature headers.
/// Maps connect/timeout failures to a transient [`PostError`] so the retry
/// loop can back off; a signing failure is fatal (poisons the stream).
async fn send_post_request(
    http: &HttpClient,
    url: &str,
    shared: &SharedState,
    signer: Option<&crate::auth::RequestSigner>,
    msg: &JsonRpcMessage,
) -> Result<reqwest::Response, PostError> {
    // In stateless mode we splice `_meta` into the request's params
    // before serialization so the signer commits to the exact bytes
    // we send. Session mode serializes the borrowed message directly.
    let body = super::wire::build_outbound_body(msg, shared);
    let mut req = http
        .request(Method::POST, url)
        .header(CONTENT_TYPE, "application/json")
        .header(ACCEPT, "application/json, text/event-stream")
        .body(body.clone());
    // Session mode replays the captured `Mcp-Session-Id`; stateless
    // mode drops the header entirely (SEP-2575).
    if matches!(shared.mode, TransportMode::Session) {
        if let Some(sid) = shared.get_session() {
            req = req.header(SESSION_HEADER, sid);
        }
    }
    // Stateless mode also emits the SEP-2243 routing headers so a
    // server can dispatch without parsing the body. Session mode
    // never emits them: a 2025-* server has no rule for the values
    // and would have to ignore them anyway, but emitting unsolicited
    // routing hints could trip a strict proxy in front of it.
    if matches!(shared.mode, TransportMode::Stateless) {
        if let Some(method) = msg.method.as_deref() {
            req = req.header(METHOD_HEADER, method);
            if let Some(name) = super::wire::name_from_params(method, msg.params.as_ref()) {
                req = req.header(NAME_HEADER, name);
            }
        }
    }
    req = apply_signature(req, signer, "POST", url, &body)?;
    req.send().await.map_err(|e| {
        let transient = e.is_connect() || e.is_timeout();
        PostError::from_reqwest(e, transient, None)
    })
}

/// Record the response headers for a request envelope (WOR-1287).
///
/// Only requests with a numeric id participate: notifications have no
/// caller to claim the entry, and this engine never allocates string ids.
fn record_response_headers(shared: &SharedState, msg: &JsonRpcMessage, resp: &reqwest::Response) {
    let Some(id) = msg.id.as_ref().and_then(|i| i.as_i64()) else {
        return;
    };
    shared.header_capture.record(id, resp.headers());
}

/// True for the statuses the transport retries: 503, 429, and any 5xx.
fn is_retryable_status(status: StatusCode) -> bool {
    status == StatusCode::SERVICE_UNAVAILABLE
        || status == StatusCode::TOO_MANY_REQUESTS
        || status.is_server_error()
}

/// Turn a retryable response into a transient [`PostError`]. We need a
/// `reqwest::Error` to satisfy the inner type; `error_for_status_ref`
/// fabricates the proper variant for the 5xx/429/503 status.
fn synthesize_retry_error(resp: reqwest::Response, retry_after: Option<u64>) -> PostError {
    let err = resp.error_for_status_ref().err();
    let inner = err.unwrap_or_else(|| {
        // Defensive: the status was retryable so error_for_status_ref should
        // always return Some. If it does not, drop the response by calling
        // error_for_status to surface a typed error.
        match resp.error_for_status() {
            Ok(_) => unreachable!("retryable status passed error_for_status"),
            Err(e) => e,
        }
    });
    PostError::from_reqwest(inner, true, retry_after)
}

/// Dispatch a non-retryable response: stream SSE, accept an empty 202, or
/// parse the JSON body and push results into `inbound_tx`.
async fn dispatch_response(
    resp: reqwest::Response,
    status: StatusCode,
    retry_after: Option<u64>,
    inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>,
) -> Result<(), PostError> {
    let content_type = resp
        .headers()
        .get(CONTENT_TYPE)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_ascii_lowercase();

    if content_type.starts_with("text/event-stream") {
        consume_sse_stream(resp, inbound_tx)
            .await
            .map_err(|e| PostError::from_reqwest(e, true, retry_after))?;
        return Ok(());
    }

    // 202 Accepted is valid for notifications: no body, nothing to dispatch.
    if status == StatusCode::ACCEPTED {
        return Ok(());
    }

    let body = resp
        .bytes()
        .await
        .map_err(|e| PostError::from_reqwest(e, true, retry_after))?;
    if body.is_empty() {
        return Ok(());
    }
    dispatch_json_bytes(&body, inbound_tx);
    Ok(())
}

/// Parse a `Retry-After` header value as a whole number of seconds.
fn parse_retry_after(resp: &reqwest::Response) -> Option<u64> {
    resp.headers()
        .get(reqwest::header::RETRY_AFTER)
        .and_then(|v| v.to_str().ok())
        .and_then(|s| s.trim().parse::<u64>().ok())
}

/// Dispatch JSON bytes that may be a single envelope or an array.
pub(super) fn dispatch_json_bytes(body: &[u8], inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>) {
    let Ok(json) = serde_json::from_slice::<Value>(body) else {
        return;
    };
    dispatch_json_value(json, inbound_tx);
}

fn dispatch_json_value(value: Value, inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>) {
    match value {
        Value::Array(items) => {
            for item in items {
                if let Ok(msg) = serde_json::from_value::<JsonRpcMessage>(item) {
                    let _ = inbound_tx.send(msg);
                }
            }
        }
        other => {
            if let Ok(msg) = serde_json::from_value::<JsonRpcMessage>(other) {
                let _ = inbound_tx.send(msg);
            }
        }
    }
}

/// Read an SSE response body and dispatch every `message` event as a
/// JSON-RPC envelope.
async fn consume_sse_stream(
    resp: reqwest::Response,
    inbound_tx: &mpsc::UnboundedSender<JsonRpcMessage>,
) -> Result<(), reqwest::Error> {
    let mut stream = resp.bytes_stream();
    let mut parser = SseParser::new();
    while let Some(chunk) = stream.next().await {
        let bytes: Bytes = chunk?;
        parser.feed(&bytes);
        for ev in parser.take_ready() {
            if ev.data.is_empty() {
                continue;
            }
            dispatch_json_bytes(ev.data.as_bytes(), inbound_tx);
        }
    }
    Ok(())
}
