//! `GrpcTransport` — implements tidalDB's `Transport` trait over gRPC.
//!
//! Each instance embeds a tokio runtime, runs a gRPC server for receiving
//! segments, and maintains a client connection pool for sending segments.
//! The sync/async bridge uses `runtime.block_on()`, which is safe because
//! callers (WAL shipper and segment receiver) run on `std::thread`, not
//! inside a tokio context.

use std::{
    collections::HashMap,
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
};

use tidaldb::replication::{
    shard::ShardId,
    transport::{Transport, TransportError, WalSegmentPayload},
};
use tokio::sync::{Notify, mpsc};

use crate::{client::PeerPool, config::GrpcTransportConfig, error::GrpcTransportError, server};

/// A gRPC-based transport for WAL segment shipping between tidalDB shards.
///
/// Implements the synchronous [`Transport`] trait by embedding a tokio runtime.
/// The gRPC server accepts incoming segments from peers and places them in an
/// internal channel. The client pool sends segments to peers via unary RPCs.
///
/// # Single-Consumer Invariant
///
/// `recv_segment` must only be called from a single thread (the segment receiver).
/// The `Mutex` around the receiver is a safety net, not a concurrency mechanism.
/// Concurrent callers would serialize on the mutex while each blocks on
/// `runtime.block_on(rx.recv())`, which is safe but wasteful.
///
/// # Deterministic receiver shutdown
///
/// A parked [`recv_segment`](Self::recv_segment) blocks on the inbound channel,
/// which only returns `None` once every `inbound_tx` sender drops — and those live
/// inside the server task. Aborting the server task on [`Drop`] eventually frees
/// them, but the receiver thread is typically a detached `std::thread`, so a clean
/// (non-process-exit) teardown could leave it parked. To make shutdown deterministic
/// the transport carries a latching `shutdown` signal that `Drop` (and the explicit
/// [`shutdown_receivers`](Self::shutdown_receivers)) trips **before** tearing the
/// runtime down: `recv_segment` checks the latch on entry and selects over the channel
/// and a wake `Notify`, so a parked receiver returns `None` immediately on shutdown and
/// its thread can exit and be joined.
pub struct GrpcTransport {
    config: GrpcTransportConfig,
    /// `Some` for the transport's whole life; taken in [`Drop`] so the runtime
    /// can be torn down with the non-blocking [`tokio::runtime::Runtime::shutdown_background`]
    /// (which consumes the runtime) instead of the default blocking `Runtime::drop`.
    runtime: Option<tokio::runtime::Runtime>,
    inbound_rx: Mutex<mpsc::Receiver<WalSegmentPayload>>,
    pool: PeerPool,
    server_handle: tokio::task::JoinHandle<Result<(), tonic::transport::Error>>,
    /// Shutdown signal for the receiver. The [`AtomicBool`] **latches** the request
    /// so a `recv_segment` that has not yet parked still observes it on entry (no
    /// missed-wakeup race), while the [`Notify`] wakes a receiver that is *already*
    /// parked. Tripped by [`shutdown_receivers`](Self::shutdown_receivers) and by
    /// [`Drop`] before the runtime is torn down, so a parked
    /// [`recv_segment`](Self::recv_segment) returns `None` deterministically instead
    /// of waiting for every server-side `inbound_tx` sender to drop.
    shutdown: Arc<ShutdownSignal>,
}

/// Latching shutdown signal: an [`AtomicBool`] that survives the not-yet-parked race
/// plus a [`Notify`] to wake an already-parked waiter. Held by [`GrpcTransport`] and
/// tripped on [`Drop`] / [`shutdown_receivers`](GrpcTransport::shutdown_receivers).
struct ShutdownSignal {
    requested: AtomicBool,
    notify: Notify,
}

impl ShutdownSignal {
    fn new() -> Self {
        Self {
            requested: AtomicBool::new(false),
            notify: Notify::new(),
        }
    }

