knafeh 1.1.0

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use boring::ssl::{SslContextBuilder, SslMethod, SslVerifyMode};
use tokio::sync::{oneshot, Mutex as TokioMutex};

use tokio_quiche::http3::driver::{
    ClientH3Driver, ClientH3Event, ClientRequestSender, H3Event, InboundFrame, IncomingH3Headers,
};
use tokio_quiche::http3::settings::Http3Settings;
use tokio_quiche::settings::{CertificateKind, ConnectionParams, Hooks, TlsCertificatePaths};
use tokio_quiche::ClientH3Controller;
use tokio_quiche::QuicConnection;

use crate::error::KnafehError;
use crate::transport::quic_wire::MAX_MESSAGE_SIZE;
use crate::transport::tls::TlsConfig;

/// Raw HTTP/3 response (headers + body) received from the server.
pub(crate) enum H3Response {
    /// Unary response — all body frames buffered into a single `Vec<u8>`.
    Complete {
        headers: Vec<quiche::h3::Header>,
        body: Vec<u8>,
    },
    /// Streaming response — headers delivered, body frames available via `recv`.
    Streaming {
        headers: Vec<quiche::h3::Header>,
        recv: tokio::sync::mpsc::Receiver<InboundFrame>,
        read_fin: bool,
    },
}

/// Shared state for demuxing responses to pending requests.
struct ConnectionState {
    /// Pending unary requests: request_id → oneshot sender.
    pending: HashMap<u64, oneshot::Sender<Result<H3Response, KnafehError>>>,
    /// Pending streaming requests: request_id → oneshot sender that delivers
    /// the full `IncomingH3Headers` so the caller can read frames.
    pending_stream: HashMap<u64, oneshot::Sender<Result<H3Response, KnafehError>>>,
    /// stream_id → request_id (populated by NewOutboundRequest event)
    stream_to_request: HashMap<u64, u64>,
}

/// Shared connection resources accessible from both the pool and call sites.
pub(crate) struct ConnectionInner {
    pub(crate) request_sender: ClientRequestSender,
    state: Arc<TokioMutex<ConnectionState>>,
    pub(crate) next_request_id: AtomicU64,
}

impl ConnectionInner {
    /// Register a pending unary request.
    pub(crate) async fn register_pending(
        &self,
        request_id: u64,
        tx: oneshot::Sender<Result<H3Response, KnafehError>>,
    ) {
        self.state.lock().await.pending.insert(request_id, tx);
    }

    /// Register a pending streaming request.
    pub(crate) async fn register_pending_stream(
        &self,
        request_id: u64,
        tx: oneshot::Sender<Result<H3Response, KnafehError>>,
    ) {
        self.state
            .lock()
            .await
            .pending_stream
            .insert(request_id, tx);
    }

    /// Remove a pending unary request (cleanup on error).
    pub(crate) async fn remove_pending(&self, request_id: u64) {
        self.state.lock().await.pending.remove(&request_id);
    }

    /// Remove a pending streaming request (cleanup on error).
    pub(crate) async fn remove_pending_stream(&self, request_id: u64) {
        self.state.lock().await.pending_stream.remove(&request_id);
    }
}

/// Internal pool entry tracking stream usage.
struct PoolEntry {
    id: u64,
    active_streams: usize,
    max_streams: usize,
    inner: Arc<ConnectionInner>,
}

/// Opaque handle to a pooled QUIC/HTTP3 connection.
pub struct ConnectionHandle {
    /// Unique connection ID.
    pub id: u64,
    /// Shared connection resources.
    pub(crate) inner: Arc<ConnectionInner>,
}

/// RAII guard that decrements `active_streams` when dropped, even on
/// cancellation or panic.
pub(crate) struct ConnectionGuard {
    pool: Arc<ClientConnectionPool>,
    conn_id: u64,
}

impl ConnectionGuard {
    pub(crate) fn new(pool: Arc<ClientConnectionPool>, conn_id: u64) -> Self {
        Self { pool, conn_id }
    }

