pub mod coprocessor;
pub mod mock_subgraphs;
pub mod otel;
pub mod s3_mock;
pub mod stdout;

use axum_server::{tls_rustls::RustlsConfig, Handle};
use bytes::Bytes;
use dashmap::DashMap;
use hive_router_plan_executor::plugin_trait::RouterPlugin;
use lazy_static::lazy_static;
use mockito::Mock;
use ntex::{
    client::ClientResponse,
    io::Sealed,
    time::Seconds,
    web::{self, test},
    ws::WsConnection,
};
use rcgen::generate_simple_self_signed;
use reqwest::header::{ACCEPT, CONTENT_TYPE};
use sonic_rs::json;
use std::{
    any::Any,
    future::Future,
    io::Write,
    marker::PhantomData,
    net::SocketAddr,
    path::PathBuf,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::{Duration, Instant},
};
use tempfile::{NamedTempFile, TempPath};
use tokio::{sync::Semaphore, time};
use tracing::{info, warn};

use hive_router::{
    add_callback_handler, background_tasks::BackgroundTasksManager, configure_app_from_config,
    configure_ntex_app, init_rustls_crypto_provider, invoke_shutdown_hooks,
    pipeline::long_lived_client_limit::LongLivedClientLimitService,
    pipeline::request_identifiers::RequestIdentifiersService,
    pipeline::request_summary::RequestSummaryService, plugins::plugins_service::PluginService,
    telemetry::Telemetry, PluginRegistry, RouterPaths, RouterSharedState, SchemaState,
};
use hive_router_config::{
    load_config, parse_yaml_config, subscriptions::CallbackConfig, HiveRouterConfig,
};
use hive_router_plan_executor::executors::websocket_client;
use subgraphs::{subgraphs_app, HTTPStreamingSubscriptionProtocol};

/// Binds a TCP listener to an OS-assigned port and returns that port number.
/// The listener is immediately dropped, so the port is free for the caller to use.
pub fn get_available_port() -> u16 {
    let listener =
        std::net::TcpListener::bind("127.0.0.1:0").expect("failed to bind to get available port");
    listener
        .local_addr()
        .expect("failed to get local address")
        .port()
}

/// Creates a Some(http::HeaderMap) from a list of key-value pairs, for use in test requests.
#[macro_export]
macro_rules! some_header_map {
    ($($key:expr => $val:expr),* $(,)?) => {{
        let mut map = ::http::HeaderMap::new();
        $(map.insert($key, $val.parse().expect("failed to parse header value"));)*
        Some(map)
    }};
}

// #[macro_export] always hoists to the crate root so we re-export it here module level
pub use some_header_map;

/// Replaces the subgraphs address in the given supergraph string with the test
/// subgraphs address and returns the modified supergraph.
///
/// It will replace all occurrences of `0.0.0.0:4200` with the test subgraphs address.
pub fn supergraph_with_subgraphs(supergraph: &str, subgraphs: &str) -> String {
    supergraph.replace("http://0.0.0.0:4200", subgraphs)
}

/// Creates a temporary supergraph file with the content of the given file but with the subgraphs
/// address replaced with the test subgraphs address.
///
/// The temp file will be automatically deleted when the returned TempPath is dropped.
pub fn supergraph_temp_file_with_subgraphs(supergraph_file: &str, subgraphs: &str) -> TempPath {
    let original =
        std::fs::read_to_string(supergraph_file).expect("failed to read supergraph file");
    let with_addr = supergraph_with_subgraphs(&original, subgraphs);

    let temp_file =
        NamedTempFile::with_suffix(".graphql").expect("failed to create temp supergraph file");
    std::fs::write(temp_file.path(), with_addr).expect("failed to write temp supergraph file");

    // close the file handle but keep the path for cleanup on drop
    // useful when running many tests in parallel to avoid hitting the open file limit
    let temp_path = temp_file.into_temp_path();

    info!(
        "Using supergraph at {} to use test subgraphs with address {}",
        temp_path
            .to_str()
            .expect("failed to convert temp path to string"),
        subgraphs.to_string()
    );

    temp_path
}

lazy_static! {
    /// Ensures only one `EnvVarsGuard` exists at a time, preventing concurrent mutation of
    /// environment variables (which are global process state and not thread-safe to modify).
    static ref ENV_VAR_SEMAPHORE: Arc<Semaphore> = Arc::new(Semaphore::new(1));
}

/// A guard that sets one or more environment variables and restores their original values (or
/// removes them) when dropped. Only one instance may exist at a time across all threads;
/// `apply()` is async and blocks until any previous guard has been dropped.
///
/// Usage: `EnvVarsGuard::new().set("key", "value").set("key2", "value2").apply().await`
pub struct EnvVarsGuard {
    pending: Vec<(String, String)>,
    vars: Vec<(String, Option<String>)>,
    permit: Option<tokio::sync::OwnedSemaphorePermit>,
}

