//! Per-method JSON-RPC handlers for the operator-facing mock.
//!
//! Pure request/response: each function maps one parsed request to a JSON-RPC
//! response value, with no transport or timing concerns (those live in the
//! serve loop in `mock/mod.rs`). Errors are returned as JSON-RPC error
//! envelopes rather than `Result::Err` so a single malformed call cannot take
//! the stdio loop down. Split out so `mock/mod.rs` stays under the module-size
//! guard.

use std::collections::HashMap;

use serde_json::{json, Value};

use super::manifest::{InputRequiredSpec, MockManifest, MockTool};

/// Route one parsed JSON-RPC request to the right handler.
///
/// Visible for tests so the per-method behaviors stay easy to pin without
/// piping bytes through `serve`. Errors are returned as JSON-RPC error
/// envelopes rather than `Result::Err` so a single malformed call cannot
/// take the stdio loop down.
pub(crate) fn dispatch(manifest: &MockManifest, method: &str, params: &Value, id: Value) -> Value {
    match method {
        "initialize" => json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": {
                "protocolVersion": "2025-06-18",
                "serverInfo": {
                    "name": manifest.name,
                    "version": mcptest_core::version::semver(),
                },
                "capabilities": capabilities(manifest),
            }
        }),
        // A notification carries no id and expects no response; the client
        // (and `ping`) just need the connection to stay healthy.
        "notifications/initialized" => json!({ "jsonrpc": "2.0", "id": id, "result": {} }),
        "ping" => json!({ "jsonrpc": "2.0", "id": id, "result": {} }),
        "tools/list" => json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": { "tools": render_tools(manifest) }
        }),
        "tools/call" => match handle_tool_call(manifest, params) {
            Ok(result) => json!({
                "jsonrpc": "2.0",
                "id": id,
                "result": result,
            }),
            Err(rpc) => error_envelope(id, rpc.code, rpc.message),
        },
        "resources/list" => json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": { "resources": render_resources(manifest) }
        }),
        "resources/read" => match handle_resource_read(manifest, params) {
            Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }),
            Err(rpc) => error_envelope(id, rpc.code, rpc.message),
        },
        "prompts/list" => json!({
            "jsonrpc": "2.0",
            "id": id,
            "result": { "prompts": render_prompts(manifest) }
        }),
        "prompts/get" => match handle_prompt_get(manifest, params) {
            Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }),
            Err(rpc) => error_envelope(id, rpc.code, rpc.message),
        },
        _ => error_envelope(id, -32601, format!("method not found: {method}")),
    }
}

/// Advertise only the primitives the manifest actually exposes, so a client's
/// capability checks line up with what `*/list` will return.
fn capabilities(manifest: &MockManifest) -> Value {
    let mut caps = serde_json::Map::new();
    // Advertise `tools.listChanged` only when a rotation catalog is declared,
    // so a client's capability check lines up with whether the mock can emit a
    // mid-session `notifications/tools/list_changed` (WOR-1471).
    let tools_list_changed = manifest.list_changed.is_some();
    caps.insert(
        "tools".to_string(),
        json!({ "listChanged": tools_list_changed }),
    );
    if !manifest.resources.is_empty() {
        caps.insert("resources".to_string(), json!({ "listChanged": false }));
    }
    if !manifest.prompts.is_empty() {
        caps.insert("prompts".to_string(), json!({ "listChanged": false }));
    }
    // A task tool means the mock speaks the Tasks extension (SEP-2663),
    // so advertise it the way a real server does (WOR-1472).
    if manifest.tools.iter().any(|t| t.task.is_some()) {
        caps.insert(
            "extensions".to_string(),
            json!({ "io.modelcontextprotocol.tasks": { "version": "1.0" } }),
        );
    }
    Value::Object(caps)
}

fn render_resources(manifest: &MockManifest) -> Vec<Value> {
    manifest
        .resources
        .iter()
        .map(|r| {
            let mut entry = serde_json::Map::new();
            entry.insert("uri".to_string(), Value::String(r.uri.clone()));
            if let Some(name) = &r.name {
                entry.insert("name".to_string(), Value::String(name.clone()));
            }
            if let Some(mime) = &r.mime_type {
                entry.insert("mimeType".to_string(), Value::String(mime.clone()));
            }
            Value::Object(entry)
        })
        .collect()
}

fn render_prompts(manifest: &MockManifest) -> Vec<Value> {
    manifest
        .prompts
        .iter()
        .map(|p| {
            let mut entry = serde_json::Map::new();
            entry.insert("name".to_string(), Value::String(p.name.clone()));
            if let Some(desc) = &p.description {
                entry.insert("description".to_string(), Value::String(desc.clone()));
            }
            Value::Object(entry)
        })
        .collect()
}

