//! Tasks-extension lifecycle for the operator-facing mock (SEP-2663).
//!
//! A tool that declares a `task` spec hands back a `CreateTaskResult`
//! instead of its result; the client then polls `tasks/get` until the
//! task reaches a terminal status, or sends `tasks/cancel`. The
//! per-connection [`TaskRegistry`] tracks each task's status and the
//! terminal `completed` payload (the value the tool would have returned
//! synchronously). Split out so `mock/mod.rs` stays under the module
//! guard.

use std::collections::HashMap;

use serde_json::{json, Value};

use mcptest_core::protocol::tasks::{create_task_result, task_get_result, Task, TaskStatus};

use super::dispatch::{error_envelope, tool_call_result};
use super::manifest::MockManifest;

/// Per-connection Tasks state: in-flight tasks keyed by id, plus the
/// monotonic counter that mints deterministic ids (`mock-task-1`, ...).
#[derive(Default)]
pub(super) struct TaskRegistry {
    tasks: HashMap<String, TaskEntry>,
    seq: u64,
}

/// One mock task's lifecycle state.
struct TaskEntry {
    /// Latched status: `Working` until it reaches a terminal state.
    status: TaskStatus,
    /// `working` polls left before the task completes.
    remaining_polls: u32,
    /// Suggested poll interval echoed on every poll, milliseconds.
    poll_interval_ms: Option<u64>,
    /// The result delivered once the task completes.
    completed_result: Value,
}

/// Build a wire [`Task`] from an id, status, and poll hint. The mock
/// leaves the advisory timestamp / ttl fields unset.
fn make_task(task_id: &str, status: TaskStatus, poll_interval_ms: Option<u64>) -> Task {
    Task {
        task_id: task_id.to_string(),
        status,
        status_message: None,
        created_at: None,
        last_updated_at: None,
        ttl_ms: None,
        poll_interval_ms,
    }
}

/// Handle a Tasks-extension method, returning the full JSON-RPC response
/// envelope, or `None` when the request is not a task method (so the
/// normal dispatcher answers it).
///
/// `tools/call` returns `Some` only when the named tool declares a `task`
/// spec; `tasks/get` and `tasks/cancel` always return `Some`.
pub(super) fn handle_task_method(
    manifest: &MockManifest,
    method: &str,
    params: &Value,
    id: &Value,
    registry: &mut TaskRegistry,
) -> Option<Value> {
    match method {
        "tools/call" => handle_task_tool_call(manifest, params, id, registry),
        "tasks/get" => Some(handle_tasks_get(params, id, registry)),
        "tasks/cancel" => Some(handle_tasks_cancel(params, id, registry)),
        _ => None,
    }
}

/// Answer a `tools/call` whose tool declares a `task` spec with a fresh
/// task handle. Returns `None` for any other tool so the caller falls
/// through to the synchronous dispatcher.
fn handle_task_tool_call(
    manifest: &MockManifest,
    params: &Value,
    id: &Value,
    registry: &mut TaskRegistry,
) -> Option<Value> {
    let name = params.get("name").and_then(Value::as_str)?;
    let tool = manifest.tools.iter().find(|t| t.name == name)?;
    let spec = tool.task.as_ref()?;
    // Compute the terminal payload now: this validates the args exactly
    // as a synchronous call would, so a bad call fails fast instead of
    // minting a task that can never complete cleanly.
    let completed_result = match tool_call_result(manifest, params) {
        Ok(result) => result,
        Err((code, message)) => return Some(error_envelope(id.clone(), code, message)),
    };
    registry.seq += 1;
    let task_id = format!("mock-task-{}", registry.seq);
    let entry = TaskEntry {
        status: TaskStatus::Working,
        remaining_polls: spec.poll_count,
        poll_interval_ms: spec.poll_interval_ms,
        completed_result,
    };
    let seed = make_task(&task_id, entry.status, entry.poll_interval_ms);
    registry.tasks.insert(task_id, entry);
    Some(json!({
        "jsonrpc": "2.0",
        "id": id,
        "result": create_task_result(&seed),
    }))
}

