//! `PartitionProxy`: a root-free, in-harness TCP relay for REAL network-partition
//! injection between OS processes (m8p10 task 05).
//!
//! # Why a proxy and not iptables/pfctl
//!
//! The ROADMAP sanctions iptables/pfctl *or* a proxy ("toxiproxy or similar").
//! `iptables`/`pfctl` require root and differ between Linux and macOS; a developer
//! laptop cannot run them unprivileged. A user-space TCP relay the test owns is
//! identical on every platform, needs no privileges, and severs connections at the
//! genuine OS transport layer: established streams are `shutdown(Both)` (the peer's
//! in-flight gRPC `ShipSegment` returns a real transport error, the circuit breaker
//! opens, HTTP fan-outs time out), and new connections are accepted-then-dropped so
//! a fresh dial through a severed proxy fails its HTTP/2 handshake. That is a true
//! network partition between the two processes — not an engine flag.
//!
//! # How it interposes (the harness rewrite hook)
//!
//! [`MultiProcCluster`](super::multiproc::MultiProcCluster) writes a per-process
//! topology where a process's OWN region lists its REAL bind addresses but every
//! PEER region lists a *published* address from the
//! [`AddrRewrite`](super::multiproc::AddrRewrite) hook. A `PartitionProxy` listens
//! on a free loopback port and relays to a peer's real target; pointing the peers'
//! published address at the proxy means peer↔peer traffic flows through it while
//! each process still self-binds its real port. Severing the proxy therefore cuts
//! ONLY the path peers use to reach that region — the test client still talks to
//! every node's real HTTP address directly (the operator's console survives the
//! partition, matching the runbook drill's "read the stale follower" step).
//!
//! # Threading / cleanup discipline
//!
//! Pure `std::net` + `std::thread`; no tokio. One acceptor thread per proxy plus
//! two pump threads per live connection (each direction). Every live downstream and
//! upstream `TcpStream` is tracked (a clone) so `sever()` can `shutdown(Both)` it.
//! [`Drop`] latches a stop flag, shuts every tracked stream, unblocks the acceptor
//! by dialing its own listener, and joins the acceptor; pump threads observe the
//! shutdown via a read/write error and exit on their own (they are detached, never
//! leaked past process exit because a closed socket unblocks their blocking I/O).

#![allow(dead_code)]

use std::{
    io::{Read, Write},
    net::{Shutdown, SocketAddr, TcpListener, TcpStream},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle},
    time::Duration,
};

/// A user-space TCP relay between peers and a region's real listener, with a latch
/// the test flips to inject / heal a real network partition.
pub struct PartitionProxy {
    /// The loopback port peers dial (published in their topology for the target).
    listen: SocketAddr,
    /// The region's real grpc/http listener this relay forwards to.
    target: SocketAddr,
    /// Latched true by [`sever`](Self::sever): the acceptor accept-and-drops, and
    /// every live stream is `shutdown(Both)`. Cleared by [`heal_link`](Self::heal_link).
    severed: Arc<AtomicBool>,
    /// Latched true on [`Drop`]: the acceptor exits and stops spawning pumps.
    stopped: Arc<AtomicBool>,
    /// Live client (downstream) + target (upstream) stream clones, so `sever()` can
    /// `shutdown(Both)` every in-flight connection. Cleared lazily on sever.
    conns: Arc<Mutex<Vec<TcpStream>>>,
    /// The acceptor thread handle, joined on [`Drop`].
    acceptor: Option<JoinHandle<()>>,
}

impl PartitionProxy {
    /// Start a relay on a free loopback port forwarding to `target`. Returns once
    /// the listener is bound (so the published address is dialable immediately).
    ///
    /// # Panics
    ///
    /// Panics if a loopback listener cannot be bound or its address read.
    #[must_use]
    pub fn start(target: SocketAddr) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("bind proxy listener");
        let listen = listener.local_addr().expect("read proxy local_addr");

        let severed = Arc::new(AtomicBool::new(false));
        let stopped = Arc::new(AtomicBool::new(false));
        let conns: Arc<Mutex<Vec<TcpStream>>> = Arc::new(Mutex::new(Vec::new()));

        let acceptor = {
            let severed = Arc::clone(&severed);
            let stopped = Arc::clone(&stopped);
            let conns = Arc::clone(&conns);
            thread::spawn(move || accept_loop(&listener, target, &severed, &stopped, &conns))
        };