impl Default for EnvVarsGuard {
    fn default() -> Self {
        Self::new()
    }
}

impl EnvVarsGuard {
    pub fn new() -> Self {
        EnvVarsGuard {
            pending: vec![],
            vars: vec![],
            permit: None,
        }
    }

    pub fn set(mut self, key: &str, value: &str) -> Self {
        self.pending.push((key.to_string(), value.to_string()));
        self
    }

    /// Applies the pending environment variable changes, returning a guard that
    /// will restore them on drop. This method is async and will block until any
    /// previous guard has been dropped to ensure that environment variable mutations
    /// are not done concurrently.
    pub async fn apply(mut self) -> Self {
        self.permit = Some(
            Arc::clone(&ENV_VAR_SEMAPHORE)
                .acquire_owned()
                .await
                .expect("env var semaphore closed"),
        );

        self.vars = self
            .pending
            .iter()
            .map(|(key, value)| {
                let original = std::env::var(key).ok();
                // SAFETY: environment variables are global state; we serialise all mutations
                // through ENV_VAR_SEMAPHORE so only one guard can set/restore vars at a time.
                unsafe { std::env::set_var(key, value) };
                (key.to_string(), original)
            })
            .collect();

        self
    }
}

impl Drop for EnvVarsGuard {
    fn drop(&mut self) {
        for (key, original) in &self.vars {
            // SAFETY: same as in `apply`; the permit is still held here and released after drop.
            unsafe {
                match original {
                    Some(value) => std::env::set_var(key, value),
                    None => std::env::remove_var(key),
                }
            }
        }
    }
}

// state markers

pub struct Built;
pub struct Started;

// subgraphs

#[derive(Clone)]
pub struct RequestLike {
    #[allow(unused)]
    pub path: String,
    pub headers: http::HeaderMap,
    #[allow(unused)]
    pub body: Option<Bytes>,
    pub http_version: http::Version,
}

pub struct ResponseLike {
    pub status: axum::http::StatusCode,
    pub headers: http::HeaderMap,
    pub body: Option<Bytes>,
}

impl ResponseLike {
    #[allow(unused)]
    pub fn new(
        status: axum::http::StatusCode,
        body: Option<String>,
        headers: Option<http::HeaderMap>,
    ) -> Self {
        Self {
            status,
            headers: headers.unwrap_or_else(http::HeaderMap::new),
            body: body.map(Bytes::from),
        }
    }
}

type OnRequest = dyn Fn(RequestLike) -> Option<ResponseLike> + Send + Sync;

pub struct TestSubgraphsBuilder {
    subscriptions_protocol: HTTPStreamingSubscriptionProtocol,
    on_request: Option<Arc<OnRequest>>,
    rustls_config: Option<RustlsConfig>,
    delay: Option<Duration>,
    http2_only: bool,
}

impl TestSubgraphsBuilder {
    pub fn new() -> Self {
        Self {
            on_request: None,
            rustls_config: None,
            delay: None,
            subscriptions_protocol: HTTPStreamingSubscriptionProtocol::default(),
            http2_only: false,
        }
    }

    pub fn with_http_streaming_subscriptions_protocol(
        mut self,
        protocol: HTTPStreamingSubscriptionProtocol,
    ) -> Self {
        self.subscriptions_protocol = protocol;
        self
    }

    pub fn with_on_request(
        mut self,
        on_request: impl Fn(RequestLike) -> Option<ResponseLike> + Send + Sync + 'static,
    ) -> Self {
        self.on_request = Some(Arc::new(on_request));
        self
    }

    #[allow(unused)]
    pub fn with_rustls_config(mut self, rustls_config: RustlsConfig) -> Self {
        self.rustls_config = Some(rustls_config);
        self
    }

    /// Adds a cooperative async delay to every subgraph request.
    /// Unlike `with_on_request` with `std::thread::sleep`, this yields
    /// back to the tokio runtime, allowing other tasks (like schema
    /// pollers) to make progress during the delay.
    #[allow(unused)]
    pub fn with_delay(mut self, delay: Duration) -> Self {
        self.delay = Some(delay);
        self
    }

    /// Enables HTTP/2 only mode (h2c) for the test subgraph server.
    /// When enabled, the server will only accept HTTP/2 connections over plain TCP.
    #[allow(unused)]
    pub fn with_http2_only(mut self) -> Self {
        self.http2_only = true;
        self
    }

    pub fn build(self) -> TestSubgraphs<Built> {
        TestSubgraphs {
            on_request: self.on_request,
            rustls_config: self.rustls_config,
            delay: self.delay,
            subscriptions_protocol: self.subscriptions_protocol,
            http2_only: self.http2_only,
            handle: None,
            _state: PhantomData,
        }
    }
}