    /// Consume the guard without releasing the connection.
    /// Used when ownership of the release is transferred elsewhere
    /// (e.g., a streaming background task).
    pub(crate) fn detach(self) -> (Arc<ClientConnectionPool>, u64) {
        let pool = Arc::clone(&self.pool);
        let id = self.conn_id;
        std::mem::forget(self);
        (pool, id)
    }
}

impl Drop for ConnectionGuard {
    fn drop(&mut self) {
        self.pool.release(self.conn_id);
    }
}

/// A pool of QUIC connections to a single endpoint.
///
/// Connections are lazily established and reused across RPC calls.
/// QUIC's native multiplexing means a single connection can handle
/// many concurrent streams, so the pool size can typically be small.
pub struct ClientConnectionPool {
    /// Maximum number of connections in the pool.
    max_size: usize,
    /// The remote endpoint address.
    endpoint: String,
    /// Hostname portion of the endpoint, used for TLS SNI.
    hostname: String,
    /// TLS configuration.
    tls_config: TlsConfig,
    /// Active connections (sync Mutex — never held across `.await`).
    connections: Mutex<Vec<PoolEntry>>,
}

impl ClientConnectionPool {
    pub fn new(endpoint: String, max_size: usize, tls_config: TlsConfig) -> Self {
        let hostname = endpoint
            .rsplit_once(':')
            .map(|(h, _)| h)
            .unwrap_or(&endpoint)
            .trim_start_matches('[')
            .trim_end_matches(']')
            .to_string();

        Self {
            max_size,
            endpoint,
            hostname,
            tls_config,
            connections: Mutex::new(Vec::new()),
        }
    }

    /// Acquire a connection from the pool.
    ///
    /// Returns an existing connection with capacity, or creates a new one.
    pub async fn acquire(&self) -> Result<ConnectionHandle, KnafehError> {
        let should_create = {
            let mut conns = self.connections.lock().unwrap();

            if let Some(entry) = conns.iter_mut().find(|e| e.active_streams < e.max_streams) {
                entry.active_streams += 1;
                return Ok(ConnectionHandle {
                    id: entry.id,
                    inner: Arc::clone(&entry.inner),
                });
            }

            conns.len() < self.max_size
        };

        if !should_create {
            return Err(KnafehError::Transport(
                "connection pool exhausted".to_string(),
            ));
        }

        let (handle, entry, quic_conn, controller, state) = self.create_connection().await?;

        let mut conns = self.connections.lock().unwrap();
        if conns.len() >= self.max_size {
            if let Some(existing) = conns.iter_mut().find(|e| e.active_streams < e.max_streams) {
                existing.active_streams += 1;
                return Ok(ConnectionHandle {
                    id: existing.id,
                    inner: Arc::clone(&existing.inner),
                });
            }
            return Err(KnafehError::Transport(
                "connection pool exhausted".to_string(),
            ));
        }
        conns.push(entry);
        tokio::spawn(async move {
            connection_event_loop(quic_conn, controller, state).await;
        });

        Ok(handle)
    }

    /// Release a connection back to the pool (decrement active streams).
    pub fn release(&self, conn_id: u64) {
        let mut conns = self.connections.lock().unwrap();
        if let Some(entry) = conns.iter_mut().find(|e| e.id == conn_id) {
            entry.active_streams = entry.active_streams.saturating_sub(1);
        }
    }

    /// Create a new QUIC/HTTP3 connection to the endpoint.
    async fn create_connection(
        &self,
    ) -> Result<
        (
            ConnectionHandle,
            PoolEntry,
            QuicConnection,
            ClientH3Controller,
            Arc<TokioMutex<ConnectionState>>,
        ),
        KnafehError,
    > {
        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
        let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);

        tracing::info!(endpoint = %self.endpoint, id, "establishing new QUIC connection");

        let addr: SocketAddr = match self.endpoint.parse() {
            Ok(addr) => addr,
            Err(_) => tokio::net::lookup_host(&self.endpoint)
                .await
                .map_err(|e| {
                    KnafehError::Transport(format!(
                        "failed to resolve endpoint '{}': {e}",
                        self.endpoint
                    ))
                })?
                .next()
                .ok_or_else(|| {
                    KnafehError::Transport(format!("no addresses found for '{}'", self.endpoint))
                })?,
        };

