//! Restart-policy respawn wrapper for `mcptest run` (WOR-1323).
//!
//! A suite that declares `run_options.restart_policy: per_file` or
//! `per_test` wants a fresh server process (or, for URL targets, a fresh
//! connection) at the matching test boundaries. The decision table lives
//! in `mcptest_core::runner::restart`; this module is the side that acts
//! on it: a [`TestExecutor`] wrapper that tears one server out of the
//! pool and reconnects it before delegating to the wrapped executor.
//!
//! The wrapper sits outside the hook bracket (`HookedExecutor`) so the
//! respawn lands before `beforeEach` runs, the order
//! `docs/test-isolation.md` commits to. The CLI forces serial execution
//! for non-default policies, so no concurrent test can race a teardown.

use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;

use async_trait::async_trait;
use mcptest_core::connector::{connect_server, ConnectError, ServerPool};
use mcptest_core::network::ProxyConfig;
use mcptest_core::runner::restart::{
    classify_boundary, restart_action, RestartAction, RestartPolicy,
};
use mcptest_core::runner::{TestExecutor, TestOutcome, TestSpec};
use mcptest_core::suite::ServerSpec;

/// Everything a respawn needs at test-boundary time, captured before the
/// plans move into the executor (WOR-1323). Owned (not borrowed) because
/// the executor stack outlives the borrows `run_live_async` holds.
pub(super) struct RestartContext {
    /// Which boundaries trigger a respawn.
    pub policy: RestartPolicy,
    /// Declared server specs keyed by suite-local name, used to
    /// reconnect after a teardown.
    pub servers: BTreeMap<String, ServerSpec>,
    /// Tool-test name to server name. Agent tests are absent on purpose:
    /// their model loop ran in the plan-building prepass, so a
    /// runner-time respawn would bounce a server nothing will call again.
    pub test_servers: HashMap<String, String>,
    /// Proxy settings `connect_server` needs to reconnect URL targets.
    pub proxy: ProxyConfig,
    /// `-- <args>` stdio passthrough, re-applied on every respawn so the
    /// fresh process gets the same argv the initial connect used.
    pub passthrough_args: Vec<String>,
}

/// Build the respawn context for the run, or `None` for `per_run` so the
/// default path keeps its unwrapped executor stack (WOR-1323). The
/// test-to-server map comes from the tool plans because they are the
/// resolved view: a plan exists exactly for the tests the runner will
/// dispatch, and only tool plans name a respawnable server.
pub(super) fn build_restart_context(
    policy: RestartPolicy,
    suite: &mcptest_core::suite::ResolvedSuite,
    plans: &HashMap<String, Arc<mcptest_core::executor::PlannedTest>>,
    proxy: &ProxyConfig,
    passthrough_args: &[String],
) -> Option<RestartContext> {
    use mcptest_core::executor::PlannedTest;
    if policy == RestartPolicy::PerRun {
        return None;
    }
    let test_servers = plans
        .iter()
        .filter_map(|(name, plan)| match plan.as_ref() {
            PlannedTest::Tool(p) => Some((name.clone(), p.server.clone())),
            PlannedTest::Compliance(p) => Some((name.clone(), p.server.clone())),
            PlannedTest::Agent(_) => None,
        })
        .collect();
    Some(RestartContext {
        policy,
        servers: suite.servers.clone(),
        test_servers,
        proxy: proxy.clone(),
        passthrough_args: passthrough_args.to_vec(),
    })
}