impl Default for TestSubgraphsBuilder {
    fn default() -> Self {
        Self::new()
    }
}

struct TestSubgraphsHandle {
    server_handle: Handle<SocketAddr>,
    addr: SocketAddr,
    state: Arc<TestSubgraphsMiddlewareState>,
    active_subscriptions: Arc<AtomicUsize>,
}

pub struct TestSubgraphs<State> {
    subscriptions_protocol: HTTPStreamingSubscriptionProtocol,
    on_request: Option<Arc<OnRequest>>,
    rustls_config: Option<RustlsConfig>,
    delay: Option<Duration>,
    http2_only: bool,
    handle: Option<TestSubgraphsHandle>,
    _state: PhantomData<State>,
}

struct TestSubgraphsMiddlewareState {
    /// A map of subgraph name to list of requests received on that subgraph.
    request_log: DashMap<String, Vec<RequestLike>>,
}

async fn record_requests(
    axum::extract::State(state): axum::extract::State<Arc<TestSubgraphsMiddlewareState>>,
    request: axum::extract::Request,
    next: axum::middleware::Next,
) -> impl axum::response::IntoResponse {
    let path = request.uri().path().to_string();
    let subgraph = path
        .trim_start_matches("/") // remove leading slash to have the path represent the subgraph
        .to_string();
    let (parts, body) = request.into_parts();
    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();

    let header_map = parts.headers.clone();
    let http_version = parts.version;
    let record = RequestLike {
        path,
        headers: header_map,
        body: if body_bytes.is_empty() {
            None
        } else {
            Some(body_bytes.clone())
        },
        http_version,
    };
    state.request_log.entry(subgraph).or_default().push(record);

    let rebuilt_body = axum::body::Body::from(body_bytes);
    let request = axum::extract::Request::from_parts(parts, rebuilt_body);
    next.run(request).await
}

async fn handle_on_request(
    axum::extract::State(on_request): axum::extract::State<Arc<OnRequest>>,
    request: axum::extract::Request,
    next: axum::middleware::Next,
) -> impl axum::response::IntoResponse {
    let path = request.uri().path().to_string();
    let (parts, body) = request.into_parts();
    let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap();

    let req = RequestLike {
        path: path.clone(),
        headers: parts.headers.clone(),
        body: if body_bytes.is_empty() {
            None
        } else {
            Some(body_bytes.clone())
        },
        http_version: parts.version,
    };

    if let Some(new_resp) = on_request(req) {
        if let Some(body) = new_resp.body {
            // short-circuit the request and respond with this body and the configured headers
            let mut response = axum::response::Response::builder()
                .status(new_resp.status)
                .body(axum::body::Body::from(body))
                .unwrap();
            *response.headers_mut() = new_resp.headers;
            return response;
        }

        if !new_resp.status.is_success() {
            // no body and non-success status, short-circuit with just the status and headers
            let mut response = axum::response::Response::builder()
                .status(new_resp.status)
                .body(axum::body::Body::empty())
                .unwrap();
            *response.headers_mut() = new_resp.headers;
            return response;
        }

        // no body provided, use body of the subgraph but change the status and headers
        let rebuilt_body = axum::body::Body::from(body_bytes);
        let rebuilt = axum::extract::Request::from_parts(parts, rebuilt_body);
        let mut subgraph_resp = next.run(rebuilt).await;
        for (name, value) in new_resp.headers {
            if let Some(name) = name {
                subgraph_resp.headers_mut().insert(name, value);
            }
        }
        return subgraph_resp;
    }

    // no override, just pass through the original request
    let rebuilt_body = axum::body::Body::from(body_bytes);
    let request = axum::extract::Request::from_parts(parts, rebuilt_body);
    next.run(request).await
}

impl TestSubgraphs<Built> {
    pub fn builder() -> TestSubgraphsBuilder {
        TestSubgraphsBuilder::new()
    }

