//! The native SGLang data-plane endpoints: `/generate` (submit a request, then
//! either fold egress frames to one unary JSON response or relay them as SSE
//! `data: {json}` … `[DONE]`, byte-compatible with Python
//! `http_server.generate_request`) and `/health` + `/health_generate` (which
//! round-trip a 1-token generate probe). Frame shaping (`meta_info`, logprob
//! tuples, cumulative vs incremental streams) lives here, as does
//! generate-request submission (`submit`); the shared `AppState` lives in the
//! parent `api_server` module.

use std::convert::Infallible;

use axum::{
    Json, Router,
    extract::State,
    extract::rejection::JsonRejection,
    http::StatusCode,
    response::{
        IntoResponse, Response,
        sse::{Event, Sse},
    },
    routing::{get, post},
};
use tokio::sync::mpsc;

use super::AppState;
use super::frame::{
    OutputAccumulator, cumulative_frame_string, error_value, frame_value, stream_frame_string,
    tag_value,
};
use super::guard::AbortGuard;
use super::submit::{pre_submit_error, submit};
use crate::environ::env_bool;
use crate::ids::Rid;
use crate::message::{EgressItem, GenerateBody, GenerateRequest, RequestKind, SamplingParams};

/// The routes this module owns, mounted by `api_server::serve`.
pub(super) fn routes() -> Router<AppState> {
    Router::new()
        .route("/generate", post(generate))
        .merge(health_routes())
}

/// `/health` + `/health_generate`. Both env knobs are resolved ONCE here, at
/// router build (server startup) — changing them on a live process needs a
/// restart. The deep-probe handler is built once with
/// `SGLANG_HEALTH_CHECK_TIMEOUT` frozen in and serves `/health_generate`
/// always; `SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION` (default true, mirroring
/// Python) decides whether `/health` shares it or is a plain 200 (routing the
/// request already proves the frontend is up).
fn health_routes() -> Router<AppState> {
    let timeout =
        std::time::Duration::from_secs(crate::environ::env_u64("SGLANG_HEALTH_CHECK_TIMEOUT", 20));
    let probe = get(move |state: State<AppState>| health_generate(state, timeout));
    let health = if env_bool("SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION", true) {
        probe.clone()
    } else {
        get(|| async { StatusCode::OK.into_response() })
    };
    Router::new()
        .route("/health", health)
        .route("/health_generate", probe)
}

/// Sentinel host that makes the KV connector no-op. Parity with
/// `sglang.srt.disaggregation.utils.FAKE_BOOTSTRAP_HOST`.
const FAKE_BOOTSTRAP_HOST: &str = "2.2.2.2";

/// `GET /health_generate` — deep health: confirm the scheduler → detok path is
/// producing output. 200 iff the egress heartbeat advances within `timeout`
/// (from `SGLANG_HEALTH_CHECK_TIMEOUT`, frozen at router build), else 503.
/// (`/health` uses the same handler when its env gate is on.)
///
/// Fires a pre-tokenized 1-token probe (`input_ids = [0]`, skips the tokenizer) so
/// an idle pipeline produces a frame, then watches the *global*
/// [`AppState::egress_activity`] counter (not the probe's own rid) — so a busy
/// server passes immediately and a backlog never false-503s (the analogue of
/// Python's `last_receive_tstamp`). The `HEALTH_CHECK` skip + `http_worker_ipc`
/// ack are irrelevant here: this single-process server owns the egress ring.
async fn health_generate(State(state): State<AppState>, timeout: std::time::Duration) -> Response {
    let baseline = state
        .egress_activity
        .load(std::sync::atomic::Ordering::Relaxed);

    // Fire the probe (the heartbeat is the signal, not its own response). A busy
    // scheduler skips it with no terminal frame, so its detok registration is
    // cleaned up only by the `AbortGuard` below.
    //
    // On a PD node the scheduler 400-aborts room-less requests, so inject the
    // same fake bootstrap pair Python uses (`FAKE_BOOTSTRAP_HOST` / room 0).
    let pd = state.server_args.is_disaggregation();
    let probe = GenerateRequest {
        // The `HEALTH_CHECK_<uuid>` rid form
        rid: Rid::new_health_check(),
        input_ids: Some(vec![0]),
        // One greedy token: the cheapest round-trip that still produces a frame.
        sampling_params: SamplingParams {
            max_new_tokens: Some(1),
            temperature: 0.0,
            ..Default::default()
        },
        stream: false,
        bootstrap_host: pd.then(|| FAKE_BOOTSTRAP_HOST.into()),
        bootstrap_room: pd.then_some(0),
        ..Default::default()
    };
    let (rid, _keepalive) =
        match submit(&state, RequestKind::Generate(Box::new(probe)), false).await {
            // Hold the receiver so the probe's sink stays open until it completes.
            Ok(v) => v,
            Err(resp) => return resp,
        };
    // Deregister on drop (never disarmed): a busy-skipped probe has no terminal
    // frame, so without this abort it leaks one detok entry per call.
    let _abort_guard = AbortGuard::new(state.senders.clone(), rid);

    // Watch the heartbeat advance (timeout frozen at router build, default 20s).
    let deadline = tokio::time::Instant::now() + timeout;
    loop {
        if state
            .egress_activity
            .load(std::sync::atomic::Ordering::Relaxed)
            != baseline
        {
            return StatusCode::OK.into_response();
        }
        if tokio::time::Instant::now() >= deadline {
            return StatusCode::SERVICE_UNAVAILABLE.into_response();
        }
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }
}

