//! Connection accept-loop helpers
//!
//! Breaks the concerns of `Server::listen_with_shutdown` into three focused types:
//!
//! - [`ConnectionLimiter`] — optional semaphore-based concurrency cap; rejects
//!   excess connections with a 503 response before they reach the service.
//! - [`ConnectionTracker`] — atomic counter of in-flight connections; used by
//!   the graceful-shutdown drain to know when it is safe to exit.
//! - [`ConnectionTask`] — spawns and drives a single HTTP/1.1 connection,
//!   wiring in the graceful-shutdown broadcast so the connection can finish its
//!   current request before closing.
//!
//! The main loop in `Server::listen_with_shutdown` becomes a thin orchestrator
//! that delegates each concern to one of these types.

use crate::core::{CoreError, StatusCode};
use crate::domain::Response;
use crate::http::to_hyper_response;
use crate::server::RuxnoService;
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper_util::rt::TokioIo;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::sync::{broadcast, OwnedSemaphorePermit, Semaphore};
use typed_builder::TypedBuilder;

// ─── ConnectionLimiter ───────────────────────────────────────────────────────

/// Optional semaphore-based connection cap.
///
/// When `max_connections` is `None` the limiter is a no-op and every call to
/// [`ConnectionLimiter::try_acquire`] returns `Ok(None)`.
///
/// When a limit is configured, [`ConnectionLimiter::try_acquire`] either returns
/// an owned permit (which must be held for the lifetime of the connection) or
/// `Err(())` to signal that the caller should reject the connection.
pub(super) struct ConnectionLimiter {
    semaphore: Option<Arc<Semaphore>>,
    production_mode: bool,
}

impl ConnectionLimiter {
    pub(super) fn new(max_connections: Option<usize>, production_mode: bool) -> Self {
        let semaphore = max_connections.map(|max| Arc::new(Semaphore::new(max)));
        Self {
            semaphore,
            production_mode,
        }
    }

    /// Try to acquire a connection permit.
    ///
    /// Returns:
    /// - `Ok(None)` — no limit configured, proceed freely.
    /// - `Ok(Some(permit))` — limit configured, permit acquired; hold until
    ///   the connection closes.
    /// - `Err(())` — limit reached; caller should reject the connection.
    pub(super) fn try_acquire(&self) -> Result<Option<OwnedSemaphorePermit>, ()> {
        match &self.semaphore {
            None => Ok(None),
            Some(sem) => match sem.clone().try_acquire_owned() {
                Ok(permit) => Ok(Some(permit)),
                Err(_) => Err(()),
            },
        }
    }

    /// Send a 503 Service Unavailable response and close the connection.
    ///
    /// Called when [`try_acquire`] returns `Err(())`. Spawns a short-lived task
    /// so the caller can continue accepting the next connection immediately.
    pub(super) fn reject_connection(&self, stream: TcpStream, peer_addr: SocketAddr) {
        eprintln!(
            "🚫 Connection limit reached, rejecting connection from {}",
            peer_addr
        );

        let production_mode = self.production_mode;
        tokio::spawn(async move {
            let io = TokioIo::new(stream);
            let svc = service_fn(move |_req: hyper::Request<Incoming>| async move {
                let domain_res = Response::from_error_with_status(
                    CoreError::custom(
                        503,
                        "Server connection limit reached. Please try again later.",
                    ),
                    StatusCode::SERVICE_UNAVAILABLE,
                    production_mode,
                );
                let mut hyper_res = to_hyper_response(domain_res);
                hyper_res.headers_mut().insert(
                    hyper::header::RETRY_AFTER,
                    hyper::header::HeaderValue::from_static("5"),
                );
                Ok::<_, std::convert::Infallible>(hyper_res)
            });

            let _ = http1::Builder::new().serve_connection(io, svc).await;
        });
    }
}

// ─── ConnectionTracker ───────────────────────────────────────────────────────

/// Atomic counter of in-flight connections.
///
/// Incremented when a connection task is spawned; decremented when it exits.
/// The graceful-shutdown drain polls [`ConnectionTracker::active`] until it
/// reaches zero (or the timeout expires).
#[derive(Clone)]
pub(super) struct ConnectionTracker {
    count: Arc<AtomicUsize>,
}