        Self {
            listen,
            target,
            severed,
            stopped,
            conns,
            acceptor: Some(acceptor),
        }
    }

    /// The address peers should dial (published in their topology for the target).
    #[must_use]
    pub const fn addr(&self) -> SocketAddr {
        self.listen
    }

    /// The real target this relay forwards to.
    #[must_use]
    pub const fn target(&self) -> SocketAddr {
        self.target
    }

    /// Whether the link is currently severed.
    #[must_use]
    pub fn is_severed(&self) -> bool {
        self.severed.load(Ordering::Acquire)
    }

    /// Inject a partition: refuse new connections (accept→drop) AND `shutdown(Both)`
    /// every live stream so in-flight peer traffic fails for real. Idempotent.
    pub fn sever(&self) {
        self.severed.store(true, Ordering::Release);
        // Tear down every tracked stream in BOTH directions. A pump blocked in
        // `read`/`write` on a shut stream returns an error and exits; the partner
        // pump's next I/O then also fails. We drain the vec — new connections after
        // a heal repopulate it.
        let mut guard = self
            .conns
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        for stream in guard.drain(..) {
            let _ = stream.shutdown(Shutdown::Both);
        }
    }

    /// Heal the link: clear the latch so new connections flow again. Existing
    /// streams stay severed (gRPC/HTTP reconnect on their own — tonic re-dials, the
    /// blocking reqwest client opens a fresh connection). Idempotent.
    pub fn heal_link(&self) {
        self.severed.store(false, Ordering::Release);
    }
}

impl Drop for PartitionProxy {
    fn drop(&mut self) {
        // Stop the acceptor and tear down every live stream so pump threads unblock.
        self.stopped.store(true, Ordering::Release);
        self.severed.store(true, Ordering::Release);
        {
            let mut guard = self
                .conns
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            for stream in guard.drain(..) {
                let _ = stream.shutdown(Shutdown::Both);
            }
        }
        // Unblock the acceptor's blocking `accept()` by dialing our own listener
        // (it will observe `stopped` and exit). A failed dial is fine — the thread
        // also re-checks `stopped` on its accept-error path.
        let _ = TcpStream::connect_timeout(&self.listen, Duration::from_millis(200));
        if let Some(handle) = self.acceptor.take() {
            let _ = handle.join();
        }
    }
}

/// The acceptor loop: for each inbound connection, dial the target and start two
/// pump threads (downstream→upstream, upstream→downstream). On `severed`, accept
/// the connection and immediately drop it (a real connection-refused-equivalent at
/// the application layer: the peer's HTTP/2 handshake never completes).
fn accept_loop(
    listener: &TcpListener,
    target: SocketAddr,
    severed: &Arc<AtomicBool>,
    stopped: &Arc<AtomicBool>,
    conns: &Arc<Mutex<Vec<TcpStream>>>,
) {
    for incoming in listener.incoming() {
        if stopped.load(Ordering::Acquire) {
            return;
        }
        let Ok(downstream) = incoming else {
            // A transient accept error: bail only if we are stopping.
            if stopped.load(Ordering::Acquire) {
                return;
            }
            continue;
        };
        if severed.load(Ordering::Acquire) {
            // Severed: drop the just-accepted connection. The peer sees its
            // connection die immediately (no relay to the target).
            let _ = downstream.shutdown(Shutdown::Both);
            continue;
        }
        // Dial the real target. A failed dial drops the downstream (peer sees a
        // dead connection) — never blocks the acceptor for other peers.
        let Ok(upstream) = TcpStream::connect_timeout(&target, Duration::from_secs(2)) else {
            let _ = downstream.shutdown(Shutdown::Both);
            continue;
        };
        // Disable Nagle so the gRPC/HTTP framing flows promptly through the relay.
        let _ = downstream.set_nodelay(true);
        let _ = upstream.set_nodelay(true);

        // Track clones of both ends so `sever()` can shut them down.
        register_conn(conns, &downstream);
        register_conn(conns, &upstream);

        spawn_pump(&downstream, &upstream);
        spawn_pump(&upstream, &downstream);
    }
}

/// Register a clone of `stream` in the live-connection set (best-effort; a failed
/// clone just means `sever()` cannot reach this end, but its partner pump will
/// still tear the pair down when ITS end is shut).
fn register_conn(conns: &Arc<Mutex<Vec<TcpStream>>>, stream: &TcpStream) {
    if let Ok(clone) = stream.try_clone() {
        conns
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .push(clone);
    }
}

// ── Harness glue: per-directed-edge proxies via the rewrite hook ──────────────

use std::collections::HashMap;

use super::multiproc::{AddrKind, AddrRewrite};

/// The directed edge a proxy interposes: the OBSERVER region (which dials the
/// published address) reaching the PEER region over one link `kind`. One
/// `PartitionProxy` sits on each `(observer → peer, kind)` edge, so the harness
/// can sever a SINGLE direction (cut eu-west→ap-south while us-east→ap-south stays
/// up — the follower↔follower partition the ROADMAP names) or every inbound edge
/// of a region at once (full node isolation).
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
struct EdgeKey {
    observer: String,
    peer: String,
    kind: AddrKind,
}

