use std::future::Future;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};

use std::collections::HashMap;

use axum::extract::{DefaultBodyLimit, State};
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response as HttpResponse};
use axum::routing::{get, post};
use axum::{Json, Router};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram};
use metrics_exporter_prometheus::{Matcher, PrometheusBuilder, PrometheusHandle};
use serde::{Deserialize, Serialize};
use tokio::sync::{Semaphore, mpsc, oneshot, watch};
use tokio::task::JoinSet;
use tonic::{Request, Response as GrpcResponse, Status};
use tonic_health::{ServingStatus, server::health_reporter};

use crate::backend::{BackendMetadata, InferenceBackend};
use crate::error::{NarsilError, Result};
use crate::schema::proto::inference_server::{Inference, InferenceServer};
use crate::schema::proto::{PredictRequest, PredictResponse};
use crate::schema::{DataType, Tensor};
use crate::schema::{InferenceRequest, InferenceResponse};

const REQUESTS_TOTAL: &str = "narsil_requests_total";
const REQUEST_LATENCY_SECONDS: &str = "narsil_request_latency_seconds";
const IN_FLIGHT_REQUESTS: &str = "narsil_in_flight_requests";
// Client-supplied batch size: how many requests a single transport call carried.
const BATCH_SIZE: &str = "narsil_batch_size";
// Effective fused batch size formed by the continuous-batching collector before it
// calls `infer_batch` on the backend. This is the batch the GPU actually sees.
const FUSED_BATCH_SIZE: &str = "narsil_fused_batch_size";
const PROMETHEUS_CONTENT_TYPE: &str = "text/plain; version=0.0.4; charset=utf-8";
const LATENCY_BUCKETS: &[f64] = &[
    0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0,
];
const BATCH_SIZE_BUCKETS: &[f64] = &[1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0];

struct InferenceJob {
    request: InferenceRequest,
    response: oneshot::Sender<Result<InferenceResponse>>,
}

struct InferenceBatchJob {
    requests: Vec<InferenceRequest>,
    response: oneshot::Sender<Vec<Result<InferenceResponse>>>,
}

enum WorkerJob {
    Infer(InferenceJob),
    InferBatch(InferenceBatchJob),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct WorkerPoolConfig {
    pub workers: usize,
    pub queue_depth: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BatchingConfig {
    pub max_batch_size: usize,
    pub batch_timeout_ms: u64,
    pub workers: usize,
    pub queue_depth: usize,
}

impl BatchingConfig {
    pub fn new(
        max_batch_size: usize,
        batch_timeout_ms: u64,
        workers: usize,
        queue_depth: usize,
    ) -> Result<Self> {
        if max_batch_size == 0 {
            return Err(NarsilError::InvalidRequest(
                "max_batch_size must be positive".to_string(),
            ));
        }
        if workers == 0 {
            return Err(NarsilError::InvalidRequest(
                "workers must be positive".to_string(),
            ));
        }
        if queue_depth == 0 {
            return Err(NarsilError::InvalidRequest(
                "queue_depth must be positive".to_string(),
            ));
        }
        Ok(Self {
            max_batch_size,
            batch_timeout_ms,
            workers,
            queue_depth,
        })
    }
}

impl WorkerPoolConfig {
    pub fn new(workers: usize, queue_depth: usize) -> Result<Self> {
        if workers == 0 {
            return Err(NarsilError::InvalidRequest(
                "inference worker count must be positive".to_string(),
            ));
        }
        if queue_depth == 0 {
            return Err(NarsilError::InvalidRequest(
                "inference queue depth must be positive".to_string(),
            ));
        }
        Ok(Self {
            workers,
            queue_depth,
        })
    }
}

#[derive(Debug)]
struct InferenceWorkerPool<B: InferenceBackend + Clone> {
    senders: Vec<mpsc::Sender<WorkerJob>>,
    next_worker: AtomicUsize,
    join_handles: Mutex<Vec<JoinHandle<()>>>,
    metadata: BackendMetadata,
    _backend: std::marker::PhantomData<B>,
}

impl<B: InferenceBackend + Clone> InferenceWorkerPool<B> {
    fn new(backend: B, config: WorkerPoolConfig) -> Result<Self> {
        let metadata = backend.metadata();
        let mut senders = Vec::with_capacity(config.workers);
        let mut join_handles = Vec::with_capacity(config.workers);

        for worker_id in 0..config.workers {
            let (sender, mut receiver) = mpsc::channel::<WorkerJob>(config.queue_depth);
            let worker_backend = backend.clone();
            let join_handle = std::thread::Builder::new()
                .name(format!("narsil-inference-{worker_id}"))
                .spawn(move || {
                    while let Some(job) = receiver.blocking_recv() {
                        match job {
                            WorkerJob::Infer(job) => {
                                if job.response.is_closed() {
                                    continue;
                                }
                                let result = worker_backend.infer(job.request);
                                let _ = job.response.send(result);
                            }
                            WorkerJob::InferBatch(job) => {
                                if job.response.is_closed() {
                                    continue;
                                }
                                let results = worker_backend.infer_batch(job.requests);
                                let _ = job.response.send(results);
                            }
                        }
                    }
                })
                .map_err(|err| {
                    NarsilError::Transport(format!(
                        "failed to spawn inference worker {worker_id}: {err}"
                    ))
                })?;

            senders.push(sender);
            join_handles.push(join_handle);
        }

        Ok(Self {
            senders,
            next_worker: AtomicUsize::new(0),
            join_handles: Mutex::new(join_handles),
            metadata,
            _backend: std::marker::PhantomData,
        })
    }

    async fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        let worker_index = self.next_worker.fetch_add(1, Ordering::Relaxed) % self.senders.len();
        let (response, receiver) = oneshot::channel();
        let job = WorkerJob::Infer(InferenceJob { request, response });

        self.senders[worker_index]
            .try_send(job)
            .map_err(|err| match err {
                mpsc::error::TrySendError::Full(_) => NarsilError::Overloaded(format!(
                    "inference worker {worker_index} queue is full"
                )),
                mpsc::error::TrySendError::Closed(_) => NarsilError::Transport(format!(
                    "inference worker {worker_index} is unavailable"
                )),
            })?;

        receiver.await.map_err(|_| {
            NarsilError::Transport(format!("inference worker {worker_index} dropped response"))
        })?
    }

    async fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        let request_count = requests.len();
        if request_count == 0 {
            return Vec::new();
        }

        let worker_index = self.next_worker.fetch_add(1, Ordering::Relaxed) % self.senders.len();
        let (response, receiver) = oneshot::channel();
        let job = WorkerJob::InferBatch(InferenceBatchJob { requests, response });

        if let Err(err) = self.senders[worker_index].try_send(job) {
            let error = match err {
                mpsc::error::TrySendError::Full(_) => NarsilError::Overloaded(format!(
                    "inference worker {worker_index} queue is full"
                )),
                mpsc::error::TrySendError::Closed(_) => NarsilError::Transport(format!(
                    "inference worker {worker_index} is unavailable"
                )),
            };
            return repeat_error(request_count, error);
        }

        receiver.await.unwrap_or_else(|_| {
            repeat_error(
                request_count,
                NarsilError::Transport(format!(
                    "inference worker {worker_index} dropped batch response"
                )),
            )
        })
    }
}

impl<B: InferenceBackend + Clone> Drop for InferenceWorkerPool<B> {
    fn drop(&mut self) {
        self.senders.clear();
        let mut handles = self
            .join_handles
            .lock()
            .expect("inference worker handles mutex poisoned");
        for handle in handles.drain(..) {
            let _ = handle.join();
        }
    }
}

/// RAII guard that clears the batch collector's `loop_alive` flag the moment the
/// collector task ends, no matter which exit path (break or panic-unwind) is taken.
struct LoopAliveGuard(Arc<AtomicBool>);

impl Drop for LoopAliveGuard {
    fn drop(&mut self) {
        self.0.store(false, Ordering::SeqCst);
    }
}