    /// Latch the shutdown request and wake any already-parked receiver.
    ///
    /// `Release` so the latch is visible to a receiver that subsequently reads it
    /// with `Acquire` on `recv_segment` entry; `notify_waiters` covers a receiver
    /// already parked on `notified()` at the moment this is called.
    fn request(&self) {
        self.requested.store(true, Ordering::Release);
        self.notify.notify_waiters();
    }

    /// Whether shutdown has been requested. `Acquire` pairs with the `Release` in
    /// [`request`](Self::request) so a not-yet-parked receiver observes the latch.
    fn is_requested(&self) -> bool {
        self.requested.load(Ordering::Acquire)
    }
}

/// Install a process-wide rustls [`CryptoProvider`] before tonic's TLS
/// builders run. tidal-net pulls rustls only transitively (via tonic), with no
/// provider feature in its own dependency closure, so under a narrow build the
/// process-level default is absent and rustls panics (could not automatically
/// determine the process-level [`CryptoProvider`]); under a wider build two
/// providers can be present, which is equally ambiguous. Installing aws-lc-rs
/// explicitly here — idempotently, ignoring the error when another component
/// (e.g. a reqwest-based client elsewhere in the process) already set a default
/// — removes that fragility for the mTLS tests, a standalone tidal-server, and
/// the replication path alike. See `tidal-net/BUILD.bazel`.
fn ensure_crypto_provider() {
    static ONCE: std::sync::Once = std::sync::Once::new();
    ONCE.call_once(|| {
        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
    });
}

impl GrpcTransport {
    /// Create and start a new gRPC transport.
    ///
    /// This starts a gRPC server on `config.listen_addr` and builds lazy
    /// connections to all configured peers.
    ///
    /// # Errors
    ///
    /// Returns an error if the server or client pool fails to initialize.
    pub fn new(config: GrpcTransportConfig) -> Result<Self, GrpcTransportError> {
        ensure_crypto_provider();

        // Validate numeric invariants BEFORE building anything (C15). In
        // particular `mpsc::channel(config.channel_capacity)` below panics on a
        // zero capacity (tokio asserts buffer > 0), so a `channel_capacity == 0`
        // config must be caught here as a typed `Internal` error rather than
        // panicking deep in the constructor. `PeerPool::new` validates again
        // (idempotent, cheap); doing it up front is what guards the real
        // constructor ordering the config doc promises.
        config.validate()?;

        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .thread_name(format!("tidal-grpc-{}", config.local_shard))
            .build()
            .map_err(|e| GrpcTransportError::Internal(format!("tokio runtime: {e}")))?;

        let (inbound_tx, inbound_rx) = mpsc::channel(config.channel_capacity);

        // Build server and client pool inside the runtime context.
        // TLS setup in tonic requires a tokio reactor to be available.
        let (server_handle, pool) = runtime.block_on(async {
            let handle = server::start_server(&config, inbound_tx)?;
            let pool = PeerPool::new(&config)?;
            Ok::<_, GrpcTransportError>((handle, pool))
        })?;

        Ok(Self {
            config,
            runtime: Some(runtime),
            inbound_rx: Mutex::new(inbound_rx),
            pool,
            server_handle,
            shutdown: Arc::new(ShutdownSignal::new()),
        })
    }

    /// The embedded tokio runtime.
    ///
    /// # Infallible by construction
    ///
    /// The `expect` here can never fire on any production path. `runtime` is set to
    /// `Some` exactly once, in [`new`](Self::new), and is only ever moved out in
    /// [`Drop::drop`] via `self.runtime.take()`. `Drop` runs at most once and is the
    /// final use of `self`; no method (`send_segment`, `recv_segment`, …) can observe
    /// `self` after `Drop` has begun, so every call to `runtime()` sees `Some`. The
    /// `Option` exists purely so `Drop` can move the runtime out and tear it down with
    /// the non-blocking `shutdown_background` (see the field doc on `runtime`), not to
    /// model a genuinely-absent runtime — hence `expect` over fallible propagation: a
    /// `None` here would be a memory-safety/lifecycle bug in this module, not a runtime
    /// condition a caller could handle.
    const fn runtime(&self) -> &tokio::runtime::Runtime {
        self.runtime
            .as_ref()
            .expect("runtime is Some for the transport's whole lifetime; only Drop takes it")
    }