        let bind_addr: SocketAddr = if addr.is_ipv4() {
            "0.0.0.0:0"
        } else {
            "[::]:0"
        }
        .parse()
        .unwrap();

        let socket = tokio::net::UdpSocket::bind(bind_addr)
            .await
            .map_err(|e| KnafehError::Transport(format!("failed to bind UDP socket: {e}")))?;
        socket
            .connect(addr)
            .await
            .map_err(|e| KnafehError::Transport(format!("failed to connect UDP socket: {e}")))?;

        let (h3_driver, controller) = ClientH3Driver::new(Http3Settings::default());
        let conn_params = self.connection_params()?;

        let quic_conn = tokio_quiche::quic::connect_with_config(
            socket,
            Some(&self.hostname),
            &conn_params,
            h3_driver,
        )
        .await
        .map_err(|e| KnafehError::Transport(format!("QUIC connect failed: {e}")))?;

        let request_sender = controller.request_sender();
        let state = Arc::new(TokioMutex::new(ConnectionState {
            pending: HashMap::new(),
            pending_stream: HashMap::new(),
            stream_to_request: HashMap::new(),
        }));

        let inner = Arc::new(ConnectionInner {
            request_sender,
            state: Arc::clone(&state),
            next_request_id: AtomicU64::new(0),
        });

        let handle = ConnectionHandle {
            id,
            inner: Arc::clone(&inner),
        };

        let entry = PoolEntry {
            id,
            active_streams: 1,
            max_streams: 100,
            inner,
        };

        Ok((handle, entry, quic_conn, controller, state))
    }

    fn connection_params(&self) -> Result<ConnectionParams<'static>, KnafehError> {
        let _ = build_client_tls_context(&self.tls_config)?;
        let settings = tokio_quiche::settings::QuicSettings {
            max_idle_timeout: Some(std::time::Duration::from_secs(30)),
            alpn: self.tls_config.alpn.clone(),
            ..Default::default()
        };

        let hooks = Hooks {
            connection_hook: Some(Arc::new(ClientTlsHook {
                tls_config: self.tls_config.clone(),
            })),
        };

        // tokio-quiche only invokes connection hooks when TLS paths are set.
        // The hook builds the complete SSL context, so these placeholders are
        // never read as files.
        let tls_cert = Some(TlsCertificatePaths {
            cert: "",
            private_key: "",
            kind: CertificateKind::X509,
        });

        Ok(ConnectionParams::new_client(settings, tls_cert, hooks))
    }

    /// Number of connections currently in the pool.
    pub fn size(&self) -> usize {
        self.connections.lock().unwrap().len()
    }

    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }

    pub fn hostname(&self) -> &str {
        &self.hostname
    }
}

struct ClientTlsHook {
    tls_config: TlsConfig,
}

impl tokio_quiche::quic::ConnectionHook for ClientTlsHook {
    fn create_custom_ssl_context_builder(
        &self,
        _settings: TlsCertificatePaths<'_>,
    ) -> Option<SslContextBuilder> {
        match build_client_tls_context(&self.tls_config) {
            Ok(builder) => Some(builder),
            Err(e) => {
                tracing::error!(error = %e, "failed to build client TLS context");
                None
            }
        }
    }
}

fn build_client_tls_context(tls_config: &TlsConfig) -> Result<SslContextBuilder, KnafehError> {
    let mut builder =
        SslContextBuilder::new(SslMethod::tls()).map_err(|e| KnafehError::Tls(e.to_string()))?;

    if tls_config.verify_peer {
        builder.set_verify(SslVerifyMode::PEER);
        if let Some(ca_path) = &tls_config.ca_path {
            builder
                .set_ca_file(ca_path)
                .map_err(|e| KnafehError::Tls(e.to_string()))?;
        } else {
            builder
                .set_default_verify_paths()
                .map_err(|e| KnafehError::Tls(e.to_string()))?;
        }
    } else {
        builder.set_verify(SslVerifyMode::NONE);
    }

    Ok(builder)
}