fn handle_resource_read(manifest: &MockManifest, params: &Value) -> Result<Value, RpcError> {
    let uri = params
        .get("uri")
        .and_then(Value::as_str)
        .ok_or_else(|| RpcError {
            code: -32602,
            message: "resources/read requires a `uri` field".into(),
        })?;
    let resource = manifest
        .resources
        .iter()
        .find(|r| r.uri == uri)
        .ok_or_else(|| RpcError {
            code: -32602,
            message: format!("unknown resource uri: {uri}"),
        })?;
    let mut content = serde_json::Map::new();
    content.insert("uri".to_string(), Value::String(resource.uri.clone()));
    if let Some(mime) = &resource.mime_type {
        content.insert("mimeType".to_string(), Value::String(mime.clone()));
    }
    content.insert(
        "text".to_string(),
        Value::String(resource.text.clone().unwrap_or_default()),
    );
    Ok(json!({ "contents": [Value::Object(content)] }))
}

fn handle_prompt_get(manifest: &MockManifest, params: &Value) -> Result<Value, RpcError> {
    let name = params
        .get("name")
        .and_then(Value::as_str)
        .ok_or_else(|| RpcError {
            code: -32602,
            message: "prompts/get requires a `name` field".into(),
        })?;
    let prompt = manifest
        .prompts
        .iter()
        .find(|p| p.name == name)
        .ok_or_else(|| RpcError {
            code: -32602,
            message: format!("unknown prompt: {name}"),
        })?;
    Ok(json!({
        "messages": [{
            "role": "user",
            "content": { "type": "text", "text": prompt.text.clone().unwrap_or_default() }
        }]
    }))
}

fn render_tools(manifest: &MockManifest) -> Vec<Value> {
    manifest
        .tools
        .iter()
        .map(|tool| {
            let mut entry = serde_json::Map::new();
            entry.insert("name".to_string(), Value::String(tool.name.clone()));
            if let Some(desc) = &tool.description {
                entry.insert("description".to_string(), Value::String(desc.clone()));
            }
            if let Some(schema) = &tool.input_schema {
                entry.insert("inputSchema".to_string(), schema.clone());
            }
            if let Some(schema) = &tool.output_schema {
                entry.insert("outputSchema".to_string(), schema.clone());
            }
            if let Some(annotations) = &tool.annotations {
                entry.insert("annotations".to_string(), annotations.clone());
            }
            Value::Object(entry)
        })
        .collect()
}

struct RpcError {
    code: i64,
    message: String,
}

/// Run a tool call and return its bare result payload (the value that
/// goes in the JSON-RPC `result`), or a `(code, message)` JSON-RPC error.
///
/// The Tasks lifecycle uses this to compute a task's terminal `completed`
/// payload: the same value a synchronous `tools/call` would have returned
/// (WOR-1472). Crate-visible so `mock/tasks.rs` can reuse the validation
/// and interpolation the synchronous path already does.
pub(super) fn tool_call_result(
    manifest: &MockManifest,
    params: &Value,
) -> Result<Value, (i64, String)> {
    handle_tool_call(manifest, params).map_err(|e| (e.code, e.message))
}

fn handle_tool_call(manifest: &MockManifest, params: &Value) -> Result<Value, RpcError> {
    let name = params
        .get("name")
        .and_then(Value::as_str)
        .ok_or_else(|| RpcError {
            code: -32602,
            message: "tools/call requires a `name` field".into(),
        })?;
    let args = params
        .get("arguments")
        .cloned()
        .unwrap_or_else(|| json!({}));
    let tool = manifest
        .tools
        .iter()
        .find(|t| t.name == name)
        .ok_or_else(|| RpcError {
            code: -32601,
            message: format!("tool `{name}` is not declared in the mock"),
        })?;
    // Input-required tools (SEP-2322 multi-round-trip): the first call
    // elicits more input via an `InputRequiredResult`; a retry that
    // echoes the `requestState` is answered normally, with the elicited
    // `inputResponses` folded into the args so the response reflects them
    // (WOR-1472). Schema validation is skipped because the elicited
    // fields travel in `inputResponses`, not `arguments`.
    if let Some(spec) = &tool.input_required {
        let echoed = params.get("requestState").and_then(Value::as_str);
        if echoed != Some(spec.state.as_str()) {
            return Ok(input_required_result(spec));
        }
        let merged = merge_input_responses(&args, params);
        return Ok(canned_response(tool, &merged));
    }
    if let Some(schema) = &tool.input_schema {
        if let Err(msg) = validate_required_args(schema, &args) {
            return Err(RpcError {
                code: -32602,
                message: msg,
            });
        }
    }
    Ok(canned_response(tool, &args))
}