    /// Signal any current or future [`recv_segment`](Self::recv_segment) to return `None`.
    ///
    /// Latches the shutdown request and wakes a receiver already blocked in
    /// `recv_segment` so it returns `None` — its exit condition per the [`Transport`]
    /// contract — letting the receiver thread exit and be joined. Because the request
    /// latches, a receiver that has not yet parked also observes it on its next entry.
    ///
    /// # Why this exists separately from [`Drop`]
    ///
    /// The receiver thread typically holds its own `Arc<dyn Transport>` clone, so
    /// [`GrpcTransport::drop`] cannot fire until that thread has already exited and
    /// released its `Arc` — a cycle: the thread parks in `recv_segment` and never
    /// releases the `Arc`, so `Drop` never runs, so the thread is never woken. Break
    /// the cycle by calling `shutdown_receivers` from the owner's shutdown path (which
    /// holds a *different* `Arc`) BEFORE joining the receiver thread and dropping its
    /// `Arc`. The notify-trip in [`Drop`] remains a backstop for the non-cyclic case
    /// (e.g. a transport owned outright, as in this crate's tests).
    ///
    /// Idempotent: the request latches, so any current or future `recv_segment`
    /// returns `None`.
    pub fn shutdown_receivers(&self) {
        self.shutdown.request();
    }

    /// Whether the embedded gRPC serve loop has terminated, for ANY reason.
    ///
    /// The serve task runs for the transport's whole lifetime; a `true` here
    /// means the listener stopped accepting — either because the transport is
    /// shutting down or because the serve loop failed (in which case the failure
    /// was already logged at `error!` inside the task; see [`server::start_server`]).
    ///
    /// Because this conflates a clean shutdown with a silent death, prefer
    /// [`serve_loop_died`](Self::serve_loop_died) for liveness decisions — it is
    /// the one that distinguishes the two (C14).
    #[must_use]
    pub fn server_terminated(&self) -> bool {
        self.server_handle.is_finished()
    }

    /// Whether the gRPC serve loop has died WITHOUT a shutdown being requested.
    ///
    /// This is the load-bearing liveness signal (C14): on a follower the receive
    /// side dies silently when the listener stops, a TLS handshake task panics,
    /// or the reactor is torn down. In all those cases the spawned server task
    /// ends, its `inbound_tx` drops, and [`recv_segment`](Self::recv_segment)
    /// returns `None` — *exactly* as it does on a clean
    /// [`shutdown_receivers`](Self::shutdown_receivers). The two are otherwise
    /// indistinguishable, so a dead follower would look like an intentionally
    /// stopped one.
    ///
    /// `serve_loop_died()` resolves the ambiguity: it is `true` only when the
    /// serve task has finished AND no shutdown was requested. Health checks and
    /// the cluster control plane poll this so a follower whose receive side has
    /// silently died is observable and demoted, not a black hole; a clean
    /// shutdown returns `false` here even though [`server_terminated`](Self::server_terminated)
    /// is `true`.
    #[must_use]
    pub fn serve_loop_died(&self) -> bool {
        self.server_handle.is_finished() && !self.shutdown.is_requested()
    }

    /// Assert we are not inside a tokio runtime (`block_on` would panic).
    #[cfg(debug_assertions)]
    fn assert_not_in_async_context() {
        debug_assert!(
            tokio::runtime::Handle::try_current().is_err(),
            "GrpcTransport methods must not be called from within a tokio runtime; \
             use std::thread instead"
        );
    }
}

