//! Integration tests for `mcptest pipe` and its budget controls.
//!
//! Dry-run tests need no server. Live-run tests stand up a wiremock
//! streamable-HTTP MCP endpoint that answers `tools/call` from canned bodies,
//! the same pattern the `tools call` and introspect tests use.

use std::io::Write;

use assert_cmd::Command;
use predicates::prelude::*;
use serde_json::{json, Value};
use tempfile::NamedTempFile;
use wiremock::matchers::method;
use wiremock::{Mock, MockServer, Request, ResponseTemplate};

fn mcptest() -> Command {
    Command::cargo_bin("mcptest").expect("binary builds")
}

fn write_pipeline(yaml: &str) -> NamedTempFile {
    let mut f = NamedTempFile::new().expect("temp file");
    f.write_all(yaml.as_bytes()).expect("write yaml");
    f.flush().expect("flush");
    f
}

fn reply(body: Value) -> impl Fn(&Request) -> ResponseTemplate {
    move |req: &Request| {
        let parsed: Value = serde_json::from_slice(&req.body).unwrap_or(Value::Null);
        let id = parsed.get("id").cloned().unwrap_or(Value::Null);
        let envelope = json!({ "jsonrpc": "2.0", "id": id, "result": body.clone() });
        ResponseTemplate::new(200)
            .set_body_string(serde_json::to_string(&envelope).unwrap())
            .insert_header("content-type", "application/json")
    }
}

async fn mount(server: &MockServer, tool_responses: Vec<(&'static str, Value)>) {
    let init = json!({
        "protocolVersion": "2025-06-18",
        "serverInfo": { "name": "fake", "version": "0.1" },
        "capabilities": {}
    });
    Mock::given(method("POST"))
        .and(wiremock::matchers::body_partial_json(
            json!({ "method": "initialize" }),
        ))
        .respond_with(reply(init))
        .mount(server)
        .await;
    for (tool_name, body) in tool_responses {
        let want = json!({ "method": "tools/call", "params": { "name": tool_name } });
        Mock::given(method("POST"))
            .and(wiremock::matchers::body_partial_json(want))
            .respond_with(reply(body))
            .mount(server)
            .await;
    }
}

#[tokio::test(flavor = "current_thread")]
async fn dry_run_with_estimate_prints_projection() {
    let file =
        write_pipeline("pipe:\n  - id: find\n    tool: search\n  - id: details\n    tool: fetch\n");
    let assertion = mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--dry-run",
            "--estimate-cost",
        ])
        .assert()
        .success();
    let stdout = String::from_utf8(assertion.get_output().stdout.clone()).unwrap();
    assert!(stdout.contains("Total estimate"), "missing total: {stdout}");
    assert!(stdout.contains("Worst-case"), "missing worst-case");
}

#[tokio::test(flavor = "current_thread")]
async fn dry_run_makes_no_tool_calls() {
    // No server mounted; a dry run must still succeed because it never
    // connects.
    let file = write_pipeline("pipe:\n  - id: a\n    tool: t\n");
    mcptest()
        .args(["pipe", file.path().to_str().unwrap(), "--dry-run"])
        .assert()
        .success();
}

#[tokio::test(flavor = "current_thread")]
async fn invalid_pipeline_fails_with_schema_error() {
    // Missing `tool` violates the schema.
    let file = write_pipeline("pipe:\n  - id: a\n");
    mcptest()
        .args(["pipe", file.path().to_str().unwrap(), "--dry-run"])
        .assert()
        .failure()
        .stderr(predicate::str::contains("schema"));
}

#[tokio::test(flavor = "multi_thread")]
async fn runs_two_steps_resolving_a_reference() {
    let server = MockServer::start().await;
    mount(
        &server,
        vec![
            ("search", json!({ "users": [{ "id": "u1" }] })),
            ("fetch", json!({ "name": "alice" })),
        ],
    )
    .await;
    let file = write_pipeline(
        "pipe:\n  - id: find\n    tool: search\n    args:\n      query: alice\n  - id: details\n    tool: fetch\n    args:\n      user_id: ${find.users[0].id}\n",
    );
    let assertion = mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
            "--format",
            "json",
        ])
        .assert()
        .success();
    let stdout = String::from_utf8(assertion.get_output().stdout.clone()).unwrap();
    let trace: Value = serde_json::from_str(stdout.trim()).expect("json trace");
    assert_eq!(trace["steps"][0]["status"], "ok");
    assert_eq!(trace["last_output"]["name"], "alice");
}