/// The two link proxies (gRPC + HTTP) for ONE directed `observer → peer` edge.
///
/// Severing BOTH = that observer cannot ship segments OR forward/aggregate over
/// HTTP to the peer; severing only `grpc` cuts replication while leaving the HTTP
/// control/forward plane intact (a narrower fault worth exercising on its own).
pub struct EdgeProxies {
    /// The gRPC relay for this directed edge (the observer's published `grpc_addr`).
    pub grpc: Arc<PartitionProxy>,
    /// The HTTP relay for this directed edge (the observer's published `http_addr`).
    pub http: Arc<PartitionProxy>,
}

impl EdgeProxies {
    /// Sever BOTH links of this directed edge.
    pub fn sever_all(&self) {
        self.grpc.sever();
        self.http.sever();
    }

    /// Sever ONLY the gRPC (replication) link, leaving HTTP control/forward up.
    pub fn sever_grpc(&self) {
        self.grpc.sever();
    }

    /// Heal BOTH links of this directed edge.
    pub fn heal_all(&self) {
        self.grpc.heal_link();
        self.http.heal_link();
    }
}

/// Every inbound link proxy of ONE region, across all observers — severing the set
/// is a complete network partition of the region from EVERY peer (no peer can ship
/// to it, forward to it, or aggregate its status). This is the "isolate one node"
/// fault used by UAT steps 3/4.
pub struct RegionProxies {
    /// All `(observer → region)` link proxies (both kinds, every observer).
    links: Vec<Arc<PartitionProxy>>,
    /// Just the gRPC link proxies, for a replication-only sever.
    grpc_links: Vec<Arc<PartitionProxy>>,
}

impl RegionProxies {
    /// Sever EVERY inbound link (gRPC + HTTP, every peer): full node isolation.
    pub fn sever_all(&self) {
        for p in &self.links {
            p.sever();
        }
    }

    /// Sever only the inbound gRPC (replication) links, every peer — HTTP stays up
    /// (status aggregation / forwards still reach the region; only WAL ships fail).
    pub fn sever_grpc(&self) {
        for p in &self.grpc_links {
            p.sever();
        }
    }

    /// Heal EVERY inbound link.
    pub fn heal_all(&self) {
        for p in &self.links {
            p.heal_link();
        }
    }
}

/// Collects the [`PartitionProxy`] instances an [`AddrRewrite`] hook creates while
/// the harness generates per-process topology files, keyed by directed edge.
///
/// Build one with [`proxied_rewrite`], pass the rewrite into
/// `ClusterOptions::with_rewrite`, start the cluster, then call [`region`] or
/// [`edge`] to recover the handles and sever/heal each link.
///
/// [`region`]: Self::region
/// [`edge`]: Self::edge
#[derive(Clone)]
pub struct ProxyController {
    /// `(observer, peer, kind) → proxy`. The rewrite closure inserts one proxy per
    /// directed edge as the per-process topology files are generated (each peer
    /// process's file declares its OWN published address for every other region).
    proxies: Arc<Mutex<HashMap<EdgeKey, Arc<PartitionProxy>>>>,
    /// The peer regions we proxy. A peer NOT in this set rewrites to identity (its
    /// real address) on every edge, so only edges INTO a chosen region are
    /// interposed.
    targets: Arc<Vec<String>>,
}

impl ProxyController {
    /// Recover ALL inbound link proxies for one proxied region (every observer,
    /// both kinds), so the test can isolate the whole region from its peers.
    ///
    /// # Panics
    ///
    /// Panics if the region was never proxied (no inbound edge recorded).
    #[must_use]
    pub fn region(&self, region: &str) -> RegionProxies {
        // Snapshot every inbound-edge proxy for `region` in a single locked
        // expression so the `MutexGuard` is dropped the moment the collect finishes
        // (keeps the significant-`Drop` guard's scope minimal).
        let matching: Vec<(AddrKind, Arc<PartitionProxy>)> = self
            .proxies
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .iter()
            .filter(|(key, _)| key.peer == region)
            .map(|(key, proxy)| (key.kind, Arc::clone(proxy)))
            .collect();
        assert!(
            !matching.is_empty(),
            "no proxies for inbound edges to region '{region}' (was it proxied?)"
        );
        let grpc_links = matching
            .iter()
            .filter(|(kind, _)| *kind == AddrKind::Grpc)
            .map(|(_, proxy)| Arc::clone(proxy))
            .collect();
        let links = matching.into_iter().map(|(_, proxy)| proxy).collect();
        RegionProxies { links, grpc_links }
    }

