//! Tasks-extension wire model (SEP-2663).
//!
//! The 2026-07-28 spec graduates Tasks from experimental core to an
//! extension and reshapes its lifecycle around statelessness: a server
//! can answer a `tools/call` with a `CreateTaskResult` (a task handle)
//! instead of the final result, and the client drives the task with
//! `tasks/get` polling and `tasks/cancel`. The result a poll returns
//! carries the task [`status`](TaskStatus) and, on the terminal
//! `completed` state, the `result` the call would have returned
//! synchronously.
//!
//! This module owns the typed wire model plus the recognizers the runner
//! uses to drive the loop: [`recognize_task_handle`] spots a
//! `CreateTaskResult` in a response, and [`parse_task_get`] reads a
//! `tasks/get` reply. The mock server uses the builders
//! ([`create_task_result`], [`task_get_result`]) to emit the same shapes.

use serde::{Deserialize, Serialize};
use serde_json::Value;

/// Discriminator value a `CreateTaskResult` carries in its `resultType`
/// field so a client can tell it apart from a standard result (SEP-2322).
pub const RESULT_TYPE_TASK: &str = "task";

/// Discriminator value the standard result shapes carry (`tasks/get`,
/// `tasks/cancel`, and any non-task call). Pinned so a server emits the
/// spec's discriminator on the poll response.
pub const RESULT_TYPE_COMPLETE: &str = "complete";

/// Lifecycle status of a task (SEP-2663).
///
/// `completed`, `cancelled`, and `failed` are terminal: once reached the
/// task's observable state does not change. `working` and
/// `input_required` are non-terminal, so a client keeps polling.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TaskStatus {
    /// The operation is in progress.
    Working,
    /// The server needs client input before the task can proceed.
    InputRequired,
    /// The operation finished; the poll response carries the `result`.
    Completed,
    /// The operation was cancelled (cancellation is cooperative).
    Cancelled,
    /// A JSON-RPC error occurred; the poll response carries the `error`.
    Failed,
}

impl TaskStatus {
    /// True once the task has reached a state it will not leave.
    pub fn is_terminal(self) -> bool {
        matches!(
            self,
            TaskStatus::Completed | TaskStatus::Cancelled | TaskStatus::Failed
        )
    }
}

/// A task handle / state object (SEP-2663 `Task`).
///
/// The same shape seeds a `CreateTaskResult` and rides every
/// `tasks/get` poll response. Only `taskId` and `status` are required on
/// the wire; the timing hints (`ttlMs`, `pollIntervalMs`) and timestamps
/// are advisory and may be absent.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Task {
    /// Stable identifier the client echoes on every `tasks/get`,
    /// `tasks/update`, and `tasks/cancel` for this task.
    pub task_id: String,
    /// Current lifecycle status.
    pub status: TaskStatus,
    /// Optional human-readable note about the current state.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub status_message: Option<String>,
    /// ISO 8601 creation timestamp, when the server supplies one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub created_at: Option<String>,
    /// ISO 8601 last-update timestamp, when the server supplies one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_updated_at: Option<String>,
    /// Time-to-live from creation in integer milliseconds. Absent means
    /// the server did not advertise one; the spec's `null` (unlimited)
    /// also deserializes here as absent.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub ttl_ms: Option<u64>,
    /// Suggested polling interval in integer milliseconds.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub poll_interval_ms: Option<u64>,
}

/// Parsed `tasks/get` reply: the [`Task`] state plus the terminal
/// payload inlined alongside it.
///
/// On `completed` the `result` field carries what the original call
/// would have returned; on `failed` the `error` field carries the
/// JSON-RPC error. Both are absent while the task is non-terminal.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TaskGetResult {
    /// The task state common to every poll.
    #[serde(flatten)]
    pub task: Task,
    /// Final result, present once `status` is `completed`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    /// JSON-RPC error, present once `status` is `failed`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<Value>,
}

