//! Shared test helpers: TestServer + TestClient + yamux injection utilities.
#![allow(dead_code)]
//!
//! # Architecture note — data plane
//!
//! The server now has a `/_data/<session_id>` WebSocket endpoint on the
//! control port.  For HTTP / TCP proxy tests we still *inject* yamux streams
//! directly into `TunnelCore::resolve_pending_conn()` to keep tests fast and
//! self-contained (no real WebSocket round-trip needed).

use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use arc_swap::ArcSwap;
use futures_util::future::poll_fn;
use futures_util::io::AsyncReadExt;
use futures_util::{SinkExt, StreamExt};
use rcgen::{CertificateParams, KeyPair};
use rustls::ClientConfig;
#[allow(unused_imports)]
use rustls::ServerConfig as RustlsConfig;
use rustunnel_server::db::Db;
use tempfile::TempDir;
use tokio::sync::mpsc;
use tokio::task::AbortHandle;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{Connector, MaybeTlsStream, WebSocketStream};
use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
use uuid::Uuid;
use yamux::{Config as YamuxConfig, Connection, Mode};

use rustunnel_server::control::mux::WsCompat;

use rustunnel_protocol::{decode_frame, encode_frame, ControlFrame, TunnelProtocol};
use rustunnel_server::config::{
    AuthSection, DatabaseSection, LimitsSection, LoggingSection, RegionSection, ServerConfig,
    ServerSection, TlsSection,
};
use rustunnel_server::core::TunnelCore;
use rustunnel_server::tls::acme::build_tls_config;

// ── TLS helpers ───────────────────────────────────────────────────────────────

/// A `rustls::ClientConfig` that accepts any server certificate.
/// For use in tests only — never in production.
pub fn insecure_client_tls() -> Arc<ClientConfig> {
    use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
    use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
    use rustls::DigitallySignedStruct;

    #[derive(Debug)]
    struct AcceptAnyCert;

    impl ServerCertVerifier for AcceptAnyCert {
        fn verify_server_cert(
            &self,
            _end_entity: &CertificateDer<'_>,
            _intermediates: &[CertificateDer<'_>],
            _server_name: &ServerName<'_>,
            _ocsp_response: &[u8],
            _now: UnixTime,
        ) -> Result<ServerCertVerified, rustls::Error> {
            Ok(ServerCertVerified::assertion())
        }

        fn verify_tls12_signature(
            &self,
            _message: &[u8],
            _cert: &CertificateDer<'_>,
            _dss: &DigitallySignedStruct,
        ) -> Result<HandshakeSignatureValid, rustls::Error> {
            Ok(HandshakeSignatureValid::assertion())
        }

        fn verify_tls13_signature(
            &self,
            _message: &[u8],
            _cert: &CertificateDer<'_>,
            _dss: &DigitallySignedStruct,
        ) -> Result<HandshakeSignatureValid, rustls::Error> {
            Ok(HandshakeSignatureValid::assertion())
        }

        fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
            vec![
                rustls::SignatureScheme::RSA_PKCS1_SHA256,
                rustls::SignatureScheme::RSA_PKCS1_SHA384,
                rustls::SignatureScheme::RSA_PKCS1_SHA512,
                rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
                rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
                rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
                rustls::SignatureScheme::RSA_PSS_SHA256,
                rustls::SignatureScheme::RSA_PSS_SHA384,
                rustls::SignatureScheme::RSA_PSS_SHA512,
                rustls::SignatureScheme::ED25519,
            ]
        }
    }

    Arc::new(
        ClientConfig::builder()
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(AcceptAnyCert))
            .with_no_client_auth(),
    )
}

// ── port selection ────────────────────────────────────────────────────────────

/// Bind port 0 to let the OS assign a free port, then release it.
/// There is a small TOCTOU window; acceptable for single-use ports (control,
/// http, https, dashboard) that are bound within the same runtime tick.
pub fn free_port() -> u16 {
    let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind port 0");
    l.local_addr().unwrap().port()
}

/// Atomically allocate a block of `count` consecutive TCP port numbers for use
/// as a tunnel port range.  Using a process-global counter ensures that
/// parallel tests within the same binary never get overlapping ranges, which
/// would cause "Address already in use" failures in the TCP edge.
pub fn alloc_tcp_port_range(count: u16) -> u16 {
    use std::sync::atomic::{AtomicU16, Ordering};
    // Start high enough to avoid conflicts with well-known services and the
    // server's default range (20000–20099).
    static COUNTER: AtomicU16 = AtomicU16::new(45_000);
    COUNTER.fetch_add(count, Ordering::SeqCst)
}

// ── cert generation ───────────────────────────────────────────────────────────

/// Generate a self-signed cert + key, write PEM files to `dir`, return paths.
pub fn generate_test_cert(dir: &TempDir) -> (String, String) {
    let cert_path = dir.path().join("cert.pem");
    let key_path = dir.path().join("key.pem");

    let params = CertificateParams::new(vec!["localhost".to_string(), "127.0.0.1".to_string()])
        .expect("cert params");
    let key_pair = KeyPair::generate().expect("key pair");
    let cert = params.self_signed(&key_pair).expect("self-signed cert");

    std::fs::write(&cert_path, cert.pem()).expect("write cert");
    std::fs::write(&key_path, key_pair.serialize_pem()).expect("write key");

    (
        cert_path.to_str().unwrap().to_string(),
        key_path.to_str().unwrap().to_string(),
    )
}