impl Transport for GrpcTransport {
    fn send_segment(&self, to: ShardId, payload: WalSegmentPayload) -> Result<(), TransportError> {
        #[cfg(debug_assertions)]
        Self::assert_not_in_async_context();

        // Validate payload size before sending.
        if payload.bytes.len() > self.config.max_payload_bytes {
            return Err(TransportError::PayloadTooLarge {
                size: payload.bytes.len(),
                max: self.config.max_payload_bytes,
            });
        }

        self.runtime()
            .block_on(self.pool.send_to(to, payload))
            .map_err(TransportError::from)
    }

    fn recv_segment(&self) -> Option<WalSegmentPayload> {
        #[cfg(debug_assertions)]
        Self::assert_not_in_async_context();

        let mut rx = match self.inbound_rx.lock() {
            Ok(rx) => rx,
            Err(poisoned) => {
                tracing::error!("inbound_rx lock poisoned; recovering guard");
                poisoned.into_inner()
            }
        };
        let shutdown = &self.shutdown;
        // Fast path: if shutdown was already requested before we ever parked, the
        // latched flag short-circuits to `None`. This closes the missed-wakeup race
        // where `shutdown_receivers`/`Drop` ran between a prior `recv_segment`
        // returning and this call constructing its `notified()` future.
        if shutdown.is_requested() {
            return None;
        }
        // Otherwise select over an inbound segment and the shutdown notify so a
        // parked receiver returns `None` the instant shutdown is signaled, rather
        // than waiting for every server-side `inbound_tx` sender to drop. `None`
        // from either arm is the receiver's exit condition (see the `Transport`
        // trait contract), so a tripped shutdown looks exactly like channel close.
        self.runtime().block_on(async {
            // Register the `notified()` future BEFORE the final flag re-check so a
            // `request()` that runs concurrently either (a) is seen by this re-check
            // or (b) finds us already registered and wakes us — never both-missed.
            let notified = shutdown.notify.notified();
            tokio::pin!(notified);
            if shutdown.is_requested() {
                return None;
            }
            tokio::select! {
                // Bias toward draining already-queued segments before observing
                // shutdown, so a shutdown that races with in-flight segments does
                // not silently discard one already sitting in the channel.
                biased;
                segment = rx.recv() => segment,
                () = &mut notified => None,
            }
        })
    }

    fn local_shard(&self) -> ShardId {
        self.config.local_shard
    }
}

impl Drop for GrpcTransport {
    fn drop(&mut self) {
        // 1. Trip the shutdown notify FIRST, while the runtime is still alive, so a
        //    parked `recv_segment` wakes on its `shutdown.notified()` arm and returns
        //    `None` deterministically — its receiver thread can then exit and join
        //    instead of leaking until process exit. `notify_waiters` wakes a waiter
        //    that is already parked; `recv_segment` re-creates its future on each call
        //    so there is no missed-wakeup window for a not-yet-parked receiver (a
        //    fresh `recv_segment` after Drop sees a closed channel once the server
        //    task is aborted below). The latch in `request()` also makes any future
        //    `recv_segment` short-circuit to `None` without parking.
        self.shutdown.request();

        // 2. Abort the server task so its `inbound_tx` is released (so a receiver
        //    that calls `recv_segment` AFTER Drop also observes channel close), then
        //    tear the runtime down WITHOUT blocking.
        //
        //    A `GrpcTransport` can be dropped from within an async context — e.g.
        //    cluster shutdown unwinding on the axum reactor thread — where the
        //    default blocking `Runtime::drop` panics ("Cannot drop a runtime in a
        //    context where blocking is not allowed"). `shutdown_background` returns
        //    immediately and reclaims the worker threads off-thread, so `Drop` is
        //    safe in any context.
        self.server_handle.abort();
        if let Some(runtime) = self.runtime.take() {
            runtime.shutdown_background();
        }
    }
}