struct BatchingWorkerPool<B: InferenceBackend + Clone> {
    sender: mpsc::Sender<InferenceJob>,
    max_batch_size: usize,
    _semaphore: Arc<Semaphore>,
    _backend: B,
    /// `true` while the collector loop is running; flipped to `false` the moment the
    /// loop exits (channel closed, semaphore closed) so readiness probes can detect a
    /// dead batch loop instead of reporting ready forever.
    loop_alive: Arc<AtomicBool>,
    _handle: tokio::task::JoinHandle<()>,
}

impl<B: InferenceBackend + Clone> Drop for BatchingWorkerPool<B> {
    fn drop(&mut self) {
        self._handle.abort();
    }
}

impl<B: InferenceBackend + Clone> BatchingWorkerPool<B> {
    fn new(backend: B, config: BatchingConfig) -> Result<Self> {
        let (sender, mut receiver) = mpsc::channel::<InferenceJob>(config.queue_depth);
        let semaphore = Arc::new(Semaphore::new(config.workers));
        let worker_semaphore = Arc::clone(&semaphore);
        let worker_backend = backend.clone();
        let max_batch_size = config.max_batch_size;
        let batch_timeout_ms = config.batch_timeout_ms;
        let loop_alive = Arc::new(AtomicBool::new(true));
        let worker_loop_alive = Arc::clone(&loop_alive);

        let handle = tokio::spawn(async move {
            // Mark the loop dead on any exit path so readiness probes stop reporting ready
            // if the collector dies (channel/semaphore closed) while the process lives on.
            let _alive_guard = LoopAliveGuard(worker_loop_alive);
            loop {
                let permit = match Arc::clone(&worker_semaphore).acquire_owned().await {
                    Ok(p) => p,
                    Err(_) => break,
                };

                let first = match receiver.recv().await {
                    Some(j) => j,
                    None => break,
                };
                let mut batch: Vec<InferenceJob> = vec![first];

                let deadline = tokio::time::Instant::now()
                    + std::time::Duration::from_millis(batch_timeout_ms);
                loop {
                    if batch.len() >= max_batch_size {
                        break;
                    }
                    match tokio::time::timeout_at(deadline, receiver.recv()).await {
                        Ok(Some(job)) => batch.push(job),
                        _ => break,
                    }
                }

                batch.retain(|j| !j.response.is_closed());
                if batch.is_empty() {
                    continue;
                }

                record_fused_batch_size(batch.len());

                let backend_clone = worker_backend.clone();
                let requests: Vec<_> = batch.iter().map(|j| j.request.clone()).collect();

                tokio::task::spawn_blocking(move || {
                    let results = backend_clone.infer_batch(requests);
                    for (job, result) in batch.into_iter().zip(results) {
                        let _ = job.response.send(result);
                    }
                    drop(permit);
                });
            }
        });

        Ok(Self {
            sender,
            max_batch_size: config.max_batch_size,
            _semaphore: semaphore,
            _backend: backend,
            loop_alive,
            _handle: handle,
        })
    }

    fn is_alive(&self) -> bool {
        self.loop_alive.load(Ordering::SeqCst)
    }

    fn submit(
        &self,
        request: InferenceRequest,
    ) -> Result<oneshot::Receiver<Result<InferenceResponse>>> {
        let (response, receiver) = oneshot::channel();
        let job = InferenceJob { request, response };
        self.sender.try_send(job).map_err(|_| {
            NarsilError::Overloaded("batch collector queue full or unavailable".to_string())
        })?;
        Ok(receiver)
    }

    async fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        let receiver = self.submit(request)?;
        receiver
            .await
            .map_err(|_| NarsilError::Transport("batch collector dropped response".to_string()))?
    }

    /// Route a multi-request batch through the *same* continuous-batching collector that
    /// single `infer` calls use. Each request is enqueued as an individual job, so the
    /// collector applies its usual admission (bounded queue) and grouping logic instead of
    /// contending for the worker semaphore directly — which previously starved direct
    /// batches whenever the collector held the only permit. Order is preserved by awaiting
    /// receivers in submission order.
    async fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        let request_count = requests.len();
        if request_count == 0 {
            return Vec::new();
        }

        // Enqueue every request first so the collector can group them; bail per-item if the
        // bounded queue rejects admission (consistent with the single-infer path).
        let slots: Vec<Slot> = requests
            .into_iter()
            .map(|request| match self.submit(request) {
                Ok(receiver) => Slot::Pending(receiver),
                Err(err) => Slot::Rejected(err),
            })
            .collect();

        let mut results = Vec::with_capacity(request_count);
        for slot in slots {
            match slot {
                Slot::Pending(receiver) => results.push(receiver.await.unwrap_or_else(|_| {
                    Err(NarsilError::Transport(
                        "batch collector dropped response".to_string(),
                    ))
                })),
                Slot::Rejected(err) => results.push(Err(err)),
            }
        }
        results
    }
}

/// Helper for [`BatchingWorkerPool::infer_batch`]: either a pending collector response or a
/// per-item admission failure, keeping per-request ordering intact.
enum Slot {
    Pending(oneshot::Receiver<Result<InferenceResponse>>),
    Rejected(NarsilError),
}

fn repeat_error(count: usize, error: NarsilError) -> Vec<Result<InferenceResponse>> {
    match error {
        NarsilError::InvalidRequest(message) => (0..count)
            .map(|_| Err(NarsilError::InvalidRequest(message.clone())))
            .collect(),
        NarsilError::Schema(message) => (0..count)
            .map(|_| Err(NarsilError::Schema(message.clone())))
            .collect(),
        NarsilError::Backend(message) => (0..count)
            .map(|_| Err(NarsilError::Backend(message.clone())))
            .collect(),
        NarsilError::Transport(message) => (0..count)
            .map(|_| Err(NarsilError::Transport(message.clone())))
            .collect(),
        NarsilError::Overloaded(message) => (0..count)
            .map(|_| Err(NarsilError::Overloaded(message.clone())))
            .collect(),
    }
}

/// Type-erased liveness/readiness probe into the (generic) inference service so that
/// the non-generic `ObservabilityState` can consult real backend/executor health.
type HealthProbe = Arc<dyn Fn() -> bool + Send + Sync>;

#[derive(Clone)]
pub struct ObservabilityState {
    /// Set to `true` once a shutdown/drain has been requested; readiness then reports
    /// not-ready so load balancers stop routing new traffic.
    draining: Arc<AtomicBool>,
    /// Real liveness probe (e.g. is the batch collector loop alive?).
    liveness: HealthProbe,
    /// Real readiness probe (e.g. backend loaded + executor live).
    readiness: HealthProbe,
    /// When `true`, `/metrics` is also mounted on the public inference HTTP gateway.
    /// Defaults to `false` so metrics are scrape-only on the private observability
    /// listener; the single-port [`http_router`] convenience helper opts in.
    expose_metrics_on_gateway: bool,
    prometheus: PrometheusHandle,
}

impl ObservabilityState {
    /// Build observability state wired to a concrete inference service so that
    /// liveness/readiness reflect the real backend and executor health.
    pub fn new<B: InferenceBackend + Clone>(service: &SharedInferenceService<B>) -> Result<Self> {
        let liveness_service = service.clone();
        let readiness_service = service.clone();
        Ok(Self {
            draining: Arc::new(AtomicBool::new(false)),
            liveness: Arc::new(move || liveness_service.is_live()),
            readiness: Arc::new(move || readiness_service.is_ready()),
            expose_metrics_on_gateway: false,
            prometheus: prometheus_handle()?,
        })
    }

    /// Opt in to also serving `/metrics` on the public inference HTTP gateway. Used by
    /// the single-port [`http_router`] convenience helper and dev/test setups; in
    /// production prefer scraping the dedicated, privately bound observability listener.
    pub fn with_metrics_on_gateway(mut self) -> Self {
        self.expose_metrics_on_gateway = true;
        self
    }

    /// Liveness: the process and its execution machinery are still healthy. Draining
    /// does not flip liveness — a draining pod is still alive while it finishes work.
    pub fn is_live(&self) -> bool {
        (self.liveness)()
    }

    /// Readiness: not draining, and the backend is loaded and the executor is live.
    pub fn is_ready(&self) -> bool {
        !self.draining.load(Ordering::SeqCst) && (self.readiness)()
    }