/// A [`TestExecutor`] that respawns each test's server at the boundaries
/// its restart policy names, then delegates (WOR-1323).
///
/// Mirrors [`super::hooks::HookedExecutor`]: wrap, act before the test,
/// delegate. Cassette-backed servers are left in place (replay is
/// deterministic, so there is no state to flush), and a test without a
/// mapped server (agent tests, synthetic rows) delegates untouched. A
/// failed reconnect fails the test naming the server, the policy, and the
/// error; the next boundary retries the respawn.
pub(super) struct RestartingExecutor<E: TestExecutor> {
    inner: E,
    pool: Arc<ServerPool>,
    ctx: RestartContext,
    /// `file` of the previously executed test; outer `None` until the
    /// first test runs. A mutex because the trait takes `&self`; the CLI
    /// forces serial execution for non-default policies, so the lock is
    /// uncontended in practice.
    prev_file: tokio::sync::Mutex<Option<Option<String>>>,
}

impl<E: TestExecutor> RestartingExecutor<E> {
    /// Wrap `inner` so each test boundary consults `ctx.policy` against
    /// the shared `pool`.
    pub(super) fn new(inner: E, pool: Arc<ServerPool>, ctx: RestartContext) -> Self {
        Self {
            inner,
            pool,
            ctx,
            prev_file: tokio::sync::Mutex::new(None),
        }
    }

    /// Classify the boundary between the previous test and this one, and
    /// record this test's file for the next call. Split out so
    /// `execute_one` stays under the function-length cap.
    async fn should_restart(&self, test: &TestSpec) -> bool {
        let mut prev = self.prev_file.lock().await;
        let boundary = classify_boundary(prev.as_ref(), &test.file);
        *prev = Some(test.file.clone());
        restart_action(self.ctx.policy, boundary) == RestartAction::Restart
    }

    /// Tear one server out of the pool and reconnect it. Cassette specs
    /// return without touching the pool: replay holds no mutable server
    /// state worth flushing, and bouncing the replay task would only
    /// reset the exchange cursor the recording was made around. A
    /// teardown error is logged and the reconnect still proceeds; the old
    /// process is unreachable either way, and the fresh connection is
    /// what the policy promised.
    async fn respawn_server(&self, name: &str) -> Result<(), ConnectError> {
        let Some(spec) = self.ctx.servers.get(name) else {
            // Server references were validated at suite-load time, so a
            // missing spec means the entry was dropped after planning;
            // there is nothing to respawn.
            tracing::debug!(
                target: "mcptest::cli",
                server = %name,
                "restart_policy: no spec for server, skipping respawn"
            );
            return Ok(());
        };
        if matches!(spec, ServerSpec::Cassette { .. }) {
            tracing::debug!(
                target: "mcptest::cli",
                server = %name,
                "restart_policy: cassette server replays deterministically, kept in place"
            );
            return Ok(());
        }
        if let Err(e) = self.pool.shutdown_one(name).await {
            tracing::warn!(
                target: "mcptest::cli",
                server = %name,
                "restart_policy: shutdown before respawn failed ({e}); reconnecting anyway"
            );
        }
        let live = connect_server(
            name,
            spec,
            "mcptest",
            mcptest_core::version::semver(),
            &self.ctx.proxy,
            &self.ctx.passthrough_args,
        )
        .await?;
        self.pool.insert(name.to_string(), live).await;
        Ok(())
    }
}