#[tokio::test(flavor = "multi_thread")]
async fn when_guard_skips_a_step() {
    let server = MockServer::start().await;
    mount(&server, vec![("search", json!({ "users": [] }))]).await;
    let file = write_pipeline(
        "pipe:\n  - id: find\n    tool: search\n  - id: details\n    tool: fetch\n    when: ${find.users} | length > 0\n    args:\n      user_id: ${find.users[0].id}\n",
    );
    let assertion = mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
            "--format",
            "json",
        ])
        .assert()
        .success();
    let stdout = String::from_utf8(assertion.get_output().stdout.clone()).unwrap();
    let trace: Value = serde_json::from_str(stdout.trim()).expect("json trace");
    assert_eq!(trace["steps"][1]["status"], "skipped");
}

#[tokio::test(flavor = "multi_thread")]
async fn cost_cap_breach_exits_nonzero() {
    let server = MockServer::start().await;
    mount(
        &server,
        vec![("a", json!({ "ok": 1 })), ("b", json!({ "ok": 2 }))],
    )
    .await;
    // Per-call cost is $0.001 in the bundled table. A cap of $0.0005 cannot
    // even afford the first call, so the run stops and exits non-zero.
    let file = write_pipeline("pipe:\n  - id: a\n    tool: a\n  - id: b\n    tool: b\n");
    mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
            "--max-cost",
            "0.0005",
        ])
        .assert()
        .failure()
        .stderr(predicate::str::contains("budget exceeded"));
}

#[tokio::test(flavor = "multi_thread")]
async fn on_budget_exceeded_warn_runs_to_completion_and_exits_zero() {
    let server = MockServer::start().await;
    mount(
        &server,
        vec![("a", json!({ "ok": 1 })), ("b", json!({ "ok": 2 }))],
    )
    .await;
    let file = write_pipeline("pipe:\n  - id: a\n    tool: a\n  - id: b\n    tool: b\n");
    mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
            "--max-cost",
            "0.0005",
            "--on-budget-exceeded",
            "warn",
        ])
        .assert()
        .success();
}

#[tokio::test(flavor = "multi_thread")]
async fn on_error_stop_aborts_the_pipeline() {
    let server = MockServer::start().await;
    // Mount initialize, but never mount the `boom` tool so the call errors.
    Mock::given(method("POST"))
        .and(wiremock::matchers::body_partial_json(
            json!({ "method": "initialize" }),
        ))
        .respond_with(reply(json!({
            "protocolVersion": "2025-06-18",
            "serverInfo": { "name": "fake", "version": "0.1" },
            "capabilities": {}
        })))
        .mount(&server)
        .await;
    Mock::given(method("POST"))
        .and(wiremock::matchers::body_partial_json(
            json!({ "method": "tools/call" }),
        ))
        .respond_with(
            ResponseTemplate::new(200).set_body_string(
                serde_json::to_string(&json!({
                    "jsonrpc": "2.0", "id": 2,
                    "error": { "code": -32000, "message": "no such tool" }
                }))
                .unwrap(),
            ),
        )
        .mount(&server)
        .await;
    let file = write_pipeline("pipe:\n  - id: boom\n    tool: boom\n");
    mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
        ])
        .assert()
        .failure();
}

#[tokio::test(flavor = "multi_thread")]
async fn var_injection_flows_into_args() {
    let server = MockServer::start().await;
    mount(&server, vec![("search", json!({ "hit": true }))]).await;
    let file =
        write_pipeline("pipe:\n  - id: find\n    tool: search\n    args:\n      query: ${var.Q}\n");
    let assertion = mcptest()
        .args([
            "pipe",
            file.path().to_str().unwrap(),
            "--url",
            &server.uri(),
            "--var",
            "Q=alice",
            "--format",
            "json",
        ])
        .assert()
        .success();
    let stdout = String::from_utf8(assertion.get_output().stdout.clone()).unwrap();
    let trace: Value = serde_json::from_str(stdout.trim()).expect("json trace");
    assert_eq!(trace["steps"][0]["status"], "ok");
}