/// Advance and report a task on a `tasks/get` poll: count down the
/// working polls, then latch `completed` with the tool's result.
fn handle_tasks_get(params: &Value, id: &Value, registry: &mut TaskRegistry) -> Value {
    let task_id = match params.get("taskId").and_then(Value::as_str) {
        Some(task_id) => task_id.to_string(),
        None => {
            return error_envelope(
                id.clone(),
                -32602,
                "tasks/get requires a `taskId` field".to_string(),
            )
        }
    };
    let entry = match registry.tasks.get_mut(&task_id) {
        Some(entry) => entry,
        None => return error_envelope(id.clone(), -32602, format!("unknown task: {task_id}")),
    };
    if entry.status == TaskStatus::Working {
        if entry.remaining_polls > 0 {
            entry.remaining_polls -= 1;
        } else {
            entry.status = TaskStatus::Completed;
        }
    }
    let task = make_task(&task_id, entry.status, entry.poll_interval_ms);
    let result = if entry.status == TaskStatus::Completed {
        Some(entry.completed_result.clone())
    } else {
        None
    };
    json!({
        "jsonrpc": "2.0",
        "id": id,
        "result": task_get_result(&task, result, None),
    })
}

/// Acknowledge a `tasks/cancel` with an empty `complete` result and
/// latch the task to `cancelled` unless it already reached a terminal
/// state (cancellation is cooperative, SEP-2663).
fn handle_tasks_cancel(params: &Value, id: &Value, registry: &mut TaskRegistry) -> Value {
    let task_id = match params.get("taskId").and_then(Value::as_str) {
        Some(task_id) => task_id.to_string(),
        None => {
            return error_envelope(
                id.clone(),
                -32602,
                "tasks/cancel requires a `taskId` field".to_string(),
            )
        }
    };
    match registry.tasks.get_mut(&task_id) {
        Some(entry) => {
            if !entry.status.is_terminal() {
                entry.status = TaskStatus::Cancelled;
            }
            json!({
                "jsonrpc": "2.0",
                "id": id,
                "result": { "resultType": "complete" },
            })
        }
        None => error_envelope(id.clone(), -32602, format!("unknown task: {task_id}")),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::commands::mock::manifest::{MockTool, TaskSpec};

    fn registry_with_task(poll_count: u32) -> (MockManifest, TaskRegistry) {
        let manifest = MockManifest {
            name: "task-mock".to_string(),
            tools: vec![MockTool {
                name: "job".to_string(),
                response: Some(json!({ "content": [{ "type": "text", "text": "ok" }] })),
                task: Some(TaskSpec {
                    poll_count,
                    poll_interval_ms: Some(25),
                }),
                ..MockTool::default()
            }],
            resources: Vec::new(),
            prompts: Vec::new(),
            list_changed: None,
        };
        (manifest, TaskRegistry::default())
    }

    fn call_params() -> Value {
        json!({ "name": "job", "arguments": {} })
    }

    #[test]
    fn a_task_tool_call_mints_a_deterministic_handle() {
        let (manifest, mut reg) = registry_with_task(1);
        let resp = handle_task_method(&manifest, "tools/call", &call_params(), &json!(1), &mut reg)
            .unwrap();
        assert_eq!(resp["result"]["taskId"], "mock-task-1");
        assert_eq!(resp["result"]["resultType"], "task");
        assert_eq!(resp["result"]["pollIntervalMs"], 25);
    }

    #[test]
    fn a_non_task_tool_call_is_left_for_the_dispatcher() {
        let (mut manifest, mut reg) = registry_with_task(1);
        manifest.tools[0].task = None;
        let result =
            handle_task_method(&manifest, "tools/call", &call_params(), &json!(1), &mut reg);
        assert!(result.is_none());
    }

    #[test]
    fn polling_counts_down_then_completes() {
        let (manifest, mut reg) = registry_with_task(1);
        handle_task_method(&manifest, "tools/call", &call_params(), &json!(1), &mut reg);
        let get = json!({ "taskId": "mock-task-1" });
        let first = handle_task_method(&manifest, "tasks/get", &get, &json!(2), &mut reg).unwrap();
        assert_eq!(first["result"]["status"], "working");
        let second = handle_task_method(&manifest, "tasks/get", &get, &json!(3), &mut reg).unwrap();
        assert_eq!(second["result"]["status"], "completed");
        assert_eq!(second["result"]["result"]["content"][0]["text"], "ok");
    }

    #[test]
    fn cancel_latches_cancelled() {
        let (manifest, mut reg) = registry_with_task(9);
        handle_task_method(&manifest, "tools/call", &call_params(), &json!(1), &mut reg);
        let by_id = json!({ "taskId": "mock-task-1" });
        handle_task_method(&manifest, "tasks/cancel", &by_id, &json!(2), &mut reg);
        let get = handle_task_method(&manifest, "tasks/get", &by_id, &json!(3), &mut reg).unwrap();
        assert_eq!(get["result"]["status"], "cancelled");
    }
}