// ── TestServer ────────────────────────────────────────────────────────────────

/// Builder-style options for spinning up a `TestServer`. Add new test-only
/// knobs here rather than growing the positional `start_on_ports` arg list.
///
/// Sensible defaults via `TestServerOpts::default()` aren't useful since
/// the port fields would clash across parallel tests; helper constructors
/// (`TestServer::start`, `start_with`, `start_with_load_balancing`)
/// allocate fresh ports and pass them in.
pub struct TestServerOpts {
    pub control_port: u16,
    pub http_port: u16,
    pub https_port: u16,
    pub dashboard_port: u16,
    pub tcp_port_range: [u16; 2],
    pub udp_port_range: [u16; 2],
    pub require_auth: bool,
    pub admin_token: String,
    pub load_balancing_enabled: bool,
    /// Phase 5 alert webhook URL — when `Some`, the server POSTs a
    /// `group_zero_healthy` payload here whenever a load-balancing group
    /// transitions to 0 healthy members.
    pub alert_webhook_url: Option<String>,
    /// Phase 6 cross-version test override — when `Some`, the server
    /// reports this string in `AuthOk.server_version` instead of
    /// `CARGO_PKG_VERSION`. Lets us simulate older edges so we can verify
    /// the client's version-gate falls through to solo registration
    /// without spamming `decode_frame` warnings.
    pub server_version_override: Option<String>,
}

/// A running server instance with all components live on random ports.
pub struct TestServer {
    pub control_port: u16,
    pub http_port: u16,
    pub https_port: u16,
    pub dashboard_port: u16,
    pub domain: String,
    pub admin_token: String,
    pub core: Arc<TunnelCore>,
    pub db: Db,
    pub config: Arc<ServerConfig>,

    // Kept alive for the test's duration.
    _temp_dir: TempDir,
    task_handles: Vec<AbortHandle>,
}

impl TestServer {
    /// Start a server with default settings (require_auth = true, random ports).
    pub async fn start() -> Self {
        Self::start_with(true, "integration-test-token").await
    }

    /// Start a server with configurable auth settings.
    pub async fn start_with(require_auth: bool, admin_token: &str) -> Self {
        let control_port = free_port();
        let http_port = free_port();
        let https_port = free_port();
        let dashboard_port = free_port();
        // Atomically reserve a 10-port range so parallel tests never overlap.
        let tcp_low = alloc_tcp_port_range(10);
        let tcp_high = tcp_low + 9;
        let udp_low = alloc_tcp_port_range(10);
        let udp_high = udp_low + 9;
        Self::start_on_ports(
            control_port,
            http_port,
            https_port,
            dashboard_port,
            [tcp_low, tcp_high],
            [udp_low, udp_high],
            require_auth,
            admin_token,
        )
        .await
    }

    /// Start a server with the TUNNEL-7 load-balancing kill switch on.
    /// Used by the Phase 2 / Phase 3 group integration tests.
    pub async fn start_with_load_balancing() -> Self {
        let control_port = free_port();
        let http_port = free_port();
        let https_port = free_port();
        let dashboard_port = free_port();
        let tcp_low = alloc_tcp_port_range(10);
        let tcp_high = tcp_low + 9;
        let udp_low = alloc_tcp_port_range(10);
        let udp_high = udp_low + 9;
        Self::start_on_ports_with_load_balancing(
            control_port,
            http_port,
            https_port,
            dashboard_port,
            [tcp_low, tcp_high],
            [udp_low, udp_high],
            true,
            "integration-test-token",
            true,
        )
        .await
    }

    /// Start a server on specific pre-allocated ports.
    /// Used by reconnect tests to restart on the same port set.
    #[allow(clippy::too_many_arguments)]
    pub async fn start_on_ports(
        control_port: u16,
        http_port: u16,
        https_port: u16,
        dashboard_port: u16,
        tcp_port_range: [u16; 2],
        udp_port_range: [u16; 2],
        require_auth: bool,
        admin_token: &str,
    ) -> Self {
        Self::start_on_ports_with_load_balancing(
            control_port,
            http_port,
            https_port,
            dashboard_port,
            tcp_port_range,
            udp_port_range,
            require_auth,
            admin_token,
            false,
        )
        .await
    }

    /// Most-flexible entry point: caller sets every port + the load-balancing
    /// kill switch. Existing callers go through `start_on_ports` (which
    /// defaults `load_balancing.enabled` to `false`); only the Phase 2 group
    /// tests flip it on.
    #[allow(clippy::too_many_arguments)]
    pub async fn start_on_ports_with_load_balancing(
        control_port: u16,
        http_port: u16,
        https_port: u16,
        dashboard_port: u16,
        tcp_port_range: [u16; 2],
        udp_port_range: [u16; 2],
        require_auth: bool,
        admin_token: &str,
        load_balancing_enabled: bool,
    ) -> Self {
        Self::start_with_opts(TestServerOpts {
            control_port,
            http_port,
            https_port,
            dashboard_port,
            tcp_port_range,
            udp_port_range,
            require_auth,
            admin_token: admin_token.to_string(),
            load_balancing_enabled,
            alert_webhook_url: None,
            server_version_override: None,
        })
        .await
    }