    /// Recover the gRPC + HTTP proxies for a SINGLE directed `observer → peer` edge,
    /// so the test can sever exactly one direction (the follower↔follower case).
    ///
    /// # Panics
    ///
    /// Panics if either link of the edge was never published (the peer was not
    /// proxied, or the observer does not dial the peer — e.g. observer == peer).
    #[must_use]
    pub fn edge(&self, observer: &str, peer: &str) -> EdgeProxies {
        let map = self
            .proxies
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let get = |kind: AddrKind| {
            map.get(&EdgeKey {
                observer: observer.to_string(),
                peer: peer.to_string(),
                kind,
            })
            .cloned()
            .unwrap_or_else(|| {
                panic!("no {kind:?} proxy for edge '{observer}' -> '{peer}' (was it proxied?)")
            })
        };
        EdgeProxies {
            grpc: get(AddrKind::Grpc),
            http: get(AddrKind::Http),
        }
    }
}

/// Build an [`AddrRewrite`] that interposes a [`PartitionProxy`] on every directed
/// edge whose PEER (target region) is in `proxied_regions`, plus a
/// [`ProxyController`] the test uses to sever/heal those links after start.
///
/// Edges whose peer is NOT listed rewrite to identity (the real address), so only
/// traffic INTO a chosen region is interposed. The rewrite is invoked once per
/// `(observer, peer, kind)` during per-process topology generation; each invocation
/// starts (or reuses) the proxy for that exact directed edge and returns its listen
/// address. Because each observer's topology file is written independently, every
/// directed edge gets its OWN relay — severing one never touches another.
#[must_use]
pub fn proxied_rewrite(proxied_regions: &[&str]) -> (AddrRewrite, ProxyController) {
    let proxies: Arc<Mutex<HashMap<EdgeKey, Arc<PartitionProxy>>>> =
        Arc::new(Mutex::new(HashMap::new()));
    let targets: Arc<Vec<String>> =
        Arc::new(proxied_regions.iter().map(|s| (*s).to_string()).collect());

    let controller = ProxyController {
        proxies: Arc::clone(&proxies),
        targets: Arc::clone(&targets),
    };

    let rewrite_proxies = Arc::clone(&proxies);
    let rewrite_targets = Arc::clone(&targets);
    let rewrite: AddrRewrite = Box::new(
        move |observer: &str, peer: &str, kind: AddrKind, real: SocketAddr| {
            if !rewrite_targets.iter().any(|t| t == peer) {
                return real.to_string(); // peer not proxied: identity
            }
            let key = EdgeKey {
                observer: observer.to_string(),
                peer: peer.to_string(),
                kind,
            };
            // Resolve (or lazily start) this edge's proxy and read its listen addr
            // under the lock, releasing the guard before the closure returns.
            let addr = {
                let mut map = rewrite_proxies
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                map.entry(key)
                    .or_insert_with(|| Arc::new(PartitionProxy::start(real)))
                    .addr()
            };
            addr.to_string()
        },
    );

    (rewrite, controller)
}

/// Fully isolate `region` from `peers`: BOTH directions, gRPC + HTTP. The
/// inbound edges come from [`ProxyController::region`]; the outbound ones are
/// `region`'s own view of each peer ([`ProxyController::edge`]) — without the
/// outbound cut a "partitioned" node keeps sending heartbeats/ships/votes,
/// which defeats partition scenarios entirely (m11p4). Requires every
/// involved region to have been proxied.
pub fn isolate_region(proxies: &ProxyController, region: &str, peers: &[&str]) {
    proxies.region(region).sever_all();
    for peer in peers {
        proxies.edge(region, peer).sever_all();
    }
}

/// Undo [`isolate_region`].
pub fn rejoin_region(proxies: &ProxyController, region: &str, peers: &[&str]) {
    proxies.region(region).heal_all();
    for peer in peers {
        proxies.edge(region, peer).heal_all();
    }
}

/// Spawn one detached unidirectional pump `from → to`. Exits when either end errors
/// (EOF, reset, or a `sever()`-driven `shutdown`), shutting the write side so the
/// partner pump unblocks. Detached: a closed socket always unblocks its blocking
/// `read`, so the thread terminates without an explicit join (and never outlives
/// the process — the streams are owned here and dropped on exit).
fn spawn_pump(from: &TcpStream, to: &TcpStream) {
    let (Ok(mut from), Ok(mut to)) = (from.try_clone(), to.try_clone()) else {
        return;
    };
    thread::spawn(move || {
        let mut buf = [0u8; 16 * 1024];
        loop {
            match from.read(&mut buf) {
                Ok(n) if n > 0 => {
                    if to.write_all(&buf[..n]).is_err() {
                        break;
                    }
                }
                // `Ok(0)` is a clean EOF; `Err(_)` is a reset / shutdown / severed
                // stream. Either way the pump retires.
                Ok(_) | Err(_) => break,
            }
        }
        // Unblock the partner pump: shut the direction we write to.
        let _ = to.shutdown(Shutdown::Both);
        let _ = from.shutdown(Shutdown::Both);
    });
}