    pub async fn start(self) -> TestSubgraphs<Started> {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("failed to bind tcp listener");
        let addr = listener.local_addr().expect("failed to get local address");
        drop(listener); // release the listener; the axum_server will bind to the same addr

        let (mut app, active_subscriptions) = subgraphs_app(self.subscriptions_protocol.clone());

        let middleware_state = Arc::new(TestSubgraphsMiddlewareState {
            request_log: DashMap::new(),
        });
        if let Some(on_request) = self.on_request.clone() {
            app = app.layer(axum::middleware::from_fn_with_state(
                on_request,
                handle_on_request,
            ));
        }
        if let Some(delay) = self.delay {
            app = app.layer(axum::middleware::from_fn(
                move |req, next: axum::middleware::Next| async move {
                    tokio::time::sleep(delay).await;
                    next.run(req).await
                },
            ));
        }
        // record_requests must be outermost so it logs the request before any blocking on_request handler runs
        app = app.layer(axum::middleware::from_fn_with_state(
            middleware_state.clone(),
            record_requests,
        ));

        let rustls_config_clone = self.rustls_config.clone();

        let server_handle = Handle::new();
        let server_handle_clone = server_handle.clone();
        tokio::spawn(async move {
            if let Some(rustls_config) = self.rustls_config {
                axum_server::bind_rustls(addr, rustls_config)
                    .handle(server_handle_clone.clone())
                    .serve(app.into_make_service())
                    .await
                    .expect("failed to start subgraphs server");
            } else if self.http2_only {
                axum_server::bind(addr)
                    .http2_only()
                    .handle(server_handle_clone.clone())
                    .serve(app.into_make_service())
                    .await
                    .expect("failed to start subgraphs h2c server");
            } else {
                axum_server::bind(addr)
                    .handle(server_handle_clone.clone())
                    .serve(app.into_make_service())
                    .await
                    .expect("failed to start subgraphs server");
            }
        });

        let addr = server_handle
            .listening()
            .await
            .expect("failed to get subgraphs server address");

        TestSubgraphs {
            on_request: self.on_request,
            rustls_config: rustls_config_clone,
            delay: self.delay,
            subscriptions_protocol: self.subscriptions_protocol,
            http2_only: self.http2_only,
            handle: Some(TestSubgraphsHandle {
                server_handle,
                addr,
                state: middleware_state,
                active_subscriptions,
            }),
            _state: PhantomData,
        }
    }
}

impl TestSubgraphs<Started> {
    #[allow(unused)]
    pub fn url(&self) -> String {
        let addr = self.handle.as_ref().expect("subgraphs not started").addr;
        let protocol = if self.rustls_config.is_some() {
            "https"
        } else {
            "http"
        };
        format!("{}://{}", protocol, addr)
    }

    /// Returns the number of currently active subscriptions on the reviews subgraph.
    pub fn active_subscriptions(&self) -> usize {
        self.handle
            .as_ref()
            .expect("subgraphs not started")
            .active_subscriptions
            .load(Ordering::SeqCst)
    }

    /// Returns the list of requests received on the given subgraph. Supply the subgarph name.
    pub fn get_requests_log(&self, subgraph: &str) -> Option<Vec<RequestLike>> {
        self.handle
            .as_ref()
            .expect("subgraphs not started")
            .state
            .request_log
            .get(subgraph)
            .map(|entry| entry.value().to_vec())
    }

    /// Replaces the subgraphs address in the given supergraph string with the test
    /// subgraphs address and returns the modified supergraph.
    ///
    /// It will replace all occurrences of `0.0.0.0:4200` with the test subgraphs address.
    pub fn supergraph(&self, supergraph: &str) -> String {
        supergraph_with_subgraphs(supergraph, &self.url())
    }
}

impl Drop for TestSubgraphsHandle {
    fn drop(&mut self) {
        self.server_handle.graceful_shutdown(None);
    }
}

// router

pub struct TestRouterBuilder {
    wait_for_healthy_on_start: bool,
    wait_for_ready_on_start: bool,
    config: Option<HiveRouterConfig>,
    plugins: Vec<Box<dyn Fn(PluginRegistry) -> PluginRegistry>>,
    subgraphs_url: Option<String>,
    port: u16,
    listener: Option<std::net::TcpListener>,
}

impl TestRouterBuilder {
    pub fn new() -> Self {
        Self {
            wait_for_healthy_on_start: true,
            wait_for_ready_on_start: true,
            config: None,
            plugins: vec![],
            subgraphs_url: None,
            port: 0,
            listener: None,
        }
    }

    pub fn inline_config(mut self, config_yaml: impl Into<String>) -> Self {
        let router_config =
            parse_yaml_config(config_yaml.into()).expect("failed to parse inline YAML config");
        self.config = Some(router_config);
        self
    }

    pub fn file_config(mut self, config_path: &str) -> Self {
        let supergraph_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(config_path);
        self.config = Some(
            load_config(Some(supergraph_path.to_str().unwrap().to_string()))
                .expect("failed to load router config from file"),
        );
        self
    }

    pub fn set_config(mut self, config: HiveRouterConfig) -> Self {
        self.config = Some(config);
        self
    }

    pub fn with_subgraphs(self, subgraphs: &TestSubgraphs<Started>) -> Self {
        self.with_subgraphs_url(subgraphs.url())
    }

    pub fn with_subgraphs_url(mut self, subgraphs_url: String) -> Self {
        self.subgraphs_url = Some(subgraphs_url);
        self
    }

    pub fn with_port(mut self, port: u16) -> Self {
        self.port = port;
        self
    }

    pub fn with_listener(mut self, listener: std::net::TcpListener) -> Self {
        self.listener = Some(listener);
        self
    }