/// `POST /generate` — the native generation endpoint. Splits the body
/// into per-request payloads (a scalar body → one, a list body → a batch) and
/// dispatches to the single or batch path; a malformed body is a 400 before
/// anything reaches the scheduler.
///
/// The body is extracted as a `Result` so a deserialization failure is answered
/// with **400** (Python's status for a bad request) carrying serde's field-level
/// message, instead of axum's default 422.
async fn generate(
    State(state): State<AppState>,
    body: Result<Json<GenerateBody>, JsonRejection>,
) -> Response {
    let body = match body {
        Ok(Json(body)) => body,
        // A body that fails to parse has no readable `stream` flag, so this one
        // can only answer unary — as Python's does (FastAPI rejects before its
        // handler runs).
        Err(rejection) => {
            return pre_submit_error(StatusCode::BAD_REQUEST, &rejection.body_text(), false);
        }
    };
    let stream = body.stream;
    // Fan `text`/`input_ids`/`sampling_params` (scalar or list) into per-request
    // payloads. `is_batch` = list form → the response is a JSON array.
    let (payloads, is_batch) = match body.into_requests() {
        Ok(v) => v,
        // The error carries its own status (a bad batch is `Validation` → 400).
        Err(e) => {
            let code = StatusCode::from_u16(e.http_status()).unwrap_or(StatusCode::BAD_REQUEST);
            return pre_submit_error(code, &e.to_string(), stream);
        }
    };
    if !is_batch {
        // `into_requests` guarantees exactly one payload for a non-batch body.
        let payload = payloads
            .into_iter()
            .next()
            .expect("into_requests yields >=1 payload");
        generate_single(&state, payload, stream).await
    } else {
        generate_batch(&state, payloads, stream).await
    }
}