    /// Begin (or clear) draining. During drain, readiness reports not-ready.
    pub fn set_draining(&self, draining: bool) {
        self.draining.store(draining, Ordering::SeqCst);
    }

    pub fn render_metrics(&self) -> String {
        self.prometheus.run_upkeep();
        self.prometheus.render()
    }
}

fn prometheus_handle() -> Result<PrometheusHandle> {
    static HANDLE: OnceLock<PrometheusHandle> = OnceLock::new();
    // `install_recorder()` can only succeed once per process (it sets the global
    // metrics recorder). Serialize initialization with a mutex so concurrent callers
    // (e.g. parallel tests) don't both race to `install_recorder()` and have the loser
    // spuriously fail; the first caller installs, the rest clone the cached handle.
    static INIT_LOCK: Mutex<()> = Mutex::new(());

    if let Some(handle) = HANDLE.get() {
        return Ok(handle.clone());
    }

    let _guard = INIT_LOCK.lock().expect("metrics init mutex poisoned");
    if let Some(handle) = HANDLE.get() {
        return Ok(handle.clone());
    }

    let builder = PrometheusBuilder::new()
        .set_buckets_for_metric(
            Matcher::Full(REQUEST_LATENCY_SECONDS.to_string()),
            LATENCY_BUCKETS,
        )
        .map_err(|err| {
            NarsilError::Transport(format!(
                "failed to configure Prometheus latency buckets: {err}"
            ))
        })?
        .set_buckets_for_metric(Matcher::Full(BATCH_SIZE.to_string()), BATCH_SIZE_BUCKETS)
        .map_err(|err| {
            NarsilError::Transport(format!(
                "failed to configure Prometheus batch buckets: {err}"
            ))
        })?
        .set_buckets_for_metric(
            Matcher::Full(FUSED_BATCH_SIZE.to_string()),
            BATCH_SIZE_BUCKETS,
        )
        .map_err(|err| {
            NarsilError::Transport(format!(
                "failed to configure Prometheus fused batch buckets: {err}"
            ))
        })?;
    let handle = builder.install_recorder().map_err(|err| {
        NarsilError::Transport(format!("failed to install metrics recorder: {err}"))
    })?;

    register_metric_descriptions();
    let _ = HANDLE.set(handle.clone());
    Ok(handle)
}

fn register_metric_descriptions() {
    describe_counter!(
        REQUESTS_TOTAL,
        "Total inference transport requests by transport, endpoint, and status."
    );
    describe_histogram!(
        REQUEST_LATENCY_SECONDS,
        "Inference transport request latency in seconds by transport, endpoint, and status."
    );
    describe_gauge!(
        IN_FLIGHT_REQUESTS,
        "Current in-flight inference transport requests by transport and endpoint."
    );
    describe_histogram!(
        BATCH_SIZE,
        "Client-supplied request batch size by transport and endpoint (requests per transport call)."
    );
    describe_histogram!(
        FUSED_BATCH_SIZE,
        "Effective fused batch size the continuous-batching collector forms before calling infer_batch."
    );
}

/// Record the effective fused batch size formed by the continuous-batching collector.
/// This is the batch that is actually handed to the backend's `infer_batch`.
fn record_fused_batch_size(fused_batch_size: usize) {
    histogram!(FUSED_BATCH_SIZE).record(fused_batch_size as f64);
}

struct InFlightMetric {
    transport: &'static str,
    endpoint: &'static str,
}

impl InFlightMetric {
    fn new(transport: &'static str, endpoint: &'static str) -> Self {
        gauge!(
            IN_FLIGHT_REQUESTS,
            "transport" => transport,
            "endpoint" => endpoint
        )
        .increment(1.0);
        Self {
            transport,
            endpoint,
        }
    }
}

impl Drop for InFlightMetric {
    fn drop(&mut self) {
        gauge!(
            IN_FLIGHT_REQUESTS,
            "transport" => self.transport,
            "endpoint" => self.endpoint
        )
        .decrement(1.0);
    }
}

fn record_request_metrics(
    transport: &'static str,
    endpoint: &'static str,
    status: &'static str,
    elapsed: Duration,
    batch_size: usize,
) {
    counter!(
        REQUESTS_TOTAL,
        "transport" => transport,
        "endpoint" => endpoint,
        "status" => status
    )
    .increment(1);
    histogram!(
        REQUEST_LATENCY_SECONDS,
        "transport" => transport,
        "endpoint" => endpoint,
        "status" => status
    )
    .record(elapsed.as_secs_f64());
    histogram!(
        BATCH_SIZE,
        "transport" => transport,
        "endpoint" => endpoint
    )
    .record(batch_size as f64);
}

async fn instrumented_infer<B: InferenceBackend + Clone>(
    service: &SharedInferenceService<B>,
    transport: &'static str,
    endpoint: &'static str,
    request: InferenceRequest,
) -> Result<InferenceResponse> {
    let _in_flight = InFlightMetric::new(transport, endpoint);
    let start = Instant::now();
    let response = service.infer(request).await;
    let status = if response.is_ok() { "ok" } else { "error" };
    record_request_metrics(transport, endpoint, status, start.elapsed(), 1);
    response
}

async fn instrumented_infer_batch<B: InferenceBackend + Clone>(
    service: &SharedInferenceService<B>,
    transport: &'static str,
    endpoint: &'static str,
    requests: Vec<InferenceRequest>,
) -> Vec<Result<InferenceResponse>> {
    let batch_size = requests.len();
    let _in_flight = InFlightMetric::new(transport, endpoint);
    let start = Instant::now();
    let responses = service.infer_batch(requests).await;
    let status = if responses.iter().all(Result::is_ok) {
        "ok"
    } else {
        "error"
    };
    record_request_metrics(transport, endpoint, status, start.elapsed(), batch_size);
    responses
}

enum InferenceExecutor<B: InferenceBackend + Clone> {
    Inline(Arc<B>),
    WorkerPool(Arc<InferenceWorkerPool<B>>),
    Batching(Arc<BatchingWorkerPool<B>>),
}

impl<B: InferenceBackend + Clone> Clone for InferenceExecutor<B> {
    fn clone(&self) -> Self {
        match self {
            Self::Inline(backend) => Self::Inline(Arc::clone(backend)),
            Self::WorkerPool(pool) => Self::WorkerPool(Arc::clone(pool)),
            Self::Batching(pool) => Self::Batching(Arc::clone(pool)),
        }
    }
}

impl<B: InferenceBackend + Clone> InferenceExecutor<B> {
    async fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        match self {
            Self::Inline(backend) => backend.infer(request),
            Self::WorkerPool(pool) => pool.infer(request).await,
            Self::Batching(pool) => pool.infer(request).await,
        }
    }

    async fn infer_batch(&self, requests: Vec<InferenceRequest>) -> Vec<Result<InferenceResponse>> {
        match self {
            Self::Inline(backend) => backend.infer_batch(requests),
            Self::WorkerPool(pool) => pool.infer_batch(requests).await,
            Self::Batching(pool) => pool.infer_batch(requests).await,
        }
    }

    /// Maximum number of requests the executor will group into a single batch, when one is
    /// configured (i.e. in [`InferenceExecutor::Batching`] mode). Used by the HTTP gateway
    /// to bound `/v1/infer_batch` payloads.
    fn max_batch_size(&self) -> Option<usize> {
        match self {
            Self::Inline(_) | Self::WorkerPool(_) => None,
            Self::Batching(pool) => Some(pool.max_batch_size),
        }
    }

    fn metadata(&self) -> BackendMetadata {
        match self {
            Self::Inline(backend) => backend.metadata(),
            Self::WorkerPool(pool) => pool.metadata.clone(),
            Self::Batching(pool) => pool._backend.metadata(),
        }
    }

    /// Liveness: the underlying execution machinery is still able to make progress.
    /// Only the continuous-batching collector has a background loop that can die out
    /// from under us; the inline and worker-pool executors run inline per request.
    fn is_live(&self) -> bool {
        match self {
            Self::Inline(_) | Self::WorkerPool(_) => true,
            Self::Batching(pool) => pool.is_alive(),
        }
    }
}