impl ConnectionTracker {
    pub(super) fn new() -> Self {
        Self {
            count: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Increment the counter (call before spawning a connection task).
    pub(super) fn increment(&self) {
        self.count.fetch_add(1, Ordering::Relaxed);
    }

    /// Decrement the counter (call when a connection task exits).
    pub(super) fn decrement(&self) {
        self.count.fetch_sub(1, Ordering::Relaxed);
    }

    /// Current number of active connections.
    pub(super) fn active(&self) -> usize {
        self.count.load(Ordering::Acquire)
    }

    /// Wait until all connections have closed, or until `timeout` elapses.
    ///
    /// Returns `true` if all connections closed within the timeout, `false`
    /// if the timeout was reached with connections still open.
    pub(super) async fn drain(&self, timeout: Duration) -> bool {
        tokio::time::timeout(timeout, async {
            while self.active() > 0 {
                tokio::time::sleep(Duration::from_millis(100)).await;
            }
        })
        .await
        .is_ok()
    }
}

// ─── ConnectionTask ──────────────────────────────────────────────────────────

/// Parameters for spawning a single HTTP/1.1 connection task.
///
/// Use [`ConnectionTask::builder()`] to construct this — the `TypedBuilder`
/// derive enforces that every field is set before calling `.spawn()`.
///
/// # Example
///
/// ```rust,ignore
/// ConnectionTask::builder()
///     .stream(stream)
///     .peer_addr(peer_addr)
///     .service(service.clone())
///     .max_body_size(max_body_size)
///     .max_headers(max_headers)
///     .permit(permit)
///     .tracker(tracker.clone())
///     .shutdown_rx(shutdown_tx.subscribe())
///     .build()
///     .spawn();
/// ```
#[derive(TypedBuilder)]
pub(super) struct ConnectionTask<E>
where
    E: Send + Sync + 'static,
{
    /// Accepted TCP stream.
    stream: TcpStream,
    /// Remote address (for error logging).
    peer_addr: SocketAddr,
    /// Cloneable service for this connection.
    service: RuxnoService<E>,
    /// Maximum request body size forwarded to `service.handle`.
    max_body_size: usize,
    /// Maximum header count forwarded to `service.handle`.
    max_headers: usize,
    /// Optional semaphore permit; dropped when the task exits.
    permit: Option<OwnedSemaphorePermit>,
    /// Connection counter; decremented when the task exits.
    tracker: ConnectionTracker,
    /// Broadcast receiver; fires when the server shuts down.
    #[builder(setter(transform = |rx: broadcast::Receiver<()>| rx))]
    shutdown_rx: broadcast::Receiver<()>,
}

impl<E> ConnectionTask<E>
where
    E: Send + Sync + 'static + Clone,
{
    /// Spawn the connection task onto the Tokio runtime.
    ///
    /// The task:
    /// 1. Holds the semaphore permit for the connection's lifetime.
    /// 2. Serves requests via `RuxnoService`.
    /// 3. Listens for the shutdown broadcast and calls `graceful_shutdown()` on
    ///    the hyper connection so the current request finishes before closing.
    /// 4. Decrements the `ConnectionTracker` on exit.
    pub(super) fn spawn(self) {
        let Self {
            stream,
            peer_addr,
            service,
            max_body_size,
            max_headers,
            permit,
            tracker,
            mut shutdown_rx,
        } = self;

        tokio::spawn(async move {
            // Hold permit for the duration of the connection.
            // Dropping it at task exit releases one slot in the semaphore.
            let _permit = permit;
            let _tracker_guard = TrackerGuard(tracker);

            let io = TokioIo::new(stream);

            let svc = service_fn(move |req: hyper::Request<Incoming>| {
                let service = service.clone();
                async move { service.handle(req, max_body_size, max_headers).await }
            });

            // `.with_upgrades()` is required for WebSocket — without it the
            // upgrade handshake hangs because hyper never polls the upgrade future.
            let connection = http1::Builder::new()
                .serve_connection(io, svc)
                .with_upgrades();

            tokio::pin!(connection);

            tokio::select! {
                result = &mut connection => {
                    if let Err(e) = result {
                        eprintln!("Connection error from {}: {}", peer_addr, e);
                    }
                }
                _ = shutdown_rx.recv() => {
                    // Signal hyper to finish the current request then close.
                    connection.as_mut().graceful_shutdown();

                    if let Err(e) = connection.await {
                        if !e.to_string().contains("connection closed") {
                            eprintln!("Connection shutdown error from {}: {}", peer_addr, e);
                        }
                    }
                }
            }

            // _permit dropped here → semaphore slot released
        });
    }
}

struct TrackerGuard(ConnectionTracker);

impl Drop for TrackerGuard {
    fn drop(&mut self) {
        self.0.decrement();
    }
}