/// Build the `InputRequiredResult` payload a first call returns: the
/// declared `inputRequests`, the opaque `requestState` the client must
/// echo, and the optional top-level `description` (SEP-2322).
fn input_required_result(spec: &InputRequiredSpec) -> Value {
    let mut obj = serde_json::Map::new();
    obj.insert("inputRequests".to_string(), spec.requests.clone());
    obj.insert(
        "requestState".to_string(),
        Value::String(spec.state.clone()),
    );
    if let Some(description) = &spec.description {
        obj.insert(
            "description".to_string(),
            Value::String(description.clone()),
        );
    }
    Value::Object(obj)
}

/// Fold a retry's `inputResponses` (the elicited `{ id, value }` answers)
/// into the call's `arguments` so the canned response can interpolate
/// them via `${args.<id>}`. The original args win on a key collision.
fn merge_input_responses(args: &Value, params: &Value) -> Value {
    let mut map = args.as_object().cloned().unwrap_or_default();
    if let Some(responses) = params.get("inputResponses").and_then(Value::as_array) {
        for response in responses {
            let id = response.get("id").and_then(Value::as_str);
            let value = response.get("value");
            if let (Some(id), Some(value)) = (id, value) {
                map.entry(id.to_string()).or_insert_with(|| value.clone());
            }
        }
    }
    Value::Object(map)
}

/// Light schema validation. Only the required-properties check fires here;
/// full JSON Schema validation is intentionally deferred so the mock stays a
/// small dependency-free surface. A missing required field is the common
/// case for prompt engineering and exercises the `-32602` code path.
fn validate_required_args(schema: &Value, args: &Value) -> Result<(), String> {
    let required = match schema.get("required").and_then(Value::as_array) {
        Some(arr) => arr,
        None => return Ok(()),
    };
    let object = match args.as_object() {
        Some(obj) => obj,
        None => return Err("arguments must be a JSON object".into()),
    };
    for entry in required {
        if let Some(field) = entry.as_str() {
            if !object.contains_key(field) {
                return Err(format!("missing required argument `{field}`"));
            }
        }
    }
    Ok(())
}

fn canned_response(tool: &MockTool, args: &Value) -> Value {
    let template = tool.response.clone().unwrap_or_else(|| {
        json!({
            "content": [{
                "type": "text",
                "text": format!("mock {}", tool.name)
            }]
        })
    });
    let arg_map: HashMap<String, String> = args
        .as_object()
        .map(|obj| {
            obj.iter()
                .map(|(k, v)| (k.clone(), render_scalar(v)))
                .collect()
        })
        .unwrap_or_default();
    interpolate(template, &arg_map)
}

fn render_scalar(value: &Value) -> String {
    match value {
        Value::Null => "null".to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Number(n) => n.to_string(),
        Value::String(s) => s.clone(),
        Value::Array(_) | Value::Object(_) => value.to_string(),
    }
}

/// Recursively replace `${args.<name>}` markers inside string leaves with the
/// matching argument value. Other JSON types pass through untouched.
fn interpolate(value: Value, args: &HashMap<String, String>) -> Value {
    match value {
        Value::String(s) => Value::String(replace_args(&s, args)),
        Value::Array(items) => {
            Value::Array(items.into_iter().map(|v| interpolate(v, args)).collect())
        }
        Value::Object(map) => {
            let mut out = serde_json::Map::with_capacity(map.len());
            for (k, v) in map {
                out.insert(k, interpolate(v, args));
            }
            Value::Object(out)
        }
        other => other,
    }
}

/// Replace `${args.<key>}` markers in `text`. Public to the crate so the serve
/// loop's tests can pin the multibyte-safe behavior directly.
pub(crate) fn replace_args(text: &str, args: &HashMap<String, String>) -> String {
    let mut out = String::with_capacity(text.len());
    let mut rest = text;
    // The `${args.<key>}` marker is pure ASCII, so we locate each one with
    // `str::find` and copy the text before it as a UTF-8 slice. Copying
    // byte-by-byte with `as char` would split multibyte characters and
    // corrupt any non-ASCII tool response.
    while let Some(start) = rest.find("${args.") {
        out.push_str(&rest[..start]);
        let after = &rest[start + "${args.".len()..];
        if let Some(end) = after.find('}') {
            if let Some(value) = args.get(&after[..end]) {
                out.push_str(value);
                rest = &after[end + 1..];
                continue;
            }
        }
        // Unresolved marker: emit it literally and keep scanning past it so a
        // later real marker in the same string still resolves.
        out.push_str("${args.");
        rest = after;
    }
    out.push_str(rest);
    out
}

/// Build a JSON-RPC error envelope. Crate-visible so the serve loop can answer
/// a parse error with the same shape the dispatcher uses.
pub(crate) fn error_envelope(id: Value, code: i64, message: String) -> Value {
    json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": {
            "code": code,
            "message": message,
        }
    })
}