#[derive(Clone)]
pub struct SharedInferenceService<B: InferenceBackend + Clone> {
    executor: InferenceExecutor<B>,
}

impl<B: InferenceBackend + Clone> SharedInferenceService<B> {
    pub fn new(backend: B) -> Self {
        Self {
            executor: InferenceExecutor::Inline(Arc::new(backend)),
        }
    }

    pub fn from_arc(backend: Arc<B>) -> Self {
        Self {
            executor: InferenceExecutor::Inline(backend),
        }
    }

    pub fn with_worker_pool(backend: B, config: WorkerPoolConfig) -> Result<Self> {
        Ok(Self {
            executor: InferenceExecutor::WorkerPool(Arc::new(InferenceWorkerPool::new(
                backend, config,
            )?)),
        })
    }

    pub fn with_batching(backend: B, config: BatchingConfig) -> Result<Self> {
        Ok(Self {
            executor: InferenceExecutor::Batching(Arc::new(BatchingWorkerPool::new(
                backend, config,
            )?)),
        })
    }

    pub async fn infer(&self, request: InferenceRequest) -> Result<InferenceResponse> {
        self.executor.infer(request).await
    }

    pub async fn infer_batch(
        &self,
        requests: Vec<InferenceRequest>,
    ) -> Vec<Result<InferenceResponse>> {
        self.executor.infer_batch(requests).await
    }

    /// Configured batch grouping limit, if the service runs in batching mode. The HTTP
    /// gateway uses this to reject oversized `/v1/infer_batch` payloads up front.
    pub fn max_batch_size(&self) -> Option<usize> {
        self.executor.max_batch_size()
    }

    /// Metadata reported by the loaded backend.
    pub fn metadata(&self) -> BackendMetadata {
        self.executor.metadata()
    }

    /// Liveness: the service's execution machinery can still make progress. For the
    /// continuous-batching executor this reflects whether the collector loop is alive.
    pub fn is_live(&self) -> bool {
        self.executor.is_live()
    }

    /// Readiness: the backend has produced valid metadata (a non-empty name, i.e. a
    /// model is actually loaded) and the executor is live and able to serve requests.
    pub fn is_ready(&self) -> bool {
        self.is_live() && !self.metadata().name.trim().is_empty()
    }
}

#[derive(Clone)]
pub struct GrpcInferenceService<B: InferenceBackend + Clone> {
    service: SharedInferenceService<B>,
}

impl<B: InferenceBackend + Clone> GrpcInferenceService<B> {
    pub fn new(backend: B) -> Self {
        Self {
            service: SharedInferenceService::new(backend),
        }
    }

    pub fn from_arc(backend: Arc<B>) -> Self {
        Self {
            service: SharedInferenceService::from_arc(backend),
        }
    }

    pub fn from_shared(service: SharedInferenceService<B>) -> Self {
        Self { service }
    }

    pub fn with_worker_pool(backend: B, config: WorkerPoolConfig) -> Result<Self> {
        Ok(Self {
            service: SharedInferenceService::with_worker_pool(backend, config)?,
        })
    }

    pub fn with_batching(backend: B, config: BatchingConfig) -> Result<Self> {
        Ok(Self {
            service: SharedInferenceService::with_batching(backend, config)?,
        })
    }

    pub fn into_server(self) -> InferenceServer<Self> {
        InferenceServer::new(self)
    }
}

#[tonic::async_trait]
impl<B: InferenceBackend + Clone> Inference for GrpcInferenceService<B> {
    async fn predict(
        &self,
        request: Request<PredictRequest>,
    ) -> std::result::Result<GrpcResponse<PredictResponse>, Status> {
        let request =
            InferenceRequest::try_from(request.into_inner()).map_err(NarsilError::into_status)?;
        let response = instrumented_infer(&self.service, "grpc", "predict", request)
            .await
            .map_err(NarsilError::into_status)?;

        Ok(GrpcResponse::new(PredictResponse::from(response)))
    }
}

#[derive(Clone)]
struct HttpState<B: InferenceBackend + Clone> {
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
}

/// Default cap on the size of an inbound HTTP/JSON request body, in bytes (2 MiB). Bounds
/// memory amplification from base64-encoded tensor payloads; oversized bodies are rejected
/// by axum with `413 Payload Too Large` before any handler runs.
pub const DEFAULT_BODY_LIMIT_BYTES: usize = 2 * 1024 * 1024;