#[async_trait]
impl<E: TestExecutor> TestExecutor for RestartingExecutor<E> {
    async fn execute_one(&self, test: &TestSpec) -> TestOutcome {
        if self.should_restart(test).await {
            // Only tool tests map to a respawnable server; anything
            // absent from the map delegates untouched.
            if let Some(server) = self.ctx.test_servers.get(&test.name) {
                if let Err(e) = self.respawn_server(server).await {
                    return TestOutcome::fail(format!(
                        "restart_policy {}: respawning server `{server}` before `{}` failed: {e}",
                        self.ctx.policy.label(),
                        test.name
                    ));
                }
            }
        }
        self.inner.execute_one(test).await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mcptest_core::report::TestStatus;

    /// Inner executor that always passes, isolating the wrapper.
    struct AlwaysPass;

    #[async_trait]
    impl TestExecutor for AlwaysPass {
        async fn execute_one(&self, _test: &TestSpec) -> TestOutcome {
            TestOutcome::pass()
        }
    }

    /// Minimal cassette: just the `initialize` handshake, enough to build
    /// a pool entry without any subprocess or network.
    const HANDSHAKE_CASSETTE: &str = r#"{
      "version": "1",
      "exchanges": [
        {
          "request": { "jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {} },
          "response": { "jsonrpc": "2.0", "id": 1, "result": { "protocolVersion": "2025-03-26", "capabilities": {}, "serverInfo": { "name": "recorded", "version": "0.1" } } }
        }
      ]
    }"#;

    fn cassette_spec(dir: &std::path::Path) -> ServerSpec {
        let path = dir.join("cassette.json");
        std::fs::write(&path, HANDSHAKE_CASSETTE).unwrap();
        ServerSpec::Cassette {
            path: path.display().to_string(),
            refresh_command: None,
            refresh_url: None,
        }
    }

    fn ctx_with(
        servers: BTreeMap<String, ServerSpec>,
        test_servers: HashMap<String, String>,
    ) -> RestartContext {
        RestartContext {
            policy: RestartPolicy::PerTest,
            servers,
            test_servers,
            proxy: ProxyConfig::default(),
            passthrough_args: Vec::new(),
        }
    }

    #[tokio::test]
    async fn cassette_backed_servers_are_kept_without_a_respawn() {
        let dir = tempfile::tempdir().unwrap();
        let spec = cassette_spec(dir.path());
        let pool = Arc::new(ServerPool::new());
        let live = connect_server(
            "srv",
            &spec,
            "mcptest",
            "0.0.0",
            &ProxyConfig::default(),
            &[],
        )
        .await
        .unwrap();
        pool.insert("srv".to_string(), live).await;
        let before = pool.get("srv").await.unwrap();

        let mut servers = BTreeMap::new();
        servers.insert("srv".to_string(), spec);
        let mut test_servers = HashMap::new();
        test_servers.insert("t".to_string(), "srv".to_string());
        let exec = RestartingExecutor::new(
            AlwaysPass,
            Arc::clone(&pool),
            ctx_with(servers, test_servers),
        );

        let outcome = exec.execute_one(&TestSpec::new("t")).await;
        assert_eq!(outcome.status, TestStatus::Pass);
        let after = pool.get("srv").await.unwrap();
        assert!(
            Arc::ptr_eq(&before, &after),
            "the cassette entry must survive the per_test boundary untouched"
        );
    }

    #[tokio::test]
    async fn a_test_without_a_mapped_server_delegates_cleanly() {
        // Agent tests (and synthetic rows) never appear in `test_servers`;
        // the wrapper must pass them straight through.
        let pool = Arc::new(ServerPool::new());
        let exec =
            RestartingExecutor::new(AlwaysPass, pool, ctx_with(BTreeMap::new(), HashMap::new()));
        let outcome = exec.execute_one(&TestSpec::new("agent case")).await;
        assert_eq!(outcome.status, TestStatus::Pass);
    }

    #[tokio::test]
    async fn a_failed_respawn_fails_the_test_naming_server_and_policy() {
        let pool = Arc::new(ServerPool::new());
        let mut servers = BTreeMap::new();
        servers.insert(
            "srv".to_string(),
            ServerSpec::Stdio {
                command: "/nonexistent/mcptest-wor1323-fixture".to_string(),
                args: Vec::new(),
                env: BTreeMap::new(),
            },
        );
        let mut test_servers = HashMap::new();
        test_servers.insert("t".to_string(), "srv".to_string());
        let exec = RestartingExecutor::new(AlwaysPass, pool, ctx_with(servers, test_servers));

        let outcome = exec.execute_one(&TestSpec::new("t")).await;
        assert_eq!(outcome.status, TestStatus::Fail);
        let message = outcome.failure.unwrap().message;
        assert!(message.contains("per_test"), "{message}");
        assert!(message.contains("`srv`"), "{message}");
    }
}