/// Factory for building a set of [`GrpcTransport`] instances, one per shard.
///
/// Analogous to `InProcessTransportFactory` but for gRPC connections.
pub struct GrpcTransportFactory;

impl GrpcTransportFactory {
    /// Build one `GrpcTransport` per configuration.
    ///
    /// Each transport gets its own tokio runtime, gRPC server, and client pool.
    ///
    /// # Errors
    ///
    /// Returns an error if any transport fails to initialize.
    pub fn build(
        configs: Vec<GrpcTransportConfig>,
    ) -> Result<HashMap<ShardId, GrpcTransport>, GrpcTransportError> {
        let mut result = HashMap::new();
        for config in configs {
            let shard = config.local_shard;
            let transport = GrpcTransport::new(config)?;
            result.insert(shard, transport);
        }
        Ok(result)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)] // test assertions on known-good fixtures
mod tests {
    use std::net::SocketAddr;

    use super::*;

    /// Bind port 0 to obtain a free, OS-assigned address (tonic cannot bind 0
    /// directly, so we resolve a concrete port up front).
    fn free_addr() -> SocketAddr {
        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
        listener.local_addr().unwrap()
    }

    /// C15: `GrpcTransport::new` with `channel_capacity == 0` must return a typed
    /// `Internal` error, NOT panic. The panic used to happen at
    /// `mpsc::channel(0)` BEFORE `PeerPool::new` ran `validate()`, so the only
    /// existing test (which called `PeerPool::new` directly) masked it. This
    /// drives the real constructor ordering.
    #[test]
    #[allow(clippy::significant_drop_tightening)]
    fn new_with_zero_channel_capacity_returns_internal_not_panic() {
        let config = GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: free_addr(),
            channel_capacity: 0,
            insecure: true,
            ..Default::default()
        };
        let result = GrpcTransport::new(config);
        assert!(
            matches!(result, Err(GrpcTransportError::Internal(_))),
            "channel_capacity == 0 must be a typed Internal error, not a panic"
        );
    }

    /// C14: a serve loop that dies WITHOUT a shutdown request must be observable
    /// via `serve_loop_died()`, distinct from a clean `shutdown_receivers()`.
    /// We abort the server task to model a silent serve-loop death and assert the
    /// transport reports `serve_loop_died() == true` while a transport that was
    /// cleanly shut down reports `false` (even though `server_terminated()` is
    /// `true` for both).
    #[test]
    #[allow(clippy::significant_drop_tightening)]
    fn serve_loop_died_distinguishes_silent_death_from_clean_shutdown() {
        // --- silent death: abort the serve task, no shutdown requested ---
        let dead = GrpcTransport::new(GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: free_addr(),
            insecure: true,
            ..Default::default()
        })
        .unwrap();
        // Kill the serve loop the way a listener-death / reactor-teardown would.
        dead.server_handle.abort();
        // Wait for the abort to take effect.
        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
        while !dead.server_terminated() && std::time::Instant::now() < deadline {
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
        assert!(
            dead.server_terminated(),
            "the aborted serve task must report server_terminated()"
        );
        assert!(
            dead.serve_loop_died(),
            "a serve loop that died WITHOUT a shutdown request must report \
             serve_loop_died() == true (C14)"
        );

        // --- clean shutdown: request shutdown, then the task ends ---
        let clean = GrpcTransport::new(GrpcTransportConfig {
            local_shard: ShardId(1),
            listen_addr: free_addr(),
            insecure: true,
            ..Default::default()
        })
        .unwrap();
        clean.shutdown_receivers();
        clean.server_handle.abort();
        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
        while !clean.server_terminated() && std::time::Instant::now() < deadline {
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
        assert!(clean.server_terminated());
        assert!(
            !clean.serve_loop_died(),
            "a serve loop that ended AFTER shutdown was requested is a clean \
             shutdown, not a silent death (C14)"
        );
    }
}