    /// Builder-style entry point so new test-only knobs (Phase 5 alert
    /// webhook, Phase 6 server-version override, …) can be added without
    /// growing the positional-arg list. Existing helpers wrap this.
    pub async fn start_with_opts(opts: TestServerOpts) -> Self {
        let TestServerOpts {
            control_port,
            http_port,
            https_port,
            dashboard_port,
            tcp_port_range,
            udp_port_range,
            require_auth,
            admin_token,
            load_balancing_enabled,
            alert_webhook_url,
            server_version_override,
        } = opts;
        let admin_token = admin_token.as_str();
        let [tcp_low, tcp_high] = tcp_port_range;

        let temp_dir = TempDir::new().expect("temp dir");
        let (cert_path, key_path) = generate_test_cert(&temp_dir);

        let pg_url = std::env::var("TEST_DATABASE_URL").unwrap_or_else(|_| {
            "postgres://rustunnel:test@localhost:5432/rustunnel_test".to_string()
        });

        let config = Arc::new(ServerConfig {
            server: ServerSection {
                domain: "localhost".to_string(),
                http_port,
                https_port,
                control_port,
                dashboard_port,
                dashboard_origin: "http://localhost:3000".to_string(),
            },
            tls: TlsSection {
                cert_path: cert_path.clone(),
                key_path: key_path.clone(),
                acme_enabled: false,
                acme_email: String::new(),
                acme_staging: false,
                acme_account_dir: temp_dir.path().to_str().unwrap().to_string(),
                cloudflare_api_token: String::new(),
                cloudflare_zone_id: String::new(),
            },
            auth: AuthSection {
                admin_token: admin_token.to_string(),
                require_auth,
            },
            database: DatabaseSection {
                url: pg_url,
                captured_path: ":memory:".to_string(),
            },
            logging: LoggingSection {
                level: "warn".to_string(),
                format: "pretty".to_string(),
                audit_log_path: None,
                sentry_dsn: None,
            },
            limits: LimitsSection {
                max_tunnels_per_session: 10,
                max_connections_per_tunnel: 100,
                rate_limit_rps: 10_000,
                ip_rate_limit_rps: 100_000,
                request_body_max_bytes: 1024 * 1024,
                tcp_port_range: [tcp_low, tcp_high],
                udp_port_range,
            },
            region: RegionSection {
                id: "test".to_string(),
                name: "Test Region".to_string(),
                location: "localhost".to_string(),
            },
            p2p: rustunnel_server::config::P2pSection {
                enabled: true,
                ..rustunnel_server::config::P2pSection::default()
            },
            load_balancing: rustunnel_server::config::LoadBalancingSection {
                enabled: load_balancing_enabled,
                alert_webhook_url: alert_webhook_url.clone(),
            },
            testing: rustunnel_server::config::TestingSection {
                override_server_version: server_version_override.clone(),
            },
        });

        // Database.
        let db = rustunnel_server::db::init_db(&config.database)
            .await
            .expect("init_db");

        // TLS.
        let tls_cfg = build_tls_config(&cert_path, &key_path).expect("build_tls_config");
        let tls_handle = Arc::new(ArcSwap::new(Arc::new(tls_cfg)));
        let tls_snapshot = tls_handle.load_full();

        // Shared tunnel core.
        let core = Arc::new(TunnelCore::new(
            config.limits.tcp_port_range,
            config.limits.udp_port_range,
            config.limits.max_tunnels_per_session,
            config.limits.max_connections_per_tunnel,
            config.limits.ip_rate_limit_rps,
        ));

        let (capture_tx, capture_rx) = mpsc::channel(256);
        let mut task_handles = Vec::new();

        // a) Control plane.
        let control_addr: SocketAddr = format!("127.0.0.1:{control_port}").parse().unwrap();
        let h = tokio::spawn({
            let core = Arc::clone(&core);
            let cfg = Arc::clone(&config);
            let tls_handle = Arc::clone(&tls_handle);
            let db = db.clone();
            async move {
                let _ = rustunnel_server::control::server::run_control_plane(
                    control_addr,
                    core,
                    cfg,
                    tls_handle,
                    rustunnel_server::audit::noop_audit(),
                    db,
                )
                .await;
            }
        });
        task_handles.push(h.abort_handle());

        // b) HTTP + HTTPS edge.
        let http_addr: SocketAddr = format!("127.0.0.1:{http_port}").parse().unwrap();
        let https_addr: SocketAddr = format!("127.0.0.1:{https_port}").parse().unwrap();
        let h = tokio::spawn({
            let core = Arc::clone(&core);
            let domain = config.server.domain.clone();
            let limits = rustunnel_server::edge::HttpEdgeConfig {
                rate_limit_rps: config.limits.rate_limit_rps,
                request_body_max_bytes: config.limits.request_body_max_bytes,
            };
            async move {
                let _ = rustunnel_server::edge::run_http_edge(
                    http_addr,
                    https_addr,
                    tls_snapshot,
                    core,
                    domain,
                    Some(capture_tx),
                    limits,
                )
                .await;
            }
        });
        task_handles.push(h.abort_handle());

        // c) TCP edge.
        let h = tokio::spawn({
            let core = Arc::clone(&core);
            async move {
                rustunnel_server::edge::run_tcp_edge(core).await;
            }
        });
        task_handles.push(h.abort_handle());

        // d) Dashboard.
        let dashboard_addr: SocketAddr = format!("127.0.0.1:{dashboard_port}").parse().unwrap();
        let h = tokio::spawn({
            let core = Arc::clone(&core);
            let db_dash = db.clone();
            let admin_token = config.auth.admin_token.clone();
            async move {
                let _ = rustunnel_server::dashboard::run_dashboard(
                    dashboard_addr,
                    core,
                    db_dash,
                    capture_rx,
                    admin_token,
                    rustunnel_server::audit::noop_audit(),
                    "http://localhost:3000".to_string(),
                    rustunnel_server::config::RegionSection {
                        id: "test".to_string(),
                        name: "Test Region".to_string(),
                        location: "localhost".to_string(),
                    },
                )
                .await;
            }
        });
        task_handles.push(h.abort_handle());

        // Wait until every public listener is actually accepting TCP
        // connections — control plane, HTTP edge, HTTPS edge, dashboard.
        // A fixed sleep is unreliable on CI; polling eliminates the race.
        // We poll all four ports because the four background tasks bind
        // their listeners independently. On a slow runner the test
        // client can outrun the HTTPS bind and get ECONNREFUSED if we
        // only wait for the control port.
        let deadline = std::time::Instant::now() + Duration::from_secs(10);
        for (label, port) in [
            ("control plane", control_port),
            ("http edge", http_port),
            ("https edge", https_port),
            ("dashboard", dashboard_port),
        ] {
            loop {
                if tokio::net::TcpStream::connect(format!("127.0.0.1:{port}"))
                    .await
                    .is_ok()
                {
                    break;
                }
                assert!(
                    std::time::Instant::now() < deadline,
                    "{label} did not start within 10 seconds"
                );
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        }

        Self {
            control_port,
            http_port,
            https_port,
            dashboard_port,
            domain: "localhost".to_string(),
            admin_token: admin_token.to_string(),
            core,
            db,
            config,
            _temp_dir: temp_dir,
            task_handles,
        }
    }

    /// Stop the server by aborting all tasks.
    pub fn stop(&self) {
        for h in &self.task_handles {
            h.abort();
        }
    }
}

impl Drop for TestServer {
    fn drop(&mut self) {
        self.stop();
    }
}

// ── TestClient ────────────────────────────────────────────────────────────────

type CtrlWs = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;

/// A test control-plane client that speaks the binary WebSocket protocol
/// directly (no CLI, no reconnect logic).
pub struct TestClient {
    ws: CtrlWs,
    pub session_id: Option<Uuid>,
    /// `server_version` field captured from `AuthOk`. Lets Phase 6 tests
    /// verify the server's `[testing] override_server_version` knob is
    /// taking effect — the client's version-gate logic itself is
    /// unit-tested in `rustunnel-client::version::tests`.
    pub server_version: Option<String>,
}

impl TestClient {
    /// Connect and authenticate with `token`.
    pub async fn connect_with_token(server: &TestServer, token: &str) -> Result<Self, String> {
        let url = format!("wss://127.0.0.1:{}/_control", server.control_port);

        let connector = Connector::Rustls(insecure_client_tls());
        let (ws, _) =
            tokio_tungstenite::connect_async_tls_with_config(&url, None, false, Some(connector))
                .await
                .map_err(|e| format!("WS connect: {e}"))?;

        let mut client = Self {
            ws,
            session_id: None,
            server_version: None,
        };

        // Send Auth frame.
        client
            .send(&ControlFrame::Auth {
                token: token.to_string(),
                client_version: "test".to_string(),
            })
            .await?;

        // Receive response.
        match client.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::AuthOk {
                session_id,
                server_version,
            } => {
                client.session_id = Some(session_id);
                client.server_version = Some(server_version);
                Ok(client)
            }
            ControlFrame::AuthError { message } => Err(format!("AuthError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Connect with the server's admin token.
    pub async fn connect(server: &TestServer) -> Result<Self, String> {
        Self::connect_with_token(server, &server.admin_token).await
    }

    /// Connect and expect the auth to fail, returning the error message.
    pub async fn connect_expect_auth_error(
        server: &TestServer,
        token: &str,
    ) -> Result<String, String> {
        let url = format!("wss://127.0.0.1:{}/_control", server.control_port);
        let connector = Connector::Rustls(insecure_client_tls());
        let (ws, _) =
            tokio_tungstenite::connect_async_tls_with_config(&url, None, false, Some(connector))
                .await
                .map_err(|e| format!("WS connect: {e}"))?;

        let mut client = Self {
            ws,
            session_id: None,
            server_version: None,
        };

        client
            .send(&ControlFrame::Auth {
                token: token.to_string(),
                client_version: "test".to_string(),
            })
            .await?;

        match client.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::AuthError { message } => Ok(message),
            ControlFrame::AuthOk { .. } => Err("expected AuthError but got AuthOk".into()),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    // ── tunnel registration ───────────────────────────────────────────────────

    /// Send `TunnelHealthy { tunnel_id }` over the control WS.
    /// Phase 4 server-side handler accepts this only from the owning
    /// session — that's why this lives on `TestClient` rather than as a
    /// free function.
    pub async fn send_tunnel_healthy(&mut self, tunnel_id: Uuid) -> Result<(), String> {
        self.send(&ControlFrame::TunnelHealthy { tunnel_id }).await
    }

    /// Send `TunnelUnhealthy { tunnel_id, reason }` over the control WS.
    pub async fn send_tunnel_unhealthy(
        &mut self,
        tunnel_id: Uuid,
        reason: &str,
    ) -> Result<(), String> {
        self.send(&ControlFrame::TunnelUnhealthy {
            tunnel_id,
            reason: reason.to_string(),
        })
        .await
    }

    /// Register an HTTP tunnel as part of a load-balancing group, with
    /// an optional per-tenant alert webhook URL. The webhook URL is
    /// part of `HealthCheckSpec` on the wire, so this helper attaches
    /// a minimal TCP probe spec carrying just the URL — Phase 5 fan-out
    /// tests use this to drive the server's per-tenant alert path.
    pub async fn register_http_tunnel_grouped_with_alert(
        &mut self,
        subdomain: Option<&str>,
        group_name: &str,
        group_key_hash: &str,
        alert_webhook_url: Option<&str>,
    ) -> Result<(Uuid, String, String), String> {
        let health_check = alert_webhook_url.map(|url| rustunnel_protocol::HealthCheckSpec {
            kind: rustunnel_protocol::HealthCheckKind::Tcp,
            interval_secs: 10,
            timeout_secs: 3,
            max_failed: 3,
            http_path: None,
            http_expect_2xx: true,
            alert_webhook_url: Some(url.to_string()),
        });
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id.clone(),
            protocol: TunnelProtocol::Http,
            subdomain: subdomain.map(str::to_string),
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: None,
            p2p_name: None,
            group: Some(group_name.to_string()),
            group_key_hash: Some(group_key_hash.to_string()),
            health_check,
        })
        .await?;
        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                public_url,
                ..
            } => {
                let sub = public_url
                    .trim_start_matches("http://")
                    .split('.')
                    .next()
                    .unwrap_or("")
                    .to_string();
                Ok((tunnel_id, sub, public_url))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register an HTTP tunnel as part of a load-balancing group.
    /// Phase 2 of TUNNEL-7. The server only honours the group fields when
    /// `[load_balancing] enabled = true`. Returns `(tunnel_id, subdomain, public_url)`.
    pub async fn register_http_tunnel_grouped(
        &mut self,
        subdomain: Option<&str>,
        group_name: &str,
        group_key_hash: &str,
    ) -> Result<(Uuid, String, String), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id.clone(),
            protocol: TunnelProtocol::Http,
            subdomain: subdomain.map(str::to_string),
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: None,
            p2p_name: None,
            group: Some(group_name.to_string()),
            group_key_hash: Some(group_key_hash.to_string()),
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                public_url,
                ..
            } => {
                let sub = public_url
                    .trim_start_matches("http://")
                    .split('.')
                    .next()
                    .unwrap_or("")
                    .to_string();
                Ok((tunnel_id, sub, public_url))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register an HTTP tunnel.  Returns `(tunnel_id, subdomain, public_url)`.
    pub async fn register_http_tunnel(
        &mut self,
        subdomain: Option<&str>,
    ) -> Result<(Uuid, String, String), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id.clone(),
            protocol: TunnelProtocol::Http,
            subdomain: subdomain.map(str::to_string),
            local_addr: "127.0.0.1:0".to_string(), // advisory only
            p2p_secret_hash: None,
            p2p_name: None,
            group: None,
            group_key_hash: None,
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                public_url,
                ..
            } => {
                // Derive the subdomain from the public_url.
                let sub = public_url
                    .trim_start_matches("http://")
                    .split('.')
                    .next()
                    .unwrap_or("")
                    .to_string();
                Ok((tunnel_id, sub, public_url))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register a TCP tunnel as part of a load-balancing group.
    /// Phase 3 of TUNNEL-7. The server returns the group's existing port
    /// for the second-and-onwards member; the first allocates a fresh one.
    /// Returns `(tunnel_id, assigned_port)`.
    pub async fn register_tcp_tunnel_grouped(
        &mut self,
        group_name: &str,
        group_key_hash: &str,
    ) -> Result<(Uuid, u16), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id,
            protocol: TunnelProtocol::Tcp,
            subdomain: None,
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: None,
            p2p_name: None,
            group: Some(group_name.to_string()),
            group_key_hash: Some(group_key_hash.to_string()),
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                assigned_port,
                ..
            } => {
                let port = assigned_port.ok_or("missing assigned_port")?;
                Ok((tunnel_id, port))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register a TCP tunnel.  Returns `(tunnel_id, assigned_port)`.
    pub async fn register_tcp_tunnel(&mut self) -> Result<(Uuid, u16), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id,
            protocol: TunnelProtocol::Tcp,
            subdomain: None,
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: None,
            p2p_name: None,
            group: None,
            group_key_hash: None,
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                assigned_port,
                ..
            } => {
                let port = assigned_port.ok_or("missing assigned_port")?;
                Ok((tunnel_id, port))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register a UDP tunnel.  Returns `(tunnel_id, assigned_port)`.
    pub async fn register_udp_tunnel(&mut self) -> Result<(Uuid, u16), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id,
            protocol: TunnelProtocol::Udp,
            subdomain: None,
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: None,
            p2p_name: None,
            group: None,
            group_key_hash: None,
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                assigned_port,
                ..
            } => {
                let port = assigned_port.ok_or("missing assigned_port")?;
                Ok((tunnel_id, port))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Register a P2P publisher tunnel.  Returns `(tunnel_id, name)`.
    pub async fn register_p2p_tunnel(
        &mut self,
        name: &str,
        secret_hash: &str,
    ) -> Result<(Uuid, String), String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::RegisterTunnel {
            request_id: req_id,
            protocol: TunnelProtocol::P2p,
            subdomain: None,
            local_addr: "127.0.0.1:0".to_string(),
            p2p_secret_hash: Some(secret_hash.to_string()),
            p2p_name: Some(name.to_string()),
            group: None,
            group_key_hash: None,
            health_check: None,
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(5)).await? {
            ControlFrame::TunnelRegistered {
                tunnel_id,
                p2p_tunnel_name,
                ..
            } => {
                let tname = p2p_tunnel_name.ok_or("missing p2p_tunnel_name")?;
                Ok((tunnel_id, tname))
            }
            ControlFrame::TunnelError { message, .. } => Err(format!("TunnelError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    /// Send a P2pConnect request. Returns the conn_id on success.
    pub async fn p2p_connect(&mut self, target: &str, secret_hash: &str) -> Result<Uuid, String> {
        let req_id = Uuid::new_v4().to_string();
        self.send(&ControlFrame::P2pConnect {
            request_id: req_id,
            target_tunnel_name: target.to_string(),
            secret_hash: secret_hash.to_string(),
        })
        .await?;

        match self.recv_timeout(Duration::from_secs(10)).await? {
            ControlFrame::P2pConnected { conn_id, .. } => Ok(conn_id),
            ControlFrame::P2pError { message, .. } => Err(format!("P2pError: {message}")),
            other => Err(format!("unexpected frame: {other:?}")),
        }
    }

    // ── frame I/O ─────────────────────────────────────────────────────────────

    pub async fn send(&mut self, frame: &ControlFrame) -> Result<(), String> {
        let bytes = encode_frame(frame);
        self.ws
            .send(Message::Binary(bytes))
            .await
            .map_err(|e| format!("send: {e}"))
    }

    pub async fn recv_timeout(&mut self, timeout: Duration) -> Result<ControlFrame, String> {
        let msg = tokio::time::timeout(timeout, self.ws.next())
            .await
            .map_err(|_| "recv timeout".to_string())?
            .ok_or("WS closed")?
            .map_err(|e| format!("recv: {e}"))?;

        match msg {
            Message::Binary(data) => decode_frame(&data).map_err(|e| format!("decode: {e}")),
            other => Err(format!("expected binary frame, got {other:?}")),
        }
    }

    /// Wait for the next `NewConnection` frame (up to 10 s).
    pub async fn wait_new_connection(&mut self) -> Result<Uuid, String> {
        let deadline = Duration::from_secs(10);
        match self.recv_timeout(deadline).await? {
            ControlFrame::NewConnection { conn_id, .. } => Ok(conn_id),
            ControlFrame::Ping { timestamp } => {
                // Reply with Pong and retry.
                self.send(&ControlFrame::Pong { timestamp }).await?;
                match self.recv_timeout(deadline).await? {
                    ControlFrame::NewConnection { conn_id, .. } => Ok(conn_id),
                    other => Err(format!("expected NewConnection, got {other:?}")),
                }
            }
            other => Err(format!("expected NewConnection, got {other:?}")),
        }
    }
}

// ── data WebSocket bridge ─────────────────────────────────────────────────────

/// Connect a `/_data/<session_id>` WebSocket to the server and run a yamux
/// bridge task that handles all proxy connections automatically.
///
/// The server runs yamux `Mode::Client` and opens one outbound stream per
/// incoming HTTP/TCP connection.  It writes a 16-byte conn_id prefix to each
/// stream so the client can route it.  This function connects a `Mode::Server`
/// yamux connection, accepts those inbound streams, reads the conn_id, and
/// bidirectionally copies data between the stream and `local_addr`.
///
/// Returns a `oneshot::Receiver<()>` that resolves once the WebSocket
/// handshake completes (the bridge is ready to handle streams).  Await it
/// in tests instead of a fixed sleep.
pub fn connect_data_bridge(
    server: &TestServer,
    session_id: Uuid,
    local_addr: SocketAddr,
) -> tokio::sync::oneshot::Receiver<()> {
    let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
    let url = format!(
        "wss://127.0.0.1:{}/_data/{}",
        server.control_port, session_id
    );
    let connector = Connector::Rustls(insecure_client_tls());

    tokio::spawn(async move {
        let (ws, _) = match tokio_tungstenite::connect_async_tls_with_config(
            &url,
            None,
            false,
            Some(connector),
        )
        .await
        {
            Ok(pair) => pair,
            Err(e) => {
                eprintln!("[test data bridge] WS connect failed: {e}");
                return;
            }
        };

        // Signal readiness immediately after WS handshake, before entering
        // the yamux accept loop.  The caller awaits this to know the bridge
        // is live and can route incoming streams.
        let _ = ready_tx.send(());

        // yamux Mode::Server: the server (Mode::Client) opens outbound streams.
        let ws_compat = WsCompat::new(ws);
        let mut conn = Connection::new(ws_compat, YamuxConfig::default(), Mode::Server);

        loop {
            match poll_fn(|cx| conn.poll_next_inbound(cx)).await {
                Some(Ok(mut stream)) => {
                    // Read the 16-byte conn_id prefix the server wrote.
                    let mut id_bytes = [0u8; 16];
                    if stream.read_exact(&mut id_bytes).await.is_err() {
                        continue;
                    }
                    let _conn_id = Uuid::from_bytes(id_bytes);

                    // Bridge the yamux stream to the local service.
                    tokio::spawn(async move {
                        match tokio::net::TcpStream::connect(local_addr).await {
                            Ok(mut local) => {
                                let mut remote = stream.compat();
                                let _ =
                                    tokio::io::copy_bidirectional(&mut local, &mut remote).await;
                            }
                            Err(e) => {
                                eprintln!("[test data bridge] connect {local_addr}: {e}");
                            }
                        }
                    });
                }
                Some(Err(e)) => {
                    eprintln!("[test data bridge] yamux error: {e}");
                    break;
                }
                None => break,
            }
        }
    });

    ready_rx
}

/// Like `connect_data_bridge` but also returns an `AbortHandle` so the caller
/// can forcibly kill the data WebSocket (simulating a NAT timeout or network
/// drop) without touching the control WS.
pub fn connect_data_bridge_abortable(
    server: &TestServer,
    session_id: Uuid,
    local_addr: SocketAddr,
) -> (tokio::sync::oneshot::Receiver<()>, AbortHandle) {
    let (ready_tx, ready_rx) = tokio::sync::oneshot::channel::<()>();
    let url = format!(
        "wss://127.0.0.1:{}/_data/{}",
        server.control_port, session_id
    );
    let connector = Connector::Rustls(insecure_client_tls());

    let handle = tokio::spawn(async move {
        let (ws, _) = match tokio_tungstenite::connect_async_tls_with_config(
            &url,
            None,
            false,
            Some(connector),
        )
        .await
        {
            Ok(pair) => pair,
            Err(e) => {
                eprintln!("[test data bridge] WS connect failed: {e}");
                return;
            }
        };

        let _ = ready_tx.send(());

        let ws_compat = WsCompat::new(ws);
        let mut conn = Connection::new(ws_compat, YamuxConfig::default(), Mode::Server);

        loop {
            match poll_fn(|cx| conn.poll_next_inbound(cx)).await {
                Some(Ok(mut stream)) => {
                    let mut id_bytes = [0u8; 16];
                    if stream.read_exact(&mut id_bytes).await.is_err() {
                        continue;
                    }
                    let _conn_id = Uuid::from_bytes(id_bytes);

                    tokio::spawn(async move {
                        match tokio::net::TcpStream::connect(local_addr).await {
                            Ok(mut local) => {
                                let mut remote = stream.compat();
                                let _ =
                                    tokio::io::copy_bidirectional(&mut local, &mut remote).await;
                            }
                            Err(e) => {
                                eprintln!("[test data bridge] connect {local_addr}: {e}");
                            }
                        }
                    });
                }
                Some(Err(e)) => {
                    eprintln!("[test data bridge] yamux error: {e}");
                    break;
                }
                None => break,
            }
        }
    });

    (ready_rx, handle.abort_handle())
}

// ── yamux stream injection (legacy, kept for reference) ──────────────────────

/// Yamux 0.13 uses a lazy SYN: the SYN frame is only sent when the opener
/// writes the first byte.  Therefore we cannot open a stream on Side A and
/// immediately accept it on Side B without anyone writing.
///
/// Solution: the SERVER-mode connection opens the stream (which the HTTP/TCP
/// edge receives via `resolve_pending_conn`).  When the edge writes its first
/// bytes (HTTP request or TCP data), the DATA+SYN frame travels through the
/// in-memory duplex to the CLIENT-mode connection, which then accepts it as
/// an inbound stream.  That inbound stream is sent to the proxy task via a
/// oneshot channel.
///
/// Returns `(edge_stream, proxy_rx)`:
/// * `edge_stream` — inject into `core.resolve_pending_conn()`.
/// * `proxy_rx` — resolves to the proxy-side stream *after* the edge writes
///   its first byte; the proxy task awaits this before bridging
///   to the local service.
pub async fn make_yamux_pair() -> (yamux::Stream, tokio::sync::oneshot::Receiver<yamux::Stream>) {
    // Large buffer so a full HTTP response fits without backpressure.
    let (io_a, io_b) = tokio::io::duplex(2 * 1024 * 1024);

    let (edge_tx, edge_rx) = tokio::sync::oneshot::channel::<yamux::Stream>();
    let (proxy_tx, proxy_rx) = tokio::sync::oneshot::channel::<yamux::Stream>();

    // Server driver: opens one outbound stream (the edge will get this stream
    // and write DATA+SYN through it).  Keepalive loop processes ACKs and data.
    tokio::spawn(async move {
        let mut conn = Connection::new(io_b.compat(), YamuxConfig::default(), Mode::Server);
        let stream = poll_fn(|cx| conn.poll_new_outbound(cx))
            .await
            .expect("yamux server open outbound");
        let _ = edge_tx.send(stream);
        // Keep driving so stream commands (DATA+SYN from edge writes) are
        // flushed to io_b and ACKs from the client side are processed.
        loop {
            match tokio::time::timeout(
                Duration::from_secs(30),
                poll_fn(|cx| conn.poll_next_inbound(cx)),
            )
            .await
            {
                Ok(None) | Ok(Some(Err(_))) => break,
                Ok(Some(Ok(_))) | Err(_) => {}
            }
        }
    });

    // Client driver: waits for the DATA+SYN frame that arrives when the edge
    // writes its first byte.  Accepts the inbound stream and hands it to the
    // proxy task.
    tokio::spawn(async move {
        let mut conn = Connection::new(io_a.compat(), YamuxConfig::default(), Mode::Client);
        // This await returns only after the server side sends DATA+SYN.
        // connection closed before a stream arrived → ignore
        if let Some(Ok(stream)) = poll_fn(|cx| conn.poll_next_inbound(cx)).await {
            let _ = proxy_tx.send(stream);
        }
        // Keep driving for subsequent data frames.
        loop {
            match tokio::time::timeout(
                Duration::from_secs(30),
                poll_fn(|cx| conn.poll_next_inbound(cx)),
            )
            .await
            {
                Ok(None) | Ok(Some(Err(_))) => break,
                Ok(Some(Ok(_))) | Err(_) => {}
            }
        }
    });

    let edge_stream = edge_rx
        .await
        .expect("yamux server driver died before opening stream");
    (edge_stream, proxy_rx)
}

/// Inject a proxy path for a single `NewConnection`:
/// 1. Creates a yamux pair (edge_stream + lazy proxy_rx).
/// 2. Delivers `edge_stream` to the waiting edge task via `resolve_pending_conn`.
/// 3. Spawns a proxy task that waits for the proxy-side stream (appears once
///    the edge writes its first byte), then bridges it to `local_addr`.
pub async fn inject_proxy(core: &Arc<TunnelCore>, conn_id: Uuid, local_addr: SocketAddr) {
    let (edge_stream, proxy_rx) = make_yamux_pair().await;

    // Deliver the stream to the waiting edge task.
    core.resolve_pending_conn(&conn_id, edge_stream);

    // Proxy: wait for the client-side yamux stream (arrives after edge's first
    // write), then bridge bidirectionally to the local service.
    tokio::spawn(async move {
        use tokio_util::compat::FuturesAsyncReadCompatExt;

        let proxy_stream = match proxy_rx.await {
            Ok(s) => s,
            Err(_) => {
                eprintln!("[test proxy] yamux proxy stream never appeared");
                return;
            }
        };
        let mut local = match tokio::net::TcpStream::connect(local_addr).await {
            Ok(s) => s,
            Err(e) => {
                eprintln!("[test proxy] connect {local_addr}: {e}");
                return;
            }
        };
        let mut remote = proxy_stream.compat();
        let _ = tokio::io::copy_bidirectional(&mut local, &mut remote).await;
    });
}

// ── reqwest helper ────────────────────────────────────────────────────────────

/// Build a `reqwest::Client` that accepts any TLS certificate.
pub fn insecure_http_client() -> reqwest::Client {
    reqwest::Client::builder()
        .danger_accept_invalid_certs(true)
        .timeout(Duration::from_secs(10))
        .build()
        .expect("reqwest client")
}

// ── tracing (one-time init) ───────────────────────────────────────────────────

/// Install a quiet tracing subscriber (warn level) if not already installed.
/// Also installs the rustls crypto provider (aws-lc-rs) once per process.
pub fn init_tracing() {
    // rustls requires a process-level CryptoProvider when both aws-lc-rs and
    // ring are present in the dependency graph.  Install it idempotently.
    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();

    use tracing_subscriber::{fmt, prelude::*, EnvFilter};
    let _ = tracing_subscriber::registry()
        .with(EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")))
        .with(fmt::layer().with_test_writer())
        .try_init();
}