/// Build the `result` block a task-augmented `tools/call` returns: the
/// [`Task`] seed flattened under a `resultType: "task"` discriminator
/// (SEP-2663 `CreateTaskResult`).
pub fn create_task_result(task: &Task) -> Value {
    let mut value =
        serde_json::to_value(task).unwrap_or_else(|_| Value::Object(Default::default()));
    if let Some(map) = value.as_object_mut() {
        map.insert(
            "resultType".to_string(),
            Value::String(RESULT_TYPE_TASK.to_string()),
        );
    }
    value
}

/// Build the `result` block a `tasks/get` poll returns: the [`Task`]
/// state flattened under a `resultType: "complete"` discriminator, with
/// the terminal `result` / `error` inlined when supplied.
pub fn task_get_result(task: &Task, result: Option<Value>, error: Option<Value>) -> Value {
    let mut value =
        serde_json::to_value(task).unwrap_or_else(|_| Value::Object(Default::default()));
    if let Some(map) = value.as_object_mut() {
        map.insert(
            "resultType".to_string(),
            Value::String(RESULT_TYPE_COMPLETE.to_string()),
        );
        if let Some(result) = result {
            map.insert("result".to_string(), result);
        }
        if let Some(error) = error {
            map.insert("error".to_string(), error);
        }
    }
    value
}

/// Recognize a `CreateTaskResult` inside a response envelope.
///
/// A response hands back a task handle when its `result.resultType` is
/// `"task"` and the result parses as a [`Task`]. Returns the typed task
/// when both hold, and `None` for a standard result so the runner does
/// not start polling a call that finished synchronously.
pub fn recognize_task_handle(response: &Value) -> Option<Task> {
    let result = response.get("result")?;
    if result.get("resultType").and_then(Value::as_str)? != RESULT_TYPE_TASK {
        return None;
    }
    serde_json::from_value(result.clone()).ok()
}