/// Default per-request timeout for the HTTP gateway. A backend that hangs past this yields a
/// `503 Service Unavailable` to the client instead of holding the connection open forever.
pub const DEFAULT_REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonPredictRequest {
    pub model_name: String,
    #[serde(default)]
    pub inputs: Vec<JsonTensor>,
    #[serde(default)]
    pub parameters: HashMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonPredictResponse {
    pub outputs: Vec<JsonTensor>,
    pub metadata: HashMap<String, String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonInferBatchRequest {
    pub requests: Vec<JsonPredictRequest>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonInferBatchResponse {
    pub responses: Vec<JsonBatchResponse>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "status")]
pub enum JsonBatchResponse {
    Ok { response: JsonPredictResponse },
    Error { error: JsonError },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonTensor {
    pub name: String,
    pub dtype: JsonDataType,
    pub shape: Vec<usize>,
    pub data: String,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum JsonDataType {
    F32,
    F64,
    I64,
    I32,
    U8,
    Bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JsonError {
    pub code: String,
    pub message: String,
}

impl TryFrom<JsonPredictRequest> for InferenceRequest {
    type Error = NarsilError;

    fn try_from(value: JsonPredictRequest) -> Result<Self> {
        if value.inputs.is_empty() {
            return Err(NarsilError::InvalidRequest(
                "predict request must contain at least one input tensor".to_string(),
            ));
        }

        Ok(Self {
            model_name: value.model_name,
            inputs: value
                .inputs
                .into_iter()
                .map(Tensor::try_from)
                .collect::<Result<Vec<_>>>()?,
            parameters: value.parameters,
        })
    }
}

impl From<InferenceResponse> for JsonPredictResponse {
    fn from(value: InferenceResponse) -> Self {
        Self {
            outputs: value.outputs.into_iter().map(JsonTensor::from).collect(),
            metadata: value.metadata,
        }
    }
}

impl TryFrom<JsonTensor> for Tensor {
    type Error = NarsilError;

    fn try_from(value: JsonTensor) -> Result<Self> {
        let data = BASE64.decode(value.data.as_bytes()).map_err(|err| {
            NarsilError::Schema(format!(
                "tensor '{}' has invalid base64 data: {err}",
                value.name
            ))
        })?;
        Self::new(value.name, DataType::from(value.dtype), value.shape, data)
    }
}

impl From<Tensor> for JsonTensor {
    fn from(value: Tensor) -> Self {
        Self {
            name: value.name,
            dtype: JsonDataType::from(value.dtype),
            shape: value.shape,
            data: BASE64.encode(value.data),
        }
    }
}

impl From<JsonDataType> for DataType {
    fn from(value: JsonDataType) -> Self {
        match value {
            JsonDataType::F32 => Self::F32,
            JsonDataType::F64 => Self::F64,
            JsonDataType::I64 => Self::I64,
            JsonDataType::I32 => Self::I32,
            JsonDataType::U8 => Self::U8,
            JsonDataType::Bool => Self::Bool,
        }
    }
}

impl From<DataType> for JsonDataType {
    fn from(value: DataType) -> Self {
        match value {
            DataType::F32 => Self::F32,
            DataType::F64 => Self::F64,
            DataType::I64 => Self::I64,
            DataType::I32 => Self::I32,
            DataType::U8 => Self::U8,
            DataType::Bool => Self::Bool,
        }
    }
}

impl From<NarsilError> for JsonError {
    fn from(value: NarsilError) -> Self {
        match value {
            NarsilError::InvalidRequest(message) => Self {
                code: "invalid_request".to_string(),
                message,
            },
            NarsilError::Schema(message) => Self {
                code: "schema_error".to_string(),
                message,
            },
            NarsilError::Backend(message) => Self {
                code: "backend_error".to_string(),
                message,
            },
            NarsilError::Transport(message) => Self {
                code: "transport_error".to_string(),
                message,
            },
            NarsilError::Overloaded(message) => Self {
                code: "overloaded".to_string(),
                message,
            },
        }
    }
}

struct HttpError(NarsilError);

impl IntoResponse for HttpError {
    fn into_response(self) -> HttpResponse {
        let status = http_status(&self.0);
        (status, Json(JsonError::from(self.0))).into_response()
    }
}

impl From<NarsilError> for HttpError {
    fn from(value: NarsilError) -> Self {
        Self(value)
    }
}

pub fn http_router<B: InferenceBackend + Clone>(service: SharedInferenceService<B>) -> Router {
    // Single-port convenience: mount `/metrics` on the gateway too.
    let observability = ObservabilityState::new(&service)
        .expect("failed to initialize Narsil observability")
        .with_metrics_on_gateway();
    http_router_with_observability(service, observability)
}

/// Inference HTTP gateway router. Serves the JSON inference endpoints plus the health
/// probes. Note: `/metrics` is intentionally *not* mounted here — metrics are scrape-only
/// on the (typically privately bound) observability listener. See
/// [`observability_router`]. When `expose_metrics` is `true`, `/metrics` is also mounted
/// here for the convenience of single-port test/dev setups (e.g. [`http_router`]).
pub fn http_router_with_observability<B: InferenceBackend + Clone>(
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
) -> Router {
    use tower_http::timeout::TimeoutLayer;

    let mut router = Router::new()
        .route("/v1/infer", post(http_infer::<B>))
        .route("/v1/infer_batch", post(http_infer_batch::<B>))
        .route("/health/live", get(http_liveness::<B>))
        .route("/health/ready", get(http_readiness::<B>));
    // Convenience: when no dedicated observability listener is configured, expose
    // `/metrics` on the inference gateway as well so single-port setups can still
    // scrape. Prefer the private observability listener in production (see README).
    if observability.expose_metrics_on_gateway {
        router = router.route("/metrics", get(http_metrics::<B>));
    }
    router
        // Explicit body cap: reject oversized payloads up front (413) rather than buffering
        // unbounded base64 tensor data.
        .layer(DefaultBodyLimit::max(DEFAULT_BODY_LIMIT_BYTES))
        // Per-request timeout so a stuck backend surfaces as 503 instead of a hung socket.
        .layer(TimeoutLayer::with_status_code(
            StatusCode::SERVICE_UNAVAILABLE,
            DEFAULT_REQUEST_TIMEOUT,
        ))
        .with_state(HttpState {
            service,
            observability,
        })
}

/// Dedicated observability listener router. This is the canonical place `/metrics` is
/// served; bind it privately (e.g. `127.0.0.1` or a sidecar-only interface). It also
/// serves the health probes for scrapers/probes pointed at the observability port.
pub fn observability_router(observability: ObservabilityState) -> Router {
    Router::new()
        .route("/metrics", get(observability_metrics))
        .route("/health/live", get(observability_liveness))
        .route("/health/ready", get(observability_readiness))
        .with_state(observability)
}

async fn http_infer<B: InferenceBackend + Clone>(
    State(state): State<HttpState<B>>,
    Json(request): Json<JsonPredictRequest>,
) -> std::result::Result<Json<JsonPredictResponse>, HttpError> {
    let request = InferenceRequest::try_from(request)?;
    let response = instrumented_infer(&state.service, "http", "infer", request).await?;
    Ok(Json(JsonPredictResponse::from(response)))
}

async fn http_infer_batch<B: InferenceBackend + Clone>(
    State(state): State<HttpState<B>>,
    Json(request): Json<JsonInferBatchRequest>,
) -> std::result::Result<Json<JsonInferBatchResponse>, HttpError> {
    if request.requests.is_empty() {
        return Err(HttpError(NarsilError::InvalidRequest(
            "infer_batch request must contain at least one request".to_string(),
        )));
    }

    // Reject batches larger than the configured grouping limit up front (4xx) so a single
    // request can't exceed what the collector will ever group. Only enforced in batching
    // mode; other executors have no fixed batch ceiling.
    if let Some(max) = state.service.max_batch_size()
        && request.requests.len() > max
    {
        return Err(HttpError(NarsilError::InvalidRequest(format!(
            "infer_batch request has {} items, exceeding the configured max_batch_size of {max}",
            request.requests.len()
        ))));
    }

    // Decode/validate each item independently so a single malformed request becomes a
    // per-item error in the response rather than failing the whole batch. Original indices
    // and ordering are preserved; only the valid items are dispatched to the backend.
    let mut responses: Vec<Option<JsonBatchResponse>> =
        (0..request.requests.len()).map(|_| None).collect();
    let mut valid_requests = Vec::new();
    let mut indices = Vec::new();

    for (i, req) in request.requests.into_iter().enumerate() {
        match InferenceRequest::try_from(req) {
            Ok(parsed) => {
                valid_requests.push(parsed);
                indices.push(i);
            }
            Err(err) => {
                responses[i] = Some(JsonBatchResponse::Error {
                    error: JsonError::from(err),
                });
            }
        }
    }

    if !valid_requests.is_empty() {
        let batch_responses =
            instrumented_infer_batch(&state.service, "http", "infer_batch", valid_requests).await;
        for (idx, res) in indices.into_iter().zip(batch_responses) {
            responses[idx] = Some(match res {
                Ok(resp) => JsonBatchResponse::Ok {
                    response: JsonPredictResponse::from(resp),
                },
                Err(err) => JsonBatchResponse::Error {
                    error: JsonError::from(err),
                },
            });
        }
    }

    // Every slot is filled: valid items by the dispatch loop, invalid items above.
    let responses = responses.into_iter().flatten().collect();
    Ok(Json(JsonInferBatchResponse { responses }))
}

async fn http_metrics<B: InferenceBackend + Clone>(
    State(state): State<HttpState<B>>,
) -> HttpResponse {
    metrics_response(&state.observability)
}

async fn http_liveness<B: InferenceBackend + Clone>(
    State(state): State<HttpState<B>>,
) -> StatusCode {
    liveness_status(&state.observability)
}

async fn http_readiness<B: InferenceBackend + Clone>(
    State(state): State<HttpState<B>>,
) -> StatusCode {
    readiness_status(&state.observability)
}

async fn observability_metrics(State(observability): State<ObservabilityState>) -> HttpResponse {
    metrics_response(&observability)
}

async fn observability_liveness(State(observability): State<ObservabilityState>) -> StatusCode {
    liveness_status(&observability)
}

async fn observability_readiness(State(observability): State<ObservabilityState>) -> StatusCode {
    readiness_status(&observability)
}

fn metrics_response(observability: &ObservabilityState) -> HttpResponse {
    (
        [(header::CONTENT_TYPE, PROMETHEUS_CONTENT_TYPE)],
        observability.render_metrics(),
    )
        .into_response()
}

fn liveness_status(observability: &ObservabilityState) -> StatusCode {
    if observability.is_live() {
        StatusCode::NO_CONTENT
    } else {
        StatusCode::SERVICE_UNAVAILABLE
    }
}

fn readiness_status(observability: &ObservabilityState) -> StatusCode {
    if observability.is_ready() {
        StatusCode::NO_CONTENT
    } else {
        StatusCode::SERVICE_UNAVAILABLE
    }
}

fn http_status(error: &NarsilError) -> StatusCode {
    match error {
        NarsilError::InvalidRequest(_) | NarsilError::Schema(_) => StatusCode::BAD_REQUEST,
        NarsilError::Backend(_) => StatusCode::INTERNAL_SERVER_ERROR,
        NarsilError::Transport(_) => StatusCode::SERVICE_UNAVAILABLE,
        NarsilError::Overloaded(_) => StatusCode::TOO_MANY_REQUESTS,
    }
}

async fn shutdown_signal() {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{SignalKind, signal};
        let mut sigterm =
            signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {},
            _ = sigterm.recv() => {},
        }
    }
    #[cfg(not(unix))]
    {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    }
    tracing::info!("shutdown signal received, draining in-flight requests");
}

async fn wait_for_shutdown(mut shutdown: watch::Receiver<bool>) {
    if *shutdown.borrow() {
        return;
    }

    while shutdown.changed().await.is_ok() {
        if *shutdown.borrow() {
            return;
        }
    }
}

fn env_socket_addr(name: &str) -> Result<Option<SocketAddr>> {
    match std::env::var(name) {
        Ok(value) if value.trim().is_empty() => Ok(None),
        Ok(value) => value
            .parse()
            .map(Some)
            .map_err(|err| NarsilError::InvalidRequest(format!("invalid {name} '{value}': {err}"))),
        Err(std::env::VarError::NotPresent) => Ok(None),
        Err(err) => Err(NarsilError::InvalidRequest(format!(
            "invalid {name}: {err}"
        ))),
    }
}

async fn serve_grpc<B, S>(
    addr: SocketAddr,
    service: SharedInferenceService<B>,
    shutdown: S,
) -> Result<()>
where
    B: InferenceBackend + Clone,
    S: Future<Output = ()> + Send + 'static,
{
    let (health_reporter, health_service) = health_reporter();
    health_reporter
        .set_serving::<InferenceServer<GrpcInferenceService<B>>>()
        .await;
    health_reporter
        .set_service_status("", ServingStatus::Serving)
        .await;

    tonic::transport::Server::builder()
        .add_service(health_service)
        .add_service(GrpcInferenceService::from_shared(service).into_server())
        .serve_with_shutdown(addr, shutdown)
        .await
        .map_err(|err| NarsilError::Transport(format!("gRPC server failed: {err}")))
}

async fn serve_http<B, S>(
    addr: SocketAddr,
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
    shutdown: S,
) -> Result<()>
where
    B: InferenceBackend + Clone,
    S: Future<Output = ()> + Send + 'static,
{
    let listener = tokio::net::TcpListener::bind(addr)
        .await
        .map_err(|err| NarsilError::Transport(format!("HTTP bind failed on {addr}: {err}")))?;
    tracing::info!(%addr, "starting Narsil HTTP/JSON gateway");
    axum::serve(
        listener,
        http_router_with_observability(service, observability),
    )
    .with_graceful_shutdown(shutdown)
    .await
    .map_err(|err| NarsilError::Transport(format!("HTTP server failed: {err}")))
}

async fn serve_observability_http<S>(
    addr: SocketAddr,
    observability: ObservabilityState,
    shutdown: S,
) -> Result<()>
where
    S: Future<Output = ()> + Send + 'static,
{
    let listener = tokio::net::TcpListener::bind(addr).await.map_err(|err| {
        NarsilError::Transport(format!("observability HTTP bind failed on {addr}: {err}"))
    })?;
    tracing::info!(%addr, "starting Narsil observability endpoint");
    axum::serve(listener, observability_router(observability))
        .with_graceful_shutdown(shutdown)
        .await
        .map_err(|err| NarsilError::Transport(format!("observability server failed: {err}")))
}

async fn serve_with_optional_http<B>(
    addr: SocketAddr,
    service: SharedInferenceService<B>,
) -> Result<()>
where
    B: InferenceBackend + Clone,
{
    let http_addr = env_socket_addr("NARSIL_HTTP_ADDR")?;
    let observability_addr = env_socket_addr("NARSIL_OBSERVABILITY_ADDR")?;
    let quic_addr = env_socket_addr("NARSIL_QUIC_ADDR")?;

    // Fail fast on a feature/runtime mismatch before binding any listener: a build
    // without `quic` cannot honor NARSIL_QUIC_ADDR, so surface it up front rather than
    // after gRPC/HTTP have already taken their ports.
    #[cfg(not(feature = "quic"))]
    if quic_addr.is_some() {
        return Err(NarsilError::InvalidRequest(
            "NARSIL_QUIC_ADDR requires building with --features quic".to_string(),
        ));
    }

    let (shutdown_tx, _) = watch::channel(false);

    // Only install the Prometheus recorder (and build `ObservabilityState`) when a
    // listener will actually render it. A gRPC-only deploy keeps collecting metrics via
    // the no-op recorder, but there is nowhere to scrape them — installing the global
    // recorder anyway can fail if the host binary already installed one (see review).
    let observability =
        if http_addr.is_some() || observability_addr.is_some() || quic_addr.is_some() {
            let observability = ObservabilityState::new(&service)?;
            // Drain readiness when shutdown is signalled so load balancers stop routing.
            let signal_tx = shutdown_tx.clone();
            let shutdown_observability = observability.clone();
            tokio::spawn(async move {
                shutdown_signal().await;
                shutdown_observability.set_draining(true);
                let _ = signal_tx.send(true);
            });
            Some(observability)
        } else {
            tracing::warn!(
                "metrics are being collected but no HTTP endpoint is configured to expose them; \
             set NARSIL_OBSERVABILITY_ADDR (recommended, private bind) to scrape /metrics, \
             or NARSIL_HTTP_ADDR for the JSON gateway + health probes"
            );
            let signal_tx = shutdown_tx.clone();
            tokio::spawn(async move {
                shutdown_signal().await;
                let _ = signal_tx.send(true);
            });
            None
        };

    // No dedicated observability listener means `/metrics` is exposed on the public
    // inference gateway as a single-port convenience (bind the gateway privately if
    // scraping there); a dedicated listener keeps `/metrics` off the gateway.
    let expose_metrics_on_gateway = http_addr.is_some() && observability_addr.is_none();
    if expose_metrics_on_gateway {
        tracing::warn!(
            "exposing /metrics on the public NARSIL_HTTP_ADDR gateway because no \
             NARSIL_OBSERVABILITY_ADDR is set; prefer a private observability listener"
        );
    }

    let mut servers = JoinSet::new();
    servers.spawn(serve_grpc(
        addr,
        service.clone(),
        wait_for_shutdown(shutdown_tx.subscribe()),
    ));

    if let Some(http_addr) = http_addr {
        let mut http_observability = observability
            .clone()
            .expect("observability is built whenever an HTTP listener is configured");
        if expose_metrics_on_gateway {
            http_observability = http_observability.with_metrics_on_gateway();
        }
        servers.spawn(serve_http(
            http_addr,
            service.clone(),
            http_observability,
            wait_for_shutdown(shutdown_tx.subscribe()),
        ));
    }

    if let Some(observability_addr) = observability_addr {
        servers.spawn(serve_observability_http(
            observability_addr,
            observability
                .clone()
                .expect("observability is built whenever an observability listener is configured"),
            wait_for_shutdown(shutdown_tx.subscribe()),
        ));
    }

    #[cfg(feature = "quic")]
    if let Some(quic_addr) = quic_addr {
        servers.spawn(serve_quic_http3(
            quic_addr,
            service,
            observability.expect("observability is built whenever a QUIC listener is configured"),
            wait_for_shutdown(shutdown_tx.subscribe()),
        ));
    }

    while let Some(result) = servers.join_next().await {
        // Any server exiting — success or failure — means we're shutting down: signal the
        // rest so we never linger with a partially-running set of transports.
        let _ = shutdown_tx.send(true);
        match result {
            Ok(Ok(())) => {}
            Ok(Err(err)) => {
                return Err(err);
            }
            Err(err) => {
                return Err(NarsilError::Transport(format!("server task failed: {err}")));
            }
        }
    }

    Ok(())
}

#[cfg(feature = "quic")]
pub async fn serve_quic<B, S>(
    addr: SocketAddr,
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
    shutdown: S,
) -> Result<()>
where
    B: InferenceBackend + Clone,
    S: Future<Output = ()> + Send + 'static,
{
    serve_quic_http3(addr, service, observability, shutdown).await
}

#[cfg(feature = "quic")]
async fn serve_quic_http3<B, S>(
    addr: SocketAddr,
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
    shutdown: S,
) -> Result<()>
where
    B: InferenceBackend + Clone,
    S: Future<Output = ()> + Send + 'static,
{
    let endpoint = quic_server_endpoint(addr)?;
    tracing::info!(%addr, "starting Narsil HTTP/3 QUIC transport");
    tokio::pin!(shutdown);

    let mut connections = JoinSet::new();

    loop {
        tokio::select! {
            _ = &mut shutdown => {
                break;
            }
            incoming = endpoint.accept() => {
                let Some(incoming) = incoming else {
                    break;
                };
                let service = service.clone();
                let observability = observability.clone();
                connections.spawn(async move {
                    match incoming.await {
                        Ok(connection) => {
                            if let Err(err) = handle_quic_connection(connection, service, observability).await {
                                tracing::warn!(error = %err, "HTTP/3 connection failed");
                            }
                        }
                        Err(err) => {
                            tracing::warn!(error = %err, "QUIC handshake failed");
                        }
                    }
                });
            }
            // Reap finished connection tasks so the JoinSet doesn't grow unbounded over
            // the lifetime of a long-running server.
            _ = connections.join_next(), if !connections.is_empty() => {}
        }
    }

    // Graceful drain: stop accepting new connections, then give in-flight HTTP/3
    // connections (and their multiplexed request streams) a bounded window to finish —
    // matching the gRPC/HTTP graceful-shutdown behavior — before tearing the endpoint
    // down. After the window we force-close so a stuck client can't block shutdown.
    const QUIC_DRAIN_TIMEOUT: Duration = Duration::from_secs(30);
    let drain = async { while connections.join_next().await.is_some() {} };
    if tokio::time::timeout(QUIC_DRAIN_TIMEOUT, drain)
        .await
        .is_err()
    {
        tracing::warn!(
            timeout_secs = QUIC_DRAIN_TIMEOUT.as_secs(),
            "QUIC drain timed out; force-closing remaining HTTP/3 connections"
        );
    }
    endpoint.close(0u32.into(), b"shutdown");
    endpoint.wait_idle().await;
    Ok(())
}

#[cfg(feature = "quic")]
const QUIC_TLS_CERT_ENV: &str = "NARSIL_QUIC_TLS_CERT";
#[cfg(feature = "quic")]
const QUIC_TLS_KEY_ENV: &str = "NARSIL_QUIC_TLS_KEY";

#[cfg(feature = "quic")]
fn quic_server_endpoint(addr: SocketAddr) -> Result<h3_quinn::quinn::Endpoint> {
    use quinn::crypto::rustls::QuicServerConfig;

    let (certs, key) = load_quic_tls_material()?;

    let mut tls_config = rustls::ServerConfig::builder_with_provider(Arc::new(
        rustls::crypto::aws_lc_rs::default_provider(),
    ))
    .with_protocol_versions(&[&rustls::version::TLS13])
    .map_err(|err| NarsilError::Transport(format!("failed to select QUIC TLS version: {err}")))?
    .with_no_client_auth()
    .with_single_cert(certs, key)
    .map_err(|err| NarsilError::Transport(format!("failed to build QUIC TLS config: {err}")))?;
    tls_config.alpn_protocols = vec![b"h3".to_vec()];

    let server_config = h3_quinn::quinn::ServerConfig::with_crypto(Arc::new(
        QuicServerConfig::try_from(tls_config)
            .map_err(|err| NarsilError::Transport(format!("failed to build QUIC config: {err}")))?,
    ));
    h3_quinn::quinn::Endpoint::server(server_config, addr)
        .map_err(|err| NarsilError::Transport(format!("QUIC bind failed on {addr}: {err}")))
}

/// Resolve the QUIC TLS certificate chain and private key.
///
/// When both `NARSIL_QUIC_TLS_CERT` and `NARSIL_QUIC_TLS_KEY` point at PEM files we load
/// a real, operator-supplied certificate. Supplying only one of the pair is treated as a
/// misconfiguration and rejected. With neither set we fall back to an ephemeral,
/// self-signed certificate **for local development only** and log a prominent warning —
/// such a certificate cannot be verified by any client and must never be relied on in
/// production.
#[cfg(feature = "quic")]
#[allow(clippy::type_complexity)]
fn load_quic_tls_material() -> Result<(
    Vec<rustls::pki_types::CertificateDer<'static>>,
    rustls::pki_types::PrivateKeyDer<'static>,
)> {
    use rustls::pki_types::PrivateKeyDer;

    let cert_path = std::env::var_os(QUIC_TLS_CERT_ENV).filter(|v| !v.is_empty());
    let key_path = std::env::var_os(QUIC_TLS_KEY_ENV).filter(|v| !v.is_empty());

    match (cert_path, key_path) {
        (Some(cert_path), Some(key_path)) => {
            let certs = load_pem_certs(cert_path.as_ref())?;
            let key = load_pem_private_key(key_path.as_ref())?;
            tracing::info!(
                cert = %std::path::Path::new(&cert_path).display(),
                "loaded operator-supplied QUIC TLS certificate"
            );
            Ok((certs, key))
        }
        (Some(_), None) => Err(NarsilError::InvalidRequest(format!(
            "{QUIC_TLS_CERT_ENV} is set but {QUIC_TLS_KEY_ENV} is not; both are required to use a real certificate"
        ))),
        (None, Some(_)) => Err(NarsilError::InvalidRequest(format!(
            "{QUIC_TLS_KEY_ENV} is set but {QUIC_TLS_CERT_ENV} is not; both are required to use a real certificate"
        ))),
        (None, None) => {
            tracing::warn!(
                "no QUIC TLS certificate configured ({QUIC_TLS_CERT_ENV}/{QUIC_TLS_KEY_ENV} unset); \
                 generating an ephemeral self-signed certificate for DEVELOPMENT ONLY — clients \
                 cannot verify it and it MUST NOT be used in production"
            );
            let certified_key = rcgen::generate_simple_self_signed(vec!["localhost".to_string()])
                .map_err(|err| {
                NarsilError::Transport(format!("failed to generate QUIC cert: {err}"))
            })?;
            let cert = certified_key.cert.into();
            let key = PrivateKeyDer::Pkcs8(certified_key.signing_key.serialize_der().into());
            Ok((vec![cert], key))
        }
    }
}

#[cfg(feature = "quic")]
fn load_pem_certs(
    path: &std::ffi::OsStr,
) -> Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
    let display = std::path::Path::new(path).display();
    let pem = std::fs::read(path).map_err(|err| {
        NarsilError::Transport(format!(
            "failed to read QUIC TLS certificate '{display}': {err}"
        ))
    })?;
    let certs = rustls_pemfile::certs(&mut pem.as_slice())
        .collect::<std::result::Result<Vec<_>, _>>()
        .map_err(|err| {
            NarsilError::Transport(format!(
                "failed to parse QUIC TLS certificate '{display}': {err}"
            ))
        })?;
    if certs.is_empty() {
        return Err(NarsilError::InvalidRequest(format!(
            "QUIC TLS certificate '{display}' contains no certificates"
        )));
    }
    Ok(certs)
}

#[cfg(feature = "quic")]
fn load_pem_private_key(
    path: &std::ffi::OsStr,
) -> Result<rustls::pki_types::PrivateKeyDer<'static>> {
    let display = std::path::Path::new(path).display();
    let pem = std::fs::read(path).map_err(|err| {
        NarsilError::Transport(format!("failed to read QUIC TLS key '{display}': {err}"))
    })?;
    rustls_pemfile::private_key(&mut pem.as_slice())
        .map_err(|err| {
            NarsilError::Transport(format!("failed to parse QUIC TLS key '{display}': {err}"))
        })?
        .ok_or_else(|| {
            NarsilError::InvalidRequest(format!("QUIC TLS key '{display}' contains no private key"))
        })
}

#[cfg(feature = "quic")]
async fn handle_quic_connection<B>(
    connection: h3_quinn::quinn::Connection,
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
) -> Result<()>
where
    B: InferenceBackend + Clone,
{
    let connection = h3_quinn::Connection::new(connection);
    let mut connection = h3::server::Connection::new(connection)
        .await
        .map_err(|err| NarsilError::Transport(format!("HTTP/3 connection failed: {err}")))?;

    // Each accepted request stream is handled on its own task so that a slow upload or a
    // long-running inference on one stream does not head-of-line block the other streams
    // multiplexed over the same QUIC connection — preserving HTTP/3 concurrency.
    let mut streams = JoinSet::new();

    loop {
        tokio::select! {
            resolver = connection.accept() => {
                match resolver
                    .map_err(|err| NarsilError::Transport(format!("HTTP/3 accept failed: {err}")))?
                {
                    Some(resolver) => {
                        let service = service.clone();
                        let observability = observability.clone();
                        streams.spawn(async move {
                            let (request, stream) = match resolver.resolve_request().await {
                                Ok(parts) => parts,
                                Err(err) => {
                                    tracing::warn!(error = %err, "HTTP/3 resolve request failed");
                                    return;
                                }
                            };
                            if let Err(err) =
                                handle_quic_request(request, stream, service, observability).await
                            {
                                tracing::warn!(error = %err, "HTTP/3 request failed");
                            }
                        });
                    }
                    None => break,
                }
            }
            _ = streams.join_next(), if !streams.is_empty() => {}
        }
    }

    // Let already-accepted request streams finish before the connection task returns.
    while streams.join_next().await.is_some() {}

    Ok(())
}

/// Maximum HTTP/3 request body Narsil will buffer in memory. QUIC request bodies are
/// fully buffered before dispatch, so without a cap a single client could stream an
/// unbounded body and OOM the server (a DoS vector). 64 MiB comfortably covers realistic
/// tensor batches while bounding per-request memory.
#[cfg(feature = "quic")]
const MAX_QUIC_BODY_SIZE: usize = 64 * 1024 * 1024;

#[cfg(feature = "quic")]
async fn handle_quic_request<B, S>(
    request: axum::http::Request<()>,
    mut stream: h3::server::RequestStream<S, bytes::Bytes>,
    service: SharedInferenceService<B>,
    observability: ObservabilityState,
) -> std::result::Result<(), h3::error::StreamError>
where
    B: InferenceBackend + Clone,
    S: h3::quic::BidiStream<bytes::Bytes>,
{
    use bytes::{Buf, Bytes};

    let method = request.method().clone();
    let path = request.uri().path().to_string();
    let mut body = Vec::new();
    let mut limit_exceeded = false;
    while let Some(mut chunk) = stream.recv_data().await? {
        if body.len() + chunk.remaining() > MAX_QUIC_BODY_SIZE {
            // Stop draining once the cap is hit so we never materialize the oversized body.
            limit_exceeded = true;
            break;
        }
        while chunk.has_remaining() {
            let bytes = chunk.copy_to_bytes(chunk.remaining());
            body.extend_from_slice(&bytes);
        }
    }

    let response = if limit_exceeded {
        quic_json_response(
            StatusCode::PAYLOAD_TOO_LARGE,
            &JsonError {
                code: "payload_too_large".to_string(),
                message: format!("request body exceeds maximum size of {MAX_QUIC_BODY_SIZE} bytes"),
            },
        )
    } else {
        dispatch_quic_request(method, &path, body, &service, &observability).await
    };
    stream
        .send_response(
            axum::http::Response::builder()
                .status(response.status)
                .header(header::CONTENT_TYPE, response.content_type)
                .body(())
                .expect("valid HTTP/3 response"),
        )
        .await?;
    if !response.body.is_empty() {
        stream.send_data(Bytes::from(response.body)).await?;
    }
    stream.finish().await
}

#[cfg(feature = "quic")]
struct QuicHttpResponse {
    status: StatusCode,
    content_type: &'static str,
    body: Vec<u8>,
}

#[cfg(feature = "quic")]
async fn dispatch_quic_request<B>(
    method: axum::http::Method,
    path: &str,
    body: Vec<u8>,
    service: &SharedInferenceService<B>,
    observability: &ObservabilityState,
) -> QuicHttpResponse
where
    B: InferenceBackend + Clone,
{
    match (method, path) {
        (axum::http::Method::POST, "/v1/infer") => quic_infer(body, service).await,
        (axum::http::Method::POST, "/v1/infer_batch") => quic_infer_batch(body, service).await,
        (axum::http::Method::GET, "/metrics") => QuicHttpResponse {
            status: StatusCode::OK,
            content_type: PROMETHEUS_CONTENT_TYPE,
            body: observability.render_metrics().into_bytes(),
        },
        (axum::http::Method::GET, "/health/live") => QuicHttpResponse {
            status: StatusCode::NO_CONTENT,
            content_type: "text/plain; charset=utf-8",
            body: Vec::new(),
        },
        (axum::http::Method::GET, "/health/ready") => QuicHttpResponse {
            status: readiness_status(observability),
            content_type: "text/plain; charset=utf-8",
            body: Vec::new(),
        },
        _ => quic_json_response(
            StatusCode::NOT_FOUND,
            &JsonError {
                code: "not_found".to_string(),
                message: "unknown HTTP/3 endpoint".to_string(),
            },
        ),
    }
}

#[cfg(feature = "quic")]
async fn quic_infer<B>(body: Vec<u8>, service: &SharedInferenceService<B>) -> QuicHttpResponse
where
    B: InferenceBackend + Clone,
{
    let request = match serde_json::from_slice::<JsonPredictRequest>(&body)
        .map_err(|err| NarsilError::InvalidRequest(format!("invalid JSON request: {err}")))
        .and_then(InferenceRequest::try_from)
    {
        Ok(request) => request,
        Err(err) => return quic_error_response(err),
    };

    match instrumented_infer(service, "quic", "infer", request).await {
        Ok(response) => quic_json_response(StatusCode::OK, &JsonPredictResponse::from(response)),
        Err(err) => quic_error_response(err),
    }
}

#[cfg(feature = "quic")]
async fn quic_infer_batch<B>(body: Vec<u8>, service: &SharedInferenceService<B>) -> QuicHttpResponse
where
    B: InferenceBackend + Clone,
{
    let request = match serde_json::from_slice::<JsonInferBatchRequest>(&body)
        .map_err(|err| NarsilError::InvalidRequest(format!("invalid JSON request: {err}")))
    {
        Ok(request) => request,
        Err(err) => return quic_error_response(err),
    };
    if request.requests.is_empty() {
        return quic_error_response(NarsilError::InvalidRequest(
            "infer_batch request must contain at least one request".to_string(),
        ));
    }
    let requests = match request
        .requests
        .into_iter()
        .map(InferenceRequest::try_from)
        .collect::<Result<Vec<_>>>()
    {
        Ok(requests) => requests,
        Err(err) => return quic_error_response(err),
    };

    let responses = instrumented_infer_batch(service, "quic", "infer_batch", requests)
        .await
        .into_iter()
        .map(|response| match response {
            Ok(response) => JsonBatchResponse::Ok {
                response: JsonPredictResponse::from(response),
            },
            Err(error) => JsonBatchResponse::Error {
                error: JsonError::from(error),
            },
        })
        .collect();

    quic_json_response(StatusCode::OK, &JsonInferBatchResponse { responses })
}

#[cfg(feature = "quic")]
fn quic_error_response(error: NarsilError) -> QuicHttpResponse {
    let status = http_status(&error);
    quic_json_response(status, &JsonError::from(error))
}

#[cfg(feature = "quic")]
fn quic_json_response<T: Serialize>(status: StatusCode, value: &T) -> QuicHttpResponse {
    match serde_json::to_vec(value) {
        Ok(body) => QuicHttpResponse {
            status,
            content_type: "application/json",
            body,
        },
        Err(err) => QuicHttpResponse {
            status: StatusCode::INTERNAL_SERVER_ERROR,
            content_type: "application/json",
            body: format!(
                r#"{{"code":"serialization_error","message":"failed to encode JSON response: {err}"}}"#
            )
            .into_bytes(),
        },
    }
}

pub async fn serve<B: InferenceBackend + Clone>(addr: SocketAddr, backend: B) -> Result<()> {
    serve_with_optional_http(addr, SharedInferenceService::new(backend)).await
}

pub async fn serve_with_worker_pool<B: InferenceBackend + Clone>(
    addr: SocketAddr,
    backend: B,
    config: WorkerPoolConfig,
) -> Result<()> {
    serve_with_optional_http(
        addr,
        SharedInferenceService::with_worker_pool(backend, config)?,
    )
    .await
}

pub async fn serve_with_batching<B: InferenceBackend + Clone>(
    addr: SocketAddr,
    backend: B,
    config: BatchingConfig,
) -> Result<()> {
    serve_with_optional_http(
        addr,
        SharedInferenceService::with_batching(backend, config)?,
    )
    .await
}