    pub fn skip_wait_for_healthy_on_start(mut self) -> Self {
        self.wait_for_healthy_on_start = false;
        self
    }

    pub fn skip_wait_for_ready_on_start(mut self) -> Self {
        self.wait_for_ready_on_start = false;
        self
    }

    pub fn register_plugin<P: RouterPlugin>(mut self) -> Self {
        self.plugins.push(Box::new(|registry: PluginRegistry| {
            registry.register::<P>()
        }));
        self
    }

    pub fn build(self) -> TestRouter<Built> {
        let mut config = self.config.unwrap_or_default();
        config.http.port = self.port; // sync with config // TODO: what if testing custom port?
        let mut _hold_until_drop: Vec<Box<dyn Any>> = vec![];

        // change the supergraph to use the test subgraphs address
        if let Some(subgraphs_url) = self.subgraphs_url {
            match &config.supergraph {
                hive_router_config::supergraph::SupergraphSource::File { path, .. } => {
                    let supergraph_path = path.as_ref().expect("supergraph file path is required");

                    let temp_path = supergraph_temp_file_with_subgraphs(
                        supergraph_path.absolute.as_str(),
                        &subgraphs_url,
                    );

                    let supergraph_file_path =
                        hive_router_config::primitives::file_path::FilePath {
                            relative: temp_path.to_str().unwrap().to_string(),
                            absolute: temp_path.to_str().unwrap().to_string(),
                        };

                    config.supergraph = hive_router_config::supergraph::SupergraphSource::File {
                        path: Some(supergraph_file_path),
                        // TODO: we disable polling, but what if it was enabled?
                        poll_interval: None,
                    };

                    _hold_until_drop.push(Box::new(temp_path));
                }
                _ => warn!("Only file-based supergraph sources are supported in tests"),
            }
        }

        TestRouter {
            wait_for_healthy_on_start: self.wait_for_healthy_on_start,
            wait_for_ready_on_start: self.wait_for_ready_on_start,
            graphql_path: config.graphql_path().to_string(),
            websocket_path: config.websocket_path().map(|s| s.to_string()),
            callback_conf: config.callback_conf().cloned(),
            port: self.port,
            listener: self.listener,
            config: Some(config),
            plugins: self.plugins,
            handle: None,
            _hold_until_drop,
            _state: PhantomData,
        }
    }
}

impl Default for TestRouterBuilder {
    fn default() -> Self {
        Self::new()
    }
}

struct TestRouterHandle {
    schema_state: Arc<SchemaState>,
    shared_state: Arc<RouterSharedState>,
    serv: test::TestServer,
    bg_tasks_manager: BackgroundTasksManager,
    telemetry: Telemetry,
}

impl Drop for TestRouterHandle {
    fn drop(&mut self) {
        // shut down backgroun tasks
        self.bg_tasks_manager.shutdown();

        // shutdown hooks and wait for complete (shutdown is async so yeah)
        let shared_state = self.shared_state.clone();
        std::thread::spawn(move || {
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("failed to build shutdown runtime")
                .block_on(invoke_shutdown_hooks(&shared_state))
        })
        .join()
        .expect("shutdown hooks panicked");

        let traces_provider = self.telemetry.traces_provider.clone();
        let metrics_provider = self.telemetry.metrics_provider.clone();
        let dispatch = tracing::dispatcher::get_default(|current| current.clone());

        std::thread::spawn(move || {
            tracing::dispatcher::with_default(&dispatch, || {
                if let Some(provider) = traces_provider {
                    tracing::info!(
                        component = "telemetry",
                        layer = "provider",
                        layer = "provider",
                        "shutdown completed"
                    );
                    let _ = provider.force_flush();
                    let _ = provider.shutdown();
                    tracing::info!(
                        component = "telemetry",
                        layer = "provider",
                        "shutdown completed"
                    );
                }

                if let Some(provider) = metrics_provider {
                    tracing::info!(
                        component = "telemetry",
                        layer = "metrics",
                        "shutdown scheduled"
                    );
                    let _ = provider.force_flush();
                    let _ = provider.shutdown();
                    tracing::info!(
                        component = "telemetry",
                        layer = "metrics",
                        "shutdown completed"
                    );
                }
            });
        })
        .join()
        .expect("tracing shutdown panicked");
    }
}

pub struct TestRouter<State> {
    wait_for_healthy_on_start: bool,
    wait_for_ready_on_start: bool,
    graphql_path: String,
    websocket_path: Option<String>,
    callback_conf: Option<CallbackConfig>,
    port: u16,
    listener: Option<std::net::TcpListener>,
    config: Option<HiveRouterConfig>,
    plugins: Vec<Box<dyn Fn(PluginRegistry) -> PluginRegistry>>,
    handle: Option<TestRouterHandle>,
    _hold_until_drop: Vec<Box<dyn Any>>,
    _state: PhantomData<State>,
}