/// Background event loop for a single QUIC/HTTP3 connection.
async fn connection_event_loop(
    _quic_conn: QuicConnection,
    mut controller: ClientH3Controller,
    state: Arc<TokioMutex<ConnectionState>>,
) {
    while let Some(event) = controller.event_receiver_mut().recv().await {
        match event {
            ClientH3Event::NewOutboundRequest {
                stream_id,
                request_id,
            } => {
                state
                    .lock()
                    .await
                    .stream_to_request
                    .insert(stream_id, request_id);
            }

            ClientH3Event::Core(H3Event::IncomingHeaders(incoming)) => {
                let state = Arc::clone(&state);
                tokio::spawn(async move {
                    dispatch_incoming_response(incoming, state).await;
                });
            }

            ClientH3Event::Core(H3Event::ConnectionError(_) | H3Event::ConnectionShutdown(_)) => {
                let mut s = state.lock().await;
                for (_, tx) in s.pending.drain() {
                    let _ = tx.send(Err(KnafehError::ConnectionClosed));
                }
                for (_, tx) in s.pending_stream.drain() {
                    let _ = tx.send(Err(KnafehError::ConnectionClosed));
                }
                break;
            }

            _ => {}
        }
    }
}

/// Route an incoming H3 response to the correct pending request.
///
/// For streaming requests, hand off the `recv` channel directly.
/// For unary requests, read all body frames first.
async fn dispatch_incoming_response(
    incoming: IncomingH3Headers,
    state: Arc<TokioMutex<ConnectionState>>,
) {
    let stream_id = incoming.stream_id;

    // Look up request_id for this stream.
    let request_id = {
        let s = state.lock().await;
        s.stream_to_request.get(&stream_id).copied()
    };

    let Some(request_id) = request_id else {
        return;
    };

    // Check if this is a streaming request.
    {
        let mut s = state.lock().await;
        if let Some(tx) = s.pending_stream.remove(&request_id) {
            s.stream_to_request.remove(&stream_id);
            let IncomingH3Headers {
                headers,
                recv,
                read_fin,
                ..
            } = incoming;
            let _ = tx.send(Ok(H3Response::Streaming {
                headers,
                recv,
                read_fin,
            }));
            return;
        }
    }

    // Unary path: read all body frames.
    read_response_and_resolve(incoming, state).await;
}

/// Read the full response body from an incoming H3 stream and resolve the
/// pending oneshot (unary path).
async fn read_response_and_resolve(
    incoming: IncomingH3Headers,
    state: Arc<TokioMutex<ConnectionState>>,
) {
    let IncomingH3Headers {
        stream_id,
        headers,
        mut recv,
        read_fin,
        ..
    } = incoming;

    let mut body = Vec::new();
    if !read_fin {
        while let Some(frame) = recv.recv().await {
            if let InboundFrame::Body(buf, fin) = frame {
                if body.len().saturating_add(buf.len()) > MAX_MESSAGE_SIZE {
                    let mut s = state.lock().await;
                    if let Some(&request_id) = s.stream_to_request.get(&stream_id) {
                        if let Some(tx) = s.pending.remove(&request_id) {
                            let _ = tx.send(Err(KnafehError::InvalidMessage(format!(
                                "response body exceeds maximum {MAX_MESSAGE_SIZE} bytes"
                            ))));
                        }
                        s.stream_to_request.remove(&stream_id);
                    }
                    return;
                }
                body.extend_from_slice(&buf);
                if fin {
                    break;
                }
            }
        }
    }

    let mut s = state.lock().await;
    if let Some(&request_id) = s.stream_to_request.get(&stream_id) {
        if let Some(tx) = s.pending.remove(&request_id) {
            let _ = tx.send(Ok(H3Response::Complete { headers, body }));
        }
        s.stream_to_request.remove(&stream_id);
    }
}