//! Imperative `mcptest tools call` with pipeline primitives.
//!
//!  shipped `tools` (list). This adds the imperative call surface plus
//! the chaining primitives that let a user chain tool calls without reaching
//! for `jq`: extract args from stdin JSON by JSONPath (`--arg name=$.path`),
//! read the whole args object from stdin (`--args-from-stdin`), capture a
//! value from the output (`--bind name=$.path`) and feed it into one chained
//! call (`--then <tool> --arg name=:bound`), and project the output before
//! printing (`--select $.path`). Multi-step chains belong in `mcptest pipe`;
//! `--then` is capped at one chained step on purpose.

use std::collections::HashMap;
use std::io::Read;

use anyhow::{Context, Result};
use serde_json::{Map, Value};

use crate::pipe::conn::Connection;
use crate::pipe::jsonpath;
use crate::pipe::pricing::PricingTable;

/// Plain-data inputs for one `tools call` invocation.
///
/// The CLI args struct lowers into this so the engine stays free of clap and
/// the tests build it directly. Each optional field maps to one flag.
#[derive(Debug, Default, Clone)]
pub struct ToolsCallArgs {
    /// MCP server endpoint.
    pub url: String,
    /// Optional bearer token forwarded on every request.
    pub bearer_token: Option<String>,
    /// Tool to call.
    pub tool: String,
    /// Literal args object as a JSON string (`--args '{...}'`).
    pub args_json: Option<String>,
    /// Read the whole args object from stdin JSON (`--args-from-stdin`).
    pub args_from_stdin: bool,
    /// `--arg name=$.path` extractions from stdin JSON, in order.
    pub arg_paths: Vec<(String, String)>,
    /// `--bind name=$.path` captures from this call's output.
    pub binds: Vec<(String, String)>,
    /// `--select $.path` projection applied before printing.
    pub select: Option<String>,
    /// Emit JSON instead of pretty output (the `--json` / `--format json`
    /// shorthand for this surface).
    pub json: bool,
    /// Chained second call.
    pub then: Option<ThenStep>,
    /// Aggregate USD ceiling across both calls in this invocation.
    /// Uses the bundled pricing table's per-call rate.
    pub max_cost: Option<f64>,
}

/// A single chained `--then` step.
#[derive(Debug, Default, Clone)]
pub struct ThenStep {
    /// Tool to call in the chained step.
    pub tool: String,
    /// `--arg name=value-or-:bound` pairs. A leading `:` references a value
    /// captured by an earlier `--bind`.
    pub args: Vec<(String, String)>,
}

/// Run the imperative tool call (plus optional chained step), printing the
/// final (optionally projected) output.
///
/// Returns the process exit code. The whole thing runs on a single-threaded
/// tokio runtime built here so the surrounding handler stays sync, matching
/// the rest of the imperative subcommands.
pub fn run(args: ToolsCallArgs) -> Result<i32> {
    let stdin_json = read_stdin_json_if_needed(&args)?;
    let initial_args = build_initial_args(&args, stdin_json.as_ref())?;

    //  parity: cap the cumulative cost across the (up to two) calls in
    // this invocation. A plain tool call has no token usage, so we attribute
    // the bundled per-call rate and refuse before a call that would breach.
    if let Some(cap) = args.max_cost {
        let per_call = PricingTable::bundled()
            .map(|t| t.tool_call_cost())
            .unwrap_or(0.0);
        let calls = if args.then.is_some() { 2.0 } else { 1.0 };
        if per_call * calls > cap + 1e-9 {
            anyhow::bail!(
                "tool call budget exceeded: {} call(s) at ${per_call:.4} each is over the --max-cost ${cap:.4} cap",
                calls as u32
            );
        }
    }

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build tokio runtime")?;
    runtime.block_on(async move {
        let conn = Connection::connect(&args.url, args.bearer_token.clone()).await?;
        let first = conn.call_tool(&args.tool, &initial_args).await;
        let output = match first {
            Ok(v) => v,
            Err(e) => {
                conn.close().await;
                return Err(e);
            }
        };

        // Capture binds, then run the optional chained step.
        let bound = capture_binds(&output, &args.binds)?;
        let final_output = match &args.then {
            Some(then) => {
                let then_args = build_then_args(&then.args, &bound)?;
                let chained = conn.call_tool(&then.tool, &then_args).await;
                match chained {
                    Ok(v) => v,
                    Err(e) => {
                        conn.close().await;
                        return Err(e);
                    }
                }
            }
            None => output,
        };
        conn.close().await;

        let projected = match &args.select {
            Some(path) => jsonpath::project(&final_output, path)
                .with_context(|| format!("--select `{path}` failed"))?,
            None => final_output,
        };
        print_output(&projected, args.json)?;
        Ok(0)
    })
}

fn read_stdin_json_if_needed(args: &ToolsCallArgs) -> Result<Option<Value>> {
    let needs_stdin = args.args_from_stdin || !args.arg_paths.is_empty();
    if !needs_stdin {
        return Ok(None);
    }
    let mut buf = String::new();
    std::io::stdin()
        .read_to_string(&mut buf)
        .context("failed to read stdin")?;
    if buf.trim().is_empty() {
        anyhow::bail!("expected JSON on stdin but it was empty");
    }
    let value: Value = serde_json::from_str(buf.trim()).context("stdin was not valid JSON")?;
    Ok(Some(value))
}