impl TestRouter<Built> {
    pub fn builder() -> TestRouterBuilder {
        TestRouterBuilder::new()
    }

    // When self-signed certificates are used, the ntex test client doesn't work
    // So we can't use it to call the healthcheck endpoints
    pub async fn start_without_healthcheck(mut self) -> TestRouter<Started> {
        init_rustls_crypto_provider();
        let config = self.config.take().unwrap();
        let (telemetry, subscriber) = Telemetry::init_testing_subscriber(&config)
            .expect("failed to initialize telemetry subscriber");
        let subscription_guard = tracing::subscriber::set_default(subscriber);
        let prometheus = telemetry
            .prometheus
            .as_ref()
            .and_then(|prom| prom.to_attached());

        let mut bg_tasks_manager = BackgroundTasksManager::new();
        let (shared_state, schema_state) = configure_app_from_config(
            config,
            telemetry.context.clone(),
            &mut bg_tasks_manager,
            self.plugins
                .iter()
                .fold(PluginRegistry::new(), |registry, register_plugin| {
                    register_plugin(registry)
                }),
        )
        .await
        .expect("failed to configure hive router from config");

        // capture the current tracing dispatch so it can be propagated to the
        // server thread spawned by test::server (which runs on a separate thread
        // and would otherwise use the no-op global subscriber)
        let serv_dispatch = tracing::dispatcher::get_default(|d| d.clone());

        let serv_shared_state = shared_state.clone();
        let serv_schema_state = schema_state.clone();
        let serv_callback_subs = schema_state.callback_subscriptions.clone();
        let serv_graphql_path = self.graphql_path.clone();
        let serv_websocket_path = self.websocket_path.clone();

        // when `listen` is set, the callback route lives on a dedicated server bound to that
        // address as a background task; otherwise it is mounted on the main server
        let serv_callback_path = match self.callback_conf {
            Some(CallbackConfig {
                listen: Some(listen),
                ref path,
                ..
            }) => {
                let cb_path = path.to_string();
                let cb_addr = listen.to_string();
                let cb_subs = schema_state.callback_subscriptions.clone();
                let cb_telemetry_context = shared_state.telemetry_context.clone();

                let server = web::HttpServer::new(async move || {
                    let cb_subs = cb_subs.clone();
                    let cb_path = cb_path.clone();
                    let telemetry_context = cb_telemetry_context.clone();
                    web::App::new()
                        .state(cb_subs)
                        .state(telemetry_context)
                        .configure(move |m| add_callback_handler(m, &cb_path))
                })
                .bind(&cb_addr)
                .expect("failed to bind callback server")
                .run();

                bg_tasks_manager.register_handle(async move {
                    server.await.ok();
                });

                None
            }
            Some(ref cb) => Some(cb.path.to_string()),
            None => None,
        };

        let paths = RouterPaths::new(
            serv_graphql_path,
            serv_websocket_path,
            serv_callback_path.clone(),
        );
        paths
            .detect_conflicts(&prometheus)
            .expect("failed to detect endpoint conflicts");

        let serv_listener = self.listener.unwrap_or(
            std::net::TcpListener::bind(format!("127.0.0.1:{}", self.port))
                .expect("failed to bind tcp listener for test server"),
        );
        let serv_port = serv_listener
            .local_addr()
            .expect("failed to get local address of test server")
            .port();
        let serv_paths = paths.clone();
        let serv_prometheus = prometheus.clone();
        let long_lived_limit = LongLivedClientLimitService::new(&shared_state.router_config);
        let mut serv_config = test::config()
            .client_timeout(Seconds(
                shared_state
                    .router_config
                    .traffic_shaping
                    .router
                    .request_timeout
                    .as_secs() as u16
                    + 1,
            ))
            .listener(serv_listener);
        if let Some(tls_config) = serv_shared_state
            .router_config
            .traffic_shaping
            .router
            .tls
            .as_ref()
        {
            let rustls_config = hive_router::tls::build_rustls_config(tls_config)
                .expect("failed to build rustls config for test router");
            serv_config = serv_config.rustls(rustls_config);
        }

        let serv = test::server_with(serv_config, move || {
            let shared_state = serv_shared_state.clone();
            let schema_state = serv_schema_state.clone();
            let paths = serv_paths.clone();
            let prometheus = serv_prometheus.clone();
            let serv_callback_path = serv_callback_path.clone();
            let callback_subs = serv_callback_subs.clone();
            let long_lived_limit = long_lived_limit.clone();

            // set the tracing dispatch on the server thread. the guard is
            // intentionally leaked: dropping it would restore the no-op default
            // dispatch, undoing what we just set. the guard is `!send` (thread-
            // local), so we can't move it back to the test thread. this is fine
            // because when the server thread exits (on testserver drop) the
            // thread-local storage is reclaimed by the os, and there is no prior
            // dispatch to restore
            let guard = tracing::dispatcher::set_default(&serv_dispatch);
            std::mem::forget(guard);

            async move {
                web::App::new()
                    .middleware(long_lived_limit)
                    .middleware(PluginService::new(
                        paths.clone(),
                        prometheus.as_ref().map(|p| p.endpoint.clone()),
                    ))
                    .middleware(RequestIdentifiersService)
                    .middleware(RequestSummaryService)
                    .state(shared_state.clone())
                    .state(schema_state)
                    .state(callback_subs)
                    .state(shared_state.telemetry_context.clone())
                    .configure(|m| configure_ntex_app(m, &paths, prometheus))
                    .configure(|m| {
                        if let Some(ref callback) = serv_callback_path {
                            add_callback_handler(m, callback);
                        }
                    })
            }
        })
        .await;

        let mut hold_until_drop = self._hold_until_drop;
        hold_until_drop.push(Box::new(subscription_guard));
        TestRouter {
            port: serv_port,
            listener: None,
            wait_for_healthy_on_start: self.wait_for_healthy_on_start,
            wait_for_ready_on_start: self.wait_for_ready_on_start,
            graphql_path: self.graphql_path,
            websocket_path: self.websocket_path,
            callback_conf: self.callback_conf,
            handle: Some(TestRouterHandle {
                schema_state,
                shared_state,
                serv,
                bg_tasks_manager,
                telemetry,
            }),
            config: self.config,
            plugins: self.plugins,
            _hold_until_drop: hold_until_drop,
            _state: PhantomData,
        }
    }