/// Parse a `tasks/get` reply into its [`TaskGetResult`].
///
/// Reads the `result` block of the poll response. Returns `None` when
/// the envelope carries no parseable task state (a malformed or
/// error-only reply), so the caller can surface it rather than loop.
pub fn parse_task_get(response: &Value) -> Option<TaskGetResult> {
    let result = response.get("result")?;
    serde_json::from_value(result.clone()).ok()
}

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

    fn working_task() -> Task {
        Task {
            task_id: "task-1".to_string(),
            status: TaskStatus::Working,
            status_message: Some("in progress".to_string()),
            created_at: Some("2026-07-28T10:30:00Z".to_string()),
            last_updated_at: Some("2026-07-28T10:30:00Z".to_string()),
            ttl_ms: Some(60_000),
            poll_interval_ms: Some(5_000),
        }
    }

    #[test]
    fn task_status_serializes_to_the_spec_snake_case_strings() {
        assert_eq!(json!(TaskStatus::Working), json!("working"));
        assert_eq!(json!(TaskStatus::InputRequired), json!("input_required"));
        assert_eq!(json!(TaskStatus::Completed), json!("completed"));
        assert_eq!(json!(TaskStatus::Cancelled), json!("cancelled"));
        assert_eq!(json!(TaskStatus::Failed), json!("failed"));
    }

    #[test]
    fn only_completed_cancelled_and_failed_are_terminal() {
        assert!(!TaskStatus::Working.is_terminal());
        assert!(!TaskStatus::InputRequired.is_terminal());
        assert!(TaskStatus::Completed.is_terminal());
        assert!(TaskStatus::Cancelled.is_terminal());
        assert!(TaskStatus::Failed.is_terminal());
    }

    #[test]
    fn task_uses_camel_case_wire_keys() {
        let value = serde_json::to_value(working_task()).unwrap();
        let obj = value.as_object().unwrap();
        assert!(obj.contains_key("taskId"));
        assert!(obj.contains_key("statusMessage"));
        assert!(obj.contains_key("createdAt"));
        assert!(obj.contains_key("lastUpdatedAt"));
        assert!(obj.contains_key("ttlMs"));
        assert!(obj.contains_key("pollIntervalMs"));
        assert!(!obj.contains_key("task_id"));
        assert!(!obj.contains_key("poll_interval_ms"));
    }

    #[test]
    fn create_task_result_carries_the_task_discriminator_and_seed_fields() {
        let value = create_task_result(&working_task());
        assert_eq!(value["resultType"], json!("task"));
        assert_eq!(value["taskId"], json!("task-1"));
        assert_eq!(value["status"], json!("working"));
        assert_eq!(value["ttlMs"], json!(60_000));
        assert_eq!(value["pollIntervalMs"], json!(5_000));
    }

    #[test]
    fn recognize_task_handle_returns_the_task_for_a_create_task_result() {
        let envelope = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "result": create_task_result(&working_task()),
        });
        let task = recognize_task_handle(&envelope).expect("recognized");
        assert_eq!(task.task_id, "task-1");
        assert_eq!(task.status, TaskStatus::Working);
        assert_eq!(task.poll_interval_ms, Some(5_000));
    }

    #[test]
    fn recognize_task_handle_returns_none_for_a_standard_result() {
        let envelope = json!({
            "jsonrpc": "2.0",
            "id": 1,
            "result": { "content": [{ "type": "text", "text": "done" }] }
        });
        assert!(recognize_task_handle(&envelope).is_none());
    }

    #[test]
    fn recognize_task_handle_ignores_a_complete_discriminator() {
        // A `tasks/get` reply is resultType "complete", not "task": it
        // must not be mistaken for a fresh handle to start polling.
        let envelope = json!({
            "jsonrpc": "2.0",
            "id": 2,
            "result": task_get_result(&working_task(), None, None),
        });
        assert!(recognize_task_handle(&envelope).is_none());
    }

    #[test]
    fn task_get_result_inlines_the_completed_result_under_a_complete_discriminator() {
        let mut done = working_task();
        done.status = TaskStatus::Completed;
        let payload = json!({ "content": [{ "type": "text", "text": "42" }] });
        let value = task_get_result(&done, Some(payload.clone()), None);
        assert_eq!(value["resultType"], json!("complete"));
        assert_eq!(value["status"], json!("completed"));
        assert_eq!(value["result"], payload);
        assert!(value.get("error").is_none());
    }

    #[test]
    fn parse_task_get_reads_status_and_terminal_result() {
        let mut done = working_task();
        done.status = TaskStatus::Completed;
        let payload = json!({ "content": [{ "type": "text", "text": "ok" }] });
        let envelope = json!({
            "jsonrpc": "2.0",
            "id": 3,
            "result": task_get_result(&done, Some(payload.clone()), None),
        });
        let parsed = parse_task_get(&envelope).expect("parsed");
        assert_eq!(parsed.task.status, TaskStatus::Completed);
        assert!(parsed.task.status.is_terminal());
        assert_eq!(parsed.result, Some(payload));
        assert!(parsed.error.is_none());
    }

    #[test]
    fn parse_task_get_reads_a_working_poll_with_no_terminal_payload() {
        let envelope = json!({
            "jsonrpc": "2.0",
            "id": 4,
            "result": task_get_result(&working_task(), None, None),
        });
        let parsed = parse_task_get(&envelope).expect("parsed");
        assert_eq!(parsed.task.status, TaskStatus::Working);
        assert!(!parsed.task.status.is_terminal());
        assert!(parsed.result.is_none());
    }

    #[test]
    fn task_round_trips_through_serde() {
        let task = working_task();
        let bytes = serde_json::to_vec(&task).unwrap();
        let parsed: Task = serde_json::from_slice(&bytes).unwrap();
        assert_eq!(parsed, task);
    }

    #[test]
    fn ttl_ms_null_deserializes_as_absent() {
        // The spec types ttlMs as `number | null`; a server that sends
        // null (unlimited) must not break the parse.
        let task: Task = serde_json::from_value(json!({
            "taskId": "t",
            "status": "working",
            "ttlMs": null
        }))
        .unwrap();
        assert_eq!(task.ttl_ms, None);
    }
}