/// Answer an error raised *before* anything was submitted, in the shape the client
/// asked for.
///
/// A single (non-batched) `/generate`: submit one request, then either stream its
/// SSE frames or fold to one unary response.
async fn generate_single(state: &AppState, req: GenerateRequest, stream: bool) -> Response {
    // `return_text_in_logprobs` is decoded on the detok shard into `*_txt`, so
    // `frame_value` just reads them — no tokenizer needed here.
    let (rid_str, mut rx) = match submit(state, RequestKind::Generate(Box::new(req)), stream).await
    {
        Ok(v) => v,
        Err(resp) => return resp,
    };
    // Abort on client disconnect: the guard fires when dropped before the request
    // finishes (axum drops the handler/SSE stream). Disarmed on a natural terminal.
    // `rid_str` is the response `meta_info.id`, reused for every frame.
    let mut guard = AbortGuard::new(state.senders.clone(), rid_str.clone());
    // Cumulative frames (SGLang default) vs per-step deltas.
    let incremental = state.server_args.incremental_streaming_output;

    if stream {
        // A single request is a 1-element batch without the `index` field — reuse
        // the same stream so the frame/abort/truncation logic lives in one place.
        use futures::StreamExt;
        let s = generation_event_stream(vec![(rid_str, rx)], guard, incremental, false)
            .map(|data| Ok::<_, Infallible>(Event::default().data(data)));
        Sse::new(s).into_response()
    } else {
        // Unary: fold to the terminal, respond once. Disarm only on a real terminal
        // (a truncation leaves the guard armed so the scheduler work is aborted).
        let (status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await;
        if terminal {
            guard.disarm(&rid_str);
        }
        (status, Json(value)).into_response()
    }
}

/// Fold a unary request to its terminal → (HTTP status, result/`error` JSON, saw-terminal);
/// `false` = truncation, caller keeps the abort guard armed. Shared by single + batch.
async fn drain_unary(
    rx: &mut mpsc::Receiver<EgressItem>,
    rid_str: &str,
) -> (StatusCode, serde_json::Value, bool) {
    let mut acc = OutputAccumulator::default();
    while let Some(item) = rx.recv().await {
        match item {
            EgressItem::Frame(out) => acc.fold(&out),
            EgressItem::Done(out) => {
                acc.fold(&out);
                let final_out = acc.into_output();
                // A validation abort carries its own HTTP status + diagnostic.
                if let Some((code, message)) = final_out
                    .finish_reason
                    .as_ref()
                    .and_then(|f| f.abort_status())
                {
                    let status =
                        StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                    return (status, error_value(code, message), true);
                }
                return (StatusCode::OK, frame_value(&final_out, rid_str), true);
            }
            EgressItem::Error(e) => {
                let code = e.http_status();
                let status =
                    StatusCode::from_u16(code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                return (status, error_value(code, &e.to_string()), true);
            }
            EgressItem::Control(_) | EgressItem::Data(_) => continue, // never on `/generate`
        }
    }
    // Sender dropped without a terminal item: the shard dropped this request (a
    // truncation — a client disconnect would have dropped the handler future).
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        error_value(500, "response truncated before completion"),
        false,
    )
}

/// Batch `/generate`: submit all sub-requests first (scheduler runs them together),
/// then either (unary) drain each in order into a JSON array, or (streaming)
/// multiplex their streams into one SSE response, each frame carrying its `index`.
/// One [`AbortGuard`] covers the batch. A failed unary item is its own
/// `{ "error": … }` entry; the batch response is 200.
async fn generate_batch(
    state: &AppState,
    requests: Vec<GenerateRequest>,
    stream: bool,
) -> Response {
    // No cross-item rid collision to worry about: `into_requests` rejected duplicate
    // rids within this batch, and `Rid::from_client` made each one unique against
    // every other in-flight request.
    let mut guard = AbortGuard::new_empty(state.senders.clone());
    let mut receivers = Vec::with_capacity(requests.len());
    for req in requests {
        match submit(state, RequestKind::Generate(Box::new(req)), stream).await {
            Ok((rid, rx)) => {
                guard.arm(rid.clone());
                receivers.push((rid, rx));
            }
            Err(resp) => return resp,
        }
    }

    if stream {
        // Multiplex the N streams (mirrors the Python `_handle_batch_request` path);
        // `guard` moves into the stream so a disconnect aborts what's unfinished.
        use futures::StreamExt;
        let incremental = state.server_args.incremental_streaming_output;
        let s = generation_event_stream(receivers, guard, incremental, true)
            .map(|data| Ok::<_, Infallible>(Event::default().data(data)));
        Sse::new(s).into_response()
    } else {
        // Unary: drain each in order (already all submitted, so they run together).
        let mut results = Vec::with_capacity(receivers.len());
        for (rid_str, mut rx) in receivers {
            let (_status, value, terminal) = drain_unary(&mut rx, rid_str.client_facing()).await;
            if terminal {
                guard.disarm(&rid_str);
            }
            results.push(value);
        }
        (StatusCode::OK, Json(serde_json::Value::Array(results))).into_response()
    }
}

/// Await the next item from `rx`, then drain whatever queued behind it (so the caller
/// can coalesce a backlog, as Python's `state.out_list` does), handing the receiver
/// back for `FuturesUnordered` to re-poll. Empty result = channel closed.
async fn recv_indexed(
    index: usize,
    mut rx: mpsc::Receiver<EgressItem>,
) -> (usize, mpsc::Receiver<EgressItem>, Vec<EgressItem>) {
    let mut items = Vec::new();
    match rx.recv().await {
        Some(item) => items.push(item),
        None => return (index, rx, items), // closed
    }
    while let Ok(item) = rx.try_recv() {
        items.push(item);
    }
    (index, rx, items)
}

/// Multiplex `receivers` (one per request) into SSE `data` strings + a final `[DONE]`;
/// `with_index` tags each frame (batch only), `incremental` = delta vs cumulative,
/// `guard` aborts unfinished on drop.
fn generation_event_stream(
    receivers: Vec<(Rid, mpsc::Receiver<EgressItem>)>,
    mut guard: AbortGuard,
    incremental: bool,
    with_index: bool,
) -> impl futures::Stream<Item = String> {
    async_stream::stream! {
        use futures::StreamExt;

        let n = receivers.len();
        let rid_strs: Vec<Rid> = receivers.iter().map(|(rid, _)| rid.clone()).collect();
        let mut accs: Vec<OutputAccumulator> =
            (0..n).map(|_| OutputAccumulator::default()).collect();

        // Batch position, tagged onto every frame (a single request omits it).
        let idx = |i: usize| with_index.then_some(i);

        // Poll all receivers concurrently; re-arm a receiver's future after each
        // non-terminal frame so its stream keeps flowing.
        let mut futs = futures::stream::FuturesUnordered::new();
        for (i, (_, rx)) in receivers.into_iter().enumerate() {
            futs.push(recv_indexed(i, rx));
        }

        while let Some((i, rx, items)) = futs.next().await {
            if items.is_empty() {
                // Channel closed with no terminal → truncation for this item;
                // leave its rid armed so the scheduler work is aborted.
                yield tag_value(error_value(500, "response truncated before completion"), idx(i));
                continue;
            }

            // Cumulative frames supersede one another, so a drained backlog collapses
            // to its last (Python's `out_list[-1]`); deltas can't be dropped.
            let mut coalesced = false; // a cumulative frame is pending
            let mut terminal = None;   // (finish_reason) of a `Done` in this batch
            let mut failed = None;     // an `Error` in this batch

            for item in items {
                match item {
                    EgressItem::Frame(out) => {
                        accs[i].fold(&out);
                        if incremental {
                            yield stream_frame_string(out, &accs[i], true, rid_strs[i].client_facing(), idx(i));
                        } else {
                            coalesced = true;
                        }
                    }
                    EgressItem::Done(out) => {
                        accs[i].fold(&out);
                        terminal = Some(out);
                    }
                    EgressItem::Error(e) => failed = Some(e),
                    EgressItem::Control(_) | EgressItem::Data(_) => {} // never on /generate
                }
            }

            if let Some(e) = failed {
                yield tag_value(error_value(e.http_status(), &e.to_string()), idx(i));
                guard.disarm(&rid_strs[i]);
            } else if let Some(out) = terminal {
                // A validation abort → an error object, not a frame. The final frame
                // carries the full cumulative state, so any coalesced ones are moot.
                yield match out.finish_reason.as_ref().and_then(|f| f.abort_status()) {
                    Some((code, message)) => tag_value(error_value(code, message), idx(i)),
                    None => stream_frame_string(out, &accs[i], incremental, rid_strs[i].client_facing(), idx(i)),
                };
                guard.disarm(&rid_strs[i]); // terminal → not re-pushed
            } else {
                if coalesced {
                    yield cumulative_frame_string(&accs[i], rid_strs[i].client_facing(), idx(i));
                }
                futs.push(recv_indexed(i, rx)); // keep this item flowing
            }
        }
        yield "[DONE]".to_string();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::message::ChunkEvent;
    use crate::tokenizer_manager::Senders;
    use futures::StreamExt;
    fn senders() -> Senders {
        Senders {
            tm: flume::unbounded().0,
            abort: flume::unbounded().0,
            tok: flume::unbounded().0,
            detok: vec![],
        }
    }

    fn frame(rid: u64, text: &str) -> EgressItem {
        EgressItem::Frame(ChunkEvent {
            rid: Rid::from(rid.to_string()),
            text: text.into(),
            completion_tokens: 1,
            ..Default::default()
        })
    }
    fn done(rid: u64, text: &str) -> EgressItem {
        EgressItem::Done(ChunkEvent {
            rid: Rid::from(rid.to_string()),
            text: text.into(),
            completion_tokens: 1,
            // Parsed from the wire map Python emits, not a hand-built enum.
            finish_reason: Some(
                serde_json::from_value(serde_json::json!({"type": "length", "length": 1}))
                    .expect("finish reason must parse"),
            ),
            ..Default::default()
        })
    }
    fn parse(s: &str) -> serde_json::Value {
        serde_json::from_str(s).expect("frame is JSON")
    }

    /// Two sub-requests' frames interleave into one stream, each tagged with its
    /// batch `index`; text accumulates per item; `[DONE]` comes only after both
    /// terminate, then the stream ends.
    #[tokio::test]
    async fn interleaves_indexes_and_accumulates() {
        let (tx0, rx0) = mpsc::channel(8);
        let (tx1, rx1) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx0), ("11".into(), rx1)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true);
        futures::pin_mut!(stream);

        // Drive deterministically: exactly one channel has data before each poll.
        tx0.send(frame(10, "a")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 0);
        assert_eq!(v["text"], "a");

        tx1.send(frame(11, "b")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 1);
        assert_eq!(v["text"], "b");

        tx0.send(done(10, "!")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 0);
        assert_eq!(v["text"], "a!", "cumulative per item");
        assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");

        tx1.send(done(11, "?")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 1);
        assert_eq!(v["text"], "b?");

        assert_eq!(stream.next().await.unwrap(), "[DONE]");
        assert!(stream.next().await.is_none());
    }

    /// A per-item error is surfaced with its `index` and doesn't end the batch;
    /// `[DONE]` still waits for the other item.
    #[tokio::test]
    async fn per_item_error_carries_index() {
        let (tx0, rx0) = mpsc::channel(8);
        let (tx1, rx1) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx0), ("11".into(), rx1)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, true);
        futures::pin_mut!(stream);

        tx0.send(EgressItem::Error(crate::error::Error::Validation(
            "bad".into(),
        )))
        .await
        .unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 0);
        assert_eq!(v["error"]["code"], 400);

        tx1.send(done(11, "ok")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["index"], 1);

        assert_eq!(stream.next().await.unwrap(), "[DONE]");
    }

    /// `incremental=true`: each frame carries this step's **delta** text/output_ids,
    /// but `meta_info.completion_tokens` stays cumulative (matching Python).
    #[tokio::test]
    async fn incremental_emits_deltas_with_cumulative_count() {
        let (tx, rx) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, true);
        futures::pin_mut!(stream);

        tx.send(frame(10, "Hello")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], "Hello");
        assert_eq!(v["meta_info"]["completion_tokens"], 1);

        tx.send(frame(10, " world")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], " world", "delta, not cumulative 'Hello world'");
        assert_eq!(
            v["meta_info"]["completion_tokens"], 2,
            "count stays cumulative"
        );

        tx.send(done(10, "!")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], "!");
        assert_eq!(v["meta_info"]["completion_tokens"], 3);
        assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");

        assert_eq!(stream.next().await.unwrap(), "[DONE]");
    }

    /// The single-request shape (`with_index=false`, one receiver) omits the
    /// `index` field entirely, and still terminates with `[DONE]`.
    #[tokio::test]
    async fn single_shape_omits_index() {
        let (tx, rx) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false);
        futures::pin_mut!(stream);

        tx.send(done(10, "hi")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], "hi");
        assert!(v.get("index").is_none(), "single response has no index");

        assert_eq!(stream.next().await.unwrap(), "[DONE]");
    }

    /// A backlog of cumulative chunks collapses to a single frame carrying the latest
    /// state — each cumulative frame supersedes the last, so emitting the intermediate
    /// ones ships the full O(T) payload again for nothing. Mirrors the Python waiter's
    /// `out = out_list[-1]`. This is the whole point of draining in `recv_indexed`.
    #[tokio::test]
    async fn cumulative_backlog_coalesces_to_latest() {
        let (tx, rx) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), false, false);
        futures::pin_mut!(stream);

        // Three chunks queued before the stream is ever polled (a client falling behind).
        tx.send(frame(10, "a")).await.unwrap();
        tx.send(frame(10, "b")).await.unwrap();
        tx.send(frame(10, "c")).await.unwrap();

        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], "abc", "one frame, full cumulative text");
        assert_eq!(v["meta_info"]["completion_tokens"], 3, "no tokens lost");

        // The terminal frame still carries everything, and only then does [DONE] land.
        tx.send(done(10, "!")).await.unwrap();
        let v = parse(&stream.next().await.unwrap());
        assert_eq!(v["text"], "abc!");
        assert_eq!(v["meta_info"]["finish_reason"]["type"], "length");
        assert_eq!(stream.next().await.unwrap(), "[DONE]");
    }

    /// Incremental frames are *deltas*, so a backlog must emit every one — dropping
    /// any would silently lose tokens. Only the cumulative protocol may coalesce.
    #[tokio::test]
    async fn incremental_backlog_emits_every_delta() {
        let (tx, rx) = mpsc::channel(8);
        let receivers = vec![("10".into(), rx)];
        let stream =
            generation_event_stream(receivers, AbortGuard::new_empty(senders()), true, false);
        futures::pin_mut!(stream);

        tx.send(frame(10, "a")).await.unwrap();
        tx.send(frame(10, "b")).await.unwrap();
        tx.send(frame(10, "c")).await.unwrap();

        for (n, expect) in [(1, "a"), (2, "b"), (3, "c")] {
            let v = parse(&stream.next().await.unwrap());
            assert_eq!(v["text"], expect, "delta {n} must not be dropped");
            assert_eq!(
                v["meta_info"]["completion_tokens"], n,
                "count stays cumulative"
            );
        }
    }
}