/// Build the first call's args from the precedence: `--args-from-stdin` reads
/// the whole object, otherwise `--args` provides a base object that
/// `--arg name=$.path` extractions overlay onto.
fn build_initial_args(args: &ToolsCallArgs, stdin: Option<&Value>) -> Result<Value> {
    if args.args_from_stdin {
        let value = stdin.context("--args-from-stdin set but no stdin JSON was read")?;
        if !value.is_object() {
            anyhow::bail!("--args-from-stdin requires a JSON object on stdin");
        }
        return Ok(value.clone());
    }
    let mut object: Map<String, Value> = match &args.args_json {
        Some(raw) => serde_json::from_str(raw).context("--args was not valid JSON")?,
        None => Map::new(),
    };
    for (name, path) in &args.arg_paths {
        let source = stdin.context("--arg requires JSON on stdin")?;
        let value = jsonpath::extract_one(source, path)
            .with_context(|| format!("--arg {name}={path} extraction failed"))?;
        object.insert(name.clone(), value);
    }
    Ok(Value::Object(object))
}

fn capture_binds(output: &Value, binds: &[(String, String)]) -> Result<HashMap<String, Value>> {
    let mut bound = HashMap::new();
    for (name, path) in binds {
        let value = jsonpath::extract_one(output, path)
            .with_context(|| format!("--bind {name}={path} extraction failed"))?;
        bound.insert(name.clone(), value);
    }
    Ok(bound)
}

/// Build the chained step's args, resolving `:bound` references against the
/// captured bind map and treating everything else as a literal string.
fn build_then_args(pairs: &[(String, String)], bound: &HashMap<String, Value>) -> Result<Value> {
    let mut object = Map::new();
    for (name, raw) in pairs {
        let value = match raw.strip_prefix(':') {
            Some(reference) => bound
                .get(reference)
                .cloned()
                .with_context(|| format!("--then arg {name} references unbound `:{reference}`"))?,
            None => Value::String(raw.clone()),
        };
        object.insert(name.clone(), value);
    }
    Ok(Value::Object(object))
}

fn print_output(value: &Value, json: bool) -> Result<()> {
    if json {
        println!("{}", serde_json::to_string_pretty(value)?);
    } else {
        match value {
            Value::String(s) => println!("{s}"),
            other => println!("{}", serde_json::to_string_pretty(other)?),
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn args_from_stdin_uses_whole_object() {
        let args = ToolsCallArgs {
            args_from_stdin: true,
            ..Default::default()
        };
        let stdin = json!({ "user_id": "abc" });
        let built = build_initial_args(&args, Some(&stdin)).unwrap();
        assert_eq!(built, json!({ "user_id": "abc" }));
    }

    #[test]
    fn arg_path_overlays_extraction_onto_args() {
        let args = ToolsCallArgs {
            args_json: Some(r#"{"limit":10}"#.to_string()),
            arg_paths: vec![("user_id".to_string(), "$.id".to_string())],
            ..Default::default()
        };
        let stdin = json!({ "id": "u1" });
        let built = build_initial_args(&args, Some(&stdin)).unwrap();
        assert_eq!(built, json!({ "limit": 10, "user_id": "u1" }));
    }

    #[test]
    fn arg_path_indexes_into_arrays() {
        let args = ToolsCallArgs {
            arg_paths: vec![("first".to_string(), "$.ids[0]".to_string())],
            ..Default::default()
        };
        let stdin = json!({ "ids": [7, 8, 9] });
        let built = build_initial_args(&args, Some(&stdin)).unwrap();
        assert_eq!(built, json!({ "first": 7 }));
    }

    #[test]
    fn binds_capture_from_output() {
        let output = json!({ "id": [42] });
        let binds = vec![("uid".to_string(), "$.id[0]".to_string())];
        let bound = capture_binds(&output, &binds).unwrap();
        assert_eq!(bound.get("uid"), Some(&json!(42)));
    }

    #[test]
    fn then_args_resolve_bound_reference() {
        let mut bound = HashMap::new();
        bound.insert("uid".to_string(), json!(42));
        let pairs = vec![("user_id".to_string(), ":uid".to_string())];
        let built = build_then_args(&pairs, &bound).unwrap();
        assert_eq!(built, json!({ "user_id": 42 }));
    }

    #[test]
    fn then_args_treat_plain_values_as_strings() {
        let bound = HashMap::new();
        let pairs = vec![("q".to_string(), "alice".to_string())];
        let built = build_then_args(&pairs, &bound).unwrap();
        assert_eq!(built, json!({ "q": "alice" }));
    }

    #[test]
    fn then_args_error_on_unbound_reference() {
        let bound = HashMap::new();
        let pairs = vec![("x".to_string(), ":ghost".to_string())];
        assert!(build_then_args(&pairs, &bound).is_err());
    }

    #[test]
    fn args_from_stdin_rejects_non_object() {
        let args = ToolsCallArgs {
            args_from_stdin: true,
            ..Default::default()
        };
        let stdin = json!([1, 2, 3]);
        assert!(build_initial_args(&args, Some(&stdin)).is_err());
    }
}