    pub async fn start(self) -> TestRouter<Started> {
        let started = self.start_without_healthcheck().await;

        if started.wait_for_healthy_on_start {
            info!("Waiting for healthcheck to pass...");
            started.wait_for_healthy(None).await;
        }

        if started.wait_for_ready_on_start {
            info!("Waiting for readiness check to pass...");
            started.wait_for_ready(None).await;
        }

        started
    }
}

impl TestRouter<Started> {
    pub fn schema_state(&self) -> &Arc<SchemaState> {
        &self.handle.as_ref().unwrap().schema_state
    }

    pub fn shared_state(&self) -> &Arc<RouterSharedState> {
        &self.handle.as_ref().unwrap().shared_state
    }

    pub fn serv(&self) -> &test::TestServer {
        &self.handle.as_ref().unwrap().serv
    }

    /// Waits for the /health endpoint to return 200 OK, with an optional timeout (defaults to 5 seconds).
    pub async fn wait_for_healthy(&self, timeout: Option<Duration>) {
        tokio::time::timeout(timeout.unwrap_or(Duration::from_secs(10)), async {
            loop {
                match self.serv().get("/health").send().await {
                    Ok(response) => {
                        if response.status() == 200 {
                            break;
                        }
                    }
                    Err(_) => {
                        // will resolve fast
                        tokio::time::sleep(Duration::from_millis(300)).await;
                    }
                }
            }
        })
        .await
        .expect("healthcheck timed out");
    }

    /// Waits for the /readiness endpoint to return 200 OK, with an optional timeout (defaults to 5 seconds).
    pub async fn wait_for_ready(&self, timeout: Option<Duration>) {
        tokio::time::timeout(timeout.unwrap_or(Duration::from_secs(10)), async {
            loop {
                match self.serv().get("/readiness").send().await {
                    Ok(response) => {
                        if response.status() == 200 {
                            break;
                        }
                    }
                    Err(_) => {
                        // readiness might take a moment, so allow for some time
                        // to avoid spamming the router and the h1 dispatcher
                        tokio::time::sleep(Duration::from_millis(600)).await;
                    }
                }
            }
        })
        .await
        .expect("readiness timed out");
    }

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

    pub async fn send_graphql_request(
        &self,
        query: &str,
        variables: Option<sonic_rs::Value>,
        headers: Option<http::HeaderMap>,
    ) -> ClientResponse {
        self.send_post_request(
            self.graphql_path(),
            json!({
              "query": query,
              "variables": variables,
            }),
            headers,
        )
        .await
    }

    pub async fn send_post_request(
        &self,
        path: &str,
        payload: sonic_rs::Value,
        headers: Option<http::HeaderMap>,
    ) -> ClientResponse {
        let mut req = self
            .serv()
            .post(path)
            .header(CONTENT_TYPE, "application/json")
            .header(ACCEPT, "application/graphql-response+json");

        if let Some(headers) = headers {
            for (key, value) in headers.iter() {
                req = req.set_header(key, value);
            }
        }

        req.send_json(&payload)
            .await
            .expect("Failed to send graphql request")
    }

    pub async fn ws(&self) -> WsConnection<Sealed> {
        let url = self.handle.as_ref().unwrap().serv.url(
            self.websocket_path
                .as_deref()
                .expect("Websocket path not set"),
        );
        let ws_url = url.as_str().replace("http://", "ws://");
        let ws_uri = ws_url.parse::<http::Uri>().expect("Failed to parse ws url");
        websocket_client::connect(&ws_uri, None)
            .await
            .expect("Failed to connect to websocket")
    }
}

pub trait ClientResponseExt {
    fn string_body(&self) -> impl Future<Output = String>;
    fn json_body(&self) -> impl Future<Output = sonic_rs::Value>;
    fn json_body_string_pretty(&self) -> impl Future<Output = String>;
    /// The difference from [`json_body_string_pretty`] is that this method uses a stable
    /// pretty-printer that does not depend on the order of fields in the JSON object.
    fn json_body_string_pretty_stable(&self) -> impl Future<Output = String>;
    /// Reads a response header and parses it as a `u64`. Used for the
    /// demand-control `X-Cost-*` headers. Header lookup is case-insensitive.
    fn cost_header(&self, name: &str) -> Option<u64>;
}

impl ClientResponseExt for ClientResponse {
    fn cost_header(&self, name: &str) -> Option<u64> {
        self.headers()
            .get(name)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.parse::<u64>().ok())
    }

    async fn string_body(&self) -> String {
        let body = self.body().await.expect("failed to read request body");
        std::str::from_utf8(&body)
            .expect("body is not valid UTF-8")
            .to_string()
    }

    async fn json_body(&self) -> sonic_rs::Value {
        let body = self.body().await.expect("failed to read request body");
        sonic_rs::from_slice(&body).expect("failed to parse request body to JSON")
    }

    async fn json_body_string_pretty(&self) -> String {
        sonic_rs::to_string_pretty(&self.json_body().await)
            .expect("failed to pretty print JSON body")
    }

    async fn json_body_string_pretty_stable(&self) -> String {
        let body = self.body().await.expect("failed to read request body");
        let stable_json: serde_json::Value =
            sonic_rs::from_slice(&body).expect("failed to parse request body to JSON");
        serde_json::to_string_pretty(&stable_json).expect("failed to pretty print canonical JSON")
    }
}

pub async fn wait_until_mock_matched(mock: &Mock) -> Result<(), String> {
    let now = Instant::now();
    let timeout = Duration::from_secs(10); // always a sane default
    loop {
        if mock.matched_async().await {
            return Ok(());
        }

        // anything less will congest the router, keep the interval chill
        time::sleep(Duration::from_millis(100)).await;

        if now.elapsed() > timeout {
            return Err(format!("timeout after {:?}", now.elapsed()));
        }
    }
}

pub struct GeneratedKeyPair {
    pub cert_file: NamedTempFile,
    pub cert_file_path: String,
    pub cert_pem: String,
    pub key_file: NamedTempFile,
    pub key_file_path: String,
    pub key_pem: String,
}

pub async fn generate_keypair() -> GeneratedKeyPair {
    let cert_key = generate_simple_self_signed(vec![
        "127.0.0.1".to_string(),
        "localhost".to_string(),
        "0.0.0.0".to_string(),
    ])
    .expect("Failed to generate self-signed certificate");

    let mut cert_file =
        NamedTempFile::new().expect("Failed to create temporary file for certificate");
    let cert = cert_key.cert;
    let cert_pem = cert.pem();
    let _ = cert_file
        .write(cert_pem.as_bytes())
        .expect("Failed to write certificate to temporary file");

    let mut key_file =
        NamedTempFile::new().expect("Failed to create temporary file for private key");
    let key = cert_key.signing_key;
    let key_pem = key.serialize_pem();
    let _ = key_file
        .write(key_pem.as_bytes())
        .expect("Failed to write private key to temporary file");

    GeneratedKeyPair {
        cert_file_path: cert_file
            .path()
            .to_str()
            .expect("Failed to convert cert file path to string")
            .to_string(),
        cert_file,
        cert_pem,
        key_file_path: key_file
            .path()
            .to_str()
            .expect("Failed to convert key file path to string")
            .to_string(),
        key_file,
        key_pem,
    }
}

pub async fn generate_tls_subgraph() -> (TestSubgraphs<Started>, GeneratedKeyPair) {
    let generated_key_pair = generate_keypair().await;
    let rustls_config = RustlsConfig::from_pem_file(
        &generated_key_pair.cert_file_path,
        &generated_key_pair.key_file_path,
    )
    .await
    .expect("Failed to create RustlsConfig from PEM files");
    let subgraphs = TestSubgraphs::builder()
        .with_rustls_config(rustls_config)
        .build()
        .start()
        .await;
    (subgraphs, generated_key_pair)
}
