//! Task management for async task execution
//!
//! Provides storage and lifecycle management for MCP tasks.

use std::collections::HashSet;
use std::sync::Arc;
use std::time::{Duration, Instant};

use chrono::Utc;
use dashmap::DashMap;
use serde_json::Value;
use state_machines::state_machine;
use thiserror::Error;
use tokio::sync::{Mutex, Notify, RwLock, mpsc, oneshot};
use uuid::Uuid;

use crate::protocol::types::{JsonRpcRequest, Task, TaskStatus};
use crate::transport::traits::JsonRpcNotification;

state_machine! {
    name: TaskFlow,
    dynamic: true,
    async: true,
    initial: Working,
    states: [Working, InputRequired, Completed, Failed, Cancelled],
    events {
        require_input {
            transition: { from: Working, to: InputRequired }
        }
        resume {
            transition: { from: InputRequired, to: Working }
        }
        complete {
            transition: { from: [Working, InputRequired], to: Completed }
        }
        fail {
            transition: { from: [Working, InputRequired], to: Failed }
        }
        cancel {
            transition: { from: [Working, InputRequired], to: Cancelled }
        }
    }
}

/// Task-related errors
#[derive(Debug, Error)]
pub enum TaskError {
    /// Task not found
    #[error("Task not found: {0}")]
    NotFound(String),

    /// Task timeout
    #[error("Task timeout after {0:?}")]
    Timeout(Duration),

    /// Task failed
    #[error("Task failed: {0}")]
    Failed(String),

    /// Task already in terminal state
    #[error("Task already in terminal status: {0:?}")]
    AlreadyTerminal(TaskStatus),

    /// Invalid lifecycle transition requested for this task
    #[error("Invalid task transition: {from:?} -> {to:?}")]
    InvalidTransition { from: TaskStatus, to: TaskStatus },

    /// Internal error
    #[error("Internal error: {0}")]
    Internal(String),
}

fn is_terminal_status(status: TaskStatus) -> bool {
    matches!(
        status,
        TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Cancelled
    )
}

fn lifecycle_status_from_state(state: &str) -> Option<TaskStatus> {
    match state {
        "Working" => Some(TaskStatus::Working),
        "InputRequired" => Some(TaskStatus::InputRequired),
        "Completed" => Some(TaskStatus::Completed),
        "Failed" => Some(TaskStatus::Failed),
        "Cancelled" => Some(TaskStatus::Cancelled),
        _ => None,
    }
}

fn transition_event(from: TaskStatus, to: TaskStatus) -> Result<Option<TaskFlowEvent>, TaskError> {
    if from == to {
        return Ok(None);
    }

    match (from, to) {
        (TaskStatus::Working, TaskStatus::InputRequired) => Ok(Some(TaskFlowEvent::RequireInput)),
        (TaskStatus::InputRequired, TaskStatus::Working) => Ok(Some(TaskFlowEvent::Resume)),
        (TaskStatus::Working | TaskStatus::InputRequired, TaskStatus::Completed) => {
            Ok(Some(TaskFlowEvent::Complete))
        }
        (TaskStatus::Working | TaskStatus::InputRequired, TaskStatus::Failed) => {
            Ok(Some(TaskFlowEvent::Fail))
        }
        (TaskStatus::Working | TaskStatus::InputRequired, TaskStatus::Cancelled) => {
            Ok(Some(TaskFlowEvent::Cancel))
        }
        _ => Err(TaskError::InvalidTransition { from, to }),
    }
}

fn cancellation_result() -> Value {
    serde_json::json!({
        "content": [{
            "type": "text",
            "text": "Cancelled by request"
        }],
        "isError": true
    })
}

/// Internal task entry with result channel
struct TaskEntry {
    /// Public task metadata
    task: Arc<RwLock<Task>>,
    /// Internal lifecycle machine used to validate async status transitions
    lifecycle: Arc<Mutex<DynamicTaskFlow<()>>>,
    /// Session ID owning this task
    session_id: String,
    /// Original request that created the task (reserved for future use)
    #[allow(dead_code)]
    original_request: JsonRpcRequest,
    /// Channel for waiting on result
    result_tx: Option<oneshot::Sender<Value>>,
    /// Stored result (when completed)
    result: Arc<RwLock<Option<Value>>>,
    /// Wakes any waiters when lifecycle or result changes
    changed: Arc<Notify>,
    /// Expiration time
    expires_at: Option<Instant>,
}

/// Task store for managing async task execution
pub struct TaskStore {
    /// All tasks indexed by task ID
    tasks: DashMap<String, TaskEntry>,
    /// Session → task IDs index
    by_session: DashMap<String, HashSet<String>>,
    /// Default TTL for new tasks
    default_ttl: Duration,
    /// Default poll interval
    default_poll_interval: Duration,
    /// Optional notification sender for `notifications/tasks/status` (2025-11-25)
    notification_tx: Option<mpsc::UnboundedSender<JsonRpcNotification>>,
}

impl TaskStore {
    /// Create a new task store
    pub fn new(default_ttl: Duration, default_poll_interval: Duration) -> Self {
        Self {
            tasks: DashMap::new(),
            by_session: DashMap::new(),
            default_ttl,
            default_poll_interval,
            notification_tx: None,
        }
    }

    /// Create a task store wired to a notification channel.
    ///
    /// When tasks transition to a non-working state the store will emit
    /// `notifications/tasks/status` to the supplied sender (per 2025-11-25 spec).
    pub fn with_notifications(
        default_ttl: Duration,
        default_poll_interval: Duration,
        notification_tx: mpsc::UnboundedSender<JsonRpcNotification>,
    ) -> Self {
        Self {
            tasks: DashMap::new(),
            by_session: DashMap::new(),
            default_ttl,
            default_poll_interval,
            notification_tx: Some(notification_tx),
        }
    }

    /// Send a `notifications/tasks/status` notification for the given task.
    ///
    /// Called after any status-changing operation. Per spec receivers MAY send this;
    /// requestors MUST NOT rely on it (polling via `tasks/get` is the primary mechanism).
    fn emit_task_status(&self, task: &Task) {
        let Some(tx) = &self.notification_tx else {
            return;
        };

        let params = serde_json::json!({
            "taskId": task.task_id,
            "status": task.status,
            "createdAt": task.created_at,
            "lastUpdatedAt": task.last_updated_at,
            "ttl": task.ttl,
            "pollInterval": task.poll_interval,
        });

        let notification = JsonRpcNotification::new("notifications/tasks/status", Some(params));
        let _ = tx.send(notification);
    }

    /// Create a new task
    pub fn create_task(
        &self,
        session_id: &str,
        original_request: JsonRpcRequest,
        requested_ttl: Option<u64>,
    ) -> (Task, oneshot::Receiver<Value>) {
        let task_id = Uuid::new_v4().to_string();
        let now = Utc::now().to_rfc3339();

        // Determine TTL (use requested or default)
        let ttl_duration = requested_ttl
            .map(Duration::from_millis)
            .unwrap_or(self.default_ttl);
        let ttl_ms = if ttl_duration == Duration::ZERO {
            None
        } else {
            Some(ttl_duration.as_millis() as u64)
        };

        let task = Task {
            task_id: task_id.clone(),
            status: TaskStatus::Working,
            status_message: None,
            created_at: now.clone(),
            last_updated_at: now,
            ttl: ttl_ms,
            poll_interval: Some(self.default_poll_interval.as_millis() as u64),
        };

        // Create result channel
        let (result_tx, result_rx) = oneshot::channel();

        let entry = TaskEntry {
            task: Arc::new(RwLock::new(task.clone())),
            lifecycle: Arc::new(Mutex::new(DynamicTaskFlow::new(()))),
            session_id: session_id.to_string(),
            original_request,
            result_tx: Some(result_tx),
            result: Arc::new(RwLock::new(None)),
            changed: Arc::new(Notify::new()),
            expires_at: ttl_ms.map(|ms| Instant::now() + Duration::from_millis(ms)),
        };

        // Store task
        self.tasks.insert(task_id.clone(), entry);

        // Index by session
        self.by_session
            .entry(session_id.to_string())
            .or_default()
            .insert(task_id);

        (task, result_rx)
    }

    /// Get task by ID
    pub async fn get_task(&self, task_id: &str) -> Option<Task> {
        let entry = self.tasks.get(task_id)?;
        Some(entry.task.read().await.clone())
    }

    /// Get task for specific session (enforces session isolation)
    pub async fn get_task_for_session(&self, task_id: &str, session_id: &str) -> Option<Task> {
        let entry = self.tasks.get(task_id)?;
        if entry.session_id == session_id {
            Some(entry.task.read().await.clone())
        } else {
            None
        }
    }

    /// List tasks for a session (with pagination)
    pub async fn list_tasks(
        &self,
        session_id: &str,
        cursor: Option<&str>,
        limit: usize,
    ) -> (Vec<Task>, Option<String>) {
        let task_ids = match self.by_session.get(session_id) {
            Some(ids) => ids.clone(),
            None => return (vec![], None),
        };

        let mut task_ids: Vec<_> = task_ids.into_iter().collect();
        task_ids.sort(); // Deterministic ordering

        // Find cursor position
        let start = if let Some(c) = cursor {
            task_ids
                .iter()
                .position(|id| id == c)
                .map(|p| p + 1)
                .unwrap_or(0)
        } else {
            0
        };

        let end = (start + limit).min(task_ids.len());
        let page_ids = &task_ids[start..end];

        // Collect tasks
        let mut tasks = Vec::new();
        for id in page_ids {
            if let Some(task) = self.get_task(id).await {
                tasks.push(task);
            }
        }

        // Next cursor
        let next_cursor = if end < task_ids.len() {
            task_ids.get(end).cloned()
        } else {
            None
        };

        (tasks, next_cursor)
    }

    async fn persist_result_value(
        &self,
        task_id: &str,
        result: Value,
    ) -> Result<Option<Value>, TaskError> {
        let (result_ref, changed) = {
            let entry = self
                .tasks
                .get(task_id)
                .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;
            (Arc::clone(&entry.result), Arc::clone(&entry.changed))
        };

        let mut stored = result_ref.write().await;
        if stored.is_some() {
            return Ok(None);
        }
        *stored = Some(result.clone());
        drop(stored);

        changed.notify_waiters();
        Ok(Some(result))
    }

    async fn send_result_to_waiter(&self, task_id: &str, result: Value) {
        if let Some(mut entry_mut) = self.tasks.get_mut(task_id)
            && let Some(tx) = entry_mut.result_tx.take()
        {
            let _ = tx.send(result);
        }
    }

    async fn transition_task(
        &self,
        task_id: &str,
        new_status: TaskStatus,
        status_message: Option<String>,
        result: Option<Value>,
    ) -> Result<Task, TaskError> {
        let (task_ref, lifecycle_ref, changed) = {
            let entry = self
                .tasks
                .get(task_id)
                .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;
            (
                Arc::clone(&entry.task),
                Arc::clone(&entry.lifecycle),
                Arc::clone(&entry.changed),
            )
        };

        let mut lifecycle = lifecycle_ref.lock().await;
        let current_status = task_ref.read().await.status;
        if is_terminal_status(current_status) {
            return Err(TaskError::AlreadyTerminal(current_status));
        }

        let lifecycle_status =
            lifecycle_status_from_state(lifecycle.current_state()).ok_or_else(|| {
                TaskError::Internal(format!(
                    "unknown task lifecycle state '{}'",
                    lifecycle.current_state()
                ))
            })?;
        if lifecycle_status != current_status {
            return Err(TaskError::Internal(format!(
                "task lifecycle out of sync: lifecycle={lifecycle_status:?}, task={current_status:?}"
            )));
        }

        if let Some(event) = transition_event(current_status, new_status)? {
            lifecycle.handle(event).await.map_err(|err| {
                TaskError::Internal(format!(
                    "task lifecycle transition {current_status:?} -> {new_status:?} failed: {err:?}"
                ))
            })?;
        }

        let resolved_status =
            lifecycle_status_from_state(lifecycle.current_state()).ok_or_else(|| {
                TaskError::Internal(format!(
                    "unknown resolved task lifecycle state '{}'",
                    lifecycle.current_state()
                ))
            })?;

        let result_to_send = if let Some(result) = result {
            self.persist_result_value(task_id, result).await?
        } else {
            None
        };

        let task_snapshot = {
            let mut task = task_ref.write().await;
            task.status = resolved_status;
            task.status_message = status_message;
            task.last_updated_at = Utc::now().to_rfc3339();
            task.clone()
        };

        drop(lifecycle);
        self.emit_task_status(&task_snapshot);
        changed.notify_waiters();

        if let Some(result) = result_to_send {
            self.send_result_to_waiter(task_id, result).await;
        }

        Ok(task_snapshot)
    }

    pub(crate) async fn complete_task(
        &self,
        task_id: &str,
        result: Value,
    ) -> Result<Task, TaskError> {
        self.transition_task(task_id, TaskStatus::Completed, None, Some(result))
            .await
    }

    pub(crate) async fn fail_task(
        &self,
        task_id: &str,
        error_message: String,
    ) -> Result<Task, TaskError> {
        let error_result = serde_json::json!({
            "content": [{
                "type": "text",
                "text": error_message.clone()
            }],
            "isError": true
        });

        self.transition_task(
            task_id,
            TaskStatus::Failed,
            Some(error_message),
            Some(error_result),
        )
        .await
    }

    /// Update task status
    pub async fn update_status(
        &self,
        task_id: &str,
        new_status: TaskStatus,
        status_message: Option<String>,
    ) -> Result<(), TaskError> {
        self.transition_task(task_id, new_status, status_message, None)
            .await
            .map(|_| ())
    }

    /// Store task result
    pub async fn store_result(&self, task_id: &str, result: Value) -> Result<(), TaskError> {
        if let Some(result) = self.persist_result_value(task_id, result).await? {
            self.send_result_to_waiter(task_id, result).await;
        }
        Ok(())
    }

    /// Get stored result
    pub async fn get_result(&self, task_id: &str) -> Option<Value> {
        let entry = self.tasks.get(task_id)?;
        entry.result.read().await.clone()
    }

    /// Wait for task result (blocks until terminal state)
    pub async fn wait_for_result(
        &self,
        task_id: &str,
        timeout: Duration,
    ) -> Result<Value, TaskError> {
        let changed = {
            let entry = self
                .tasks
                .get(task_id)
                .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;
            Arc::clone(&entry.changed)
        };
        let deadline = tokio::time::Instant::now() + timeout;

        loop {
            let notified = changed.notified();
            let task = self
                .get_task(task_id)
                .await
                .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;

            if is_terminal_status(task.status) {
                if let Some(result) = self.get_result(task_id).await {
                    return Ok(result);
                } else {
                    return Err(TaskError::Failed(
                        "Task in terminal state but no result stored".to_string(),
                    ));
                }
            }

            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            if remaining.is_zero() {
                return Err(TaskError::Timeout(timeout));
            }

            if tokio::time::timeout(remaining, notified).await.is_err() {
                return Err(TaskError::Timeout(timeout));
            }
        }
    }

    /// Cancel a task
    pub async fn cancel_task(&self, task_id: &str, session_id: &str) -> Result<Task, TaskError> {
        let entry = self
            .tasks
            .get(task_id)
            .ok_or_else(|| TaskError::NotFound(task_id.to_string()))?;
        if entry.session_id != session_id {
            return Err(TaskError::NotFound(task_id.to_string()));
        }
        drop(entry);

        self.transition_task(
            task_id,
            TaskStatus::Cancelled,
            Some("Cancelled by request".to_string()),
            Some(cancellation_result()),
        )
        .await
    }

    /// Cleanup expired tasks (should be called periodically)
    pub async fn cleanup_expired(&self) {
        let now = Instant::now();
        let mut to_remove = Vec::new();

        for entry in self.tasks.iter() {
            if let Some(expires_at) = entry.expires_at
                && now >= expires_at
            {
                to_remove.push((entry.key().clone(), Arc::clone(&entry.changed)));
            }
        }

        for (task_id, changed) in to_remove {
            if let Some((_, entry)) = self.tasks.remove(&task_id) {
                // Remove from session index
                if let Some(mut ids) = self.by_session.get_mut(&entry.session_id) {
                    ids.remove(&task_id);
                }
                changed.notify_waiters();
            }
        }
    }

    /// Spawn background cleanup task
    pub fn spawn_cleanup_task(self: Arc<Self>, interval: Duration) {
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(interval);
            loop {
                interval.tick().await;
                self.cleanup_expired().await;
            }
        });
    }
}

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

    #[tokio::test]
    async fn test_create_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        assert_eq!(task.status, TaskStatus::Working);
        assert!(task.task_id.len() > 0);
        assert_eq!(task.ttl, Some(60000));
    }

    #[tokio::test]
    async fn test_get_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);
        let retrieved = store.get_task(&task.task_id).await.unwrap();

        assert_eq!(retrieved.task_id, task.task_id);
        assert_eq!(retrieved.status, TaskStatus::Working);
    }

    #[tokio::test]
    async fn test_session_isolation() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        // Should not be visible from different session
        assert!(
            store
                .get_task_for_session(&task.task_id, "session2")
                .await
                .is_none()
        );

        // Should be visible from same session
        assert!(
            store
                .get_task_for_session(&task.task_id, "session1")
                .await
                .is_some()
        );
    }

    #[tokio::test]
    async fn test_update_status() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        store
            .update_status(
                &task.task_id,
                TaskStatus::Completed,
                Some("Done".to_string()),
            )
            .await
            .unwrap();

        let updated = store.get_task(&task.task_id).await.unwrap();
        assert_eq!(updated.status, TaskStatus::Completed);
        assert_eq!(updated.status_message, Some("Done".to_string()));
    }

    #[tokio::test]
    async fn test_input_required_then_resume() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        store
            .update_status(
                &task.task_id,
                TaskStatus::InputRequired,
                Some("Need user input".to_string()),
            )
            .await
            .unwrap();

        let waiting = store.get_task(&task.task_id).await.unwrap();
        assert_eq!(waiting.status, TaskStatus::InputRequired);
        assert_eq!(waiting.status_message, Some("Need user input".to_string()));

        store
            .update_status(&task.task_id, TaskStatus::Working, None)
            .await
            .unwrap();

        let resumed = store.get_task(&task.task_id).await.unwrap();
        assert_eq!(resumed.status, TaskStatus::Working);
        assert_eq!(resumed.status_message, None);
    }

    #[tokio::test]
    async fn test_store_and_get_result() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        let result = json!({"answer": 42});
        store
            .store_result(&task.task_id, result.clone())
            .await
            .unwrap();

        let retrieved = store.get_result(&task.task_id).await.unwrap();
        assert_eq!(retrieved, result);
    }

    #[tokio::test]
    async fn test_cancel_task() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));

        let (task, _rx) = store.create_task("session1", req, None);

        let cancelled = store.cancel_task(&task.task_id, "session1").await.unwrap();
        assert_eq!(cancelled.status, TaskStatus::Cancelled);
        let result = store
            .get_result(&task.task_id)
            .await
            .expect("cancelled task should store a result payload");
        assert_eq!(result["isError"], json!(true));
    }

    #[tokio::test]
    async fn test_wait_for_result_unblocks_on_complete_task() {
        let store = Arc::new(TaskStore::new(
            Duration::from_secs(60),
            Duration::from_secs(5),
        ));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);
        let task_id = task.task_id.clone();

        let waiter_store = Arc::clone(&store);
        let waiter = tokio::spawn(async move {
            waiter_store
                .wait_for_result(&task_id, Duration::from_secs(1))
                .await
        });

        tokio::task::yield_now().await;
        store
            .complete_task(&task.task_id, json!({"answer": 42}))
            .await
            .unwrap();

        let result = waiter.await.unwrap().unwrap();
        assert_eq!(result, json!({"answer": 42}));
    }

    #[tokio::test]
    async fn test_wait_for_result_unblocks_on_fail_task() {
        let store = Arc::new(TaskStore::new(
            Duration::from_secs(60),
            Duration::from_secs(5),
        ));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);
        let task_id = task.task_id.clone();

        let waiter_store = Arc::clone(&store);
        let waiter = tokio::spawn(async move {
            waiter_store
                .wait_for_result(&task_id, Duration::from_secs(1))
                .await
        });

        tokio::task::yield_now().await;
        store
            .fail_task(&task.task_id, "boom".to_string())
            .await
            .unwrap();

        let result = waiter.await.unwrap().unwrap();
        assert_eq!(result["isError"], json!(true));
        assert_eq!(result["content"][0]["text"], json!("boom"));
    }

    #[tokio::test]
    async fn test_concurrent_complete_and_cancel_preserves_single_terminal_outcome() {
        let store = Arc::new(TaskStore::new(
            Duration::from_secs(60),
            Duration::from_secs(5),
        ));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);
        let task_id = task.task_id.clone();

        let complete_store = Arc::clone(&store);
        let complete_task_id = task_id.clone();
        let complete = tokio::spawn(async move {
            complete_store
                .complete_task(&complete_task_id, json!({"answer": 42}))
                .await
        });

        let cancel_store = Arc::clone(&store);
        let cancel_task_id = task_id.clone();
        let cancel =
            tokio::spawn(
                async move { cancel_store.cancel_task(&cancel_task_id, "session1").await },
            );

        let complete_result = complete.await.unwrap();
        let cancel_result = cancel.await.unwrap();

        match (&complete_result, &cancel_result) {
            (Ok(task), Err(TaskError::AlreadyTerminal(TaskStatus::Completed))) => {
                assert_eq!(task.status, TaskStatus::Completed);
            }
            (Err(TaskError::AlreadyTerminal(TaskStatus::Cancelled)), Ok(task)) => {
                assert_eq!(task.status, TaskStatus::Cancelled);
            }
            other => panic!("unexpected concurrent transition outcome: {other:?}"),
        }

        let final_task = store.get_task(&task_id).await.unwrap();
        let final_result = store.get_result(&task_id).await.unwrap();

        match final_task.status {
            TaskStatus::Completed => assert_eq!(final_result, json!({"answer": 42})),
            TaskStatus::Cancelled => assert_eq!(final_result, cancellation_result()),
            status => panic!("unexpected terminal status after race: {status:?}"),
        }
    }

    #[tokio::test]
    async fn test_wait_for_result_unblocks_when_expired_task_is_cleaned_up() {
        let store = Arc::new(TaskStore::new(
            Duration::from_millis(10),
            Duration::from_secs(5),
        ));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);
        let task_id = task.task_id.clone();

        let waiter_store = Arc::clone(&store);
        let waiter = tokio::spawn(async move {
            waiter_store
                .wait_for_result(&task_id, Duration::from_secs(1))
                .await
        });

        tokio::task::yield_now().await;
        tokio::time::sleep(Duration::from_millis(30)).await;
        store.cleanup_expired().await;

        let err = waiter.await.unwrap().unwrap_err();
        match err {
            TaskError::NotFound(id) => assert_eq!(id, task.task_id),
            other => panic!("expected NotFound after cleanup, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn test_list_tasks() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));

        // Create multiple tasks
        for _ in 0..5 {
            let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
            store.create_task("session1", req, None);
        }

        let (tasks, cursor) = store.list_tasks("session1", None, 10).await;
        assert_eq!(tasks.len(), 5);
        assert!(cursor.is_none());
    }

    #[tokio::test]
    async fn test_task_status_notification_on_update() {
        let (tx, mut rx) = mpsc::unbounded_channel();
        let store =
            TaskStore::with_notifications(Duration::from_secs(60), Duration::from_secs(5), tx);
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);

        store
            .update_status(&task.task_id, TaskStatus::Completed, None)
            .await
            .unwrap();

        let notification = tokio::time::timeout(Duration::from_millis(250), rx.recv())
            .await
            .expect("timed out waiting for task status notification")
            .expect("notification channel closed unexpectedly");

        assert_eq!(notification.method, "notifications/tasks/status");
        let params = notification.params.expect("expected notification params");
        assert_eq!(params["status"], json!("completed"));
        assert_eq!(params["taskId"], json!(task.task_id));
    }

    #[tokio::test]
    async fn test_task_status_notification_on_cancel() {
        let (tx, mut rx) = mpsc::unbounded_channel();
        let store =
            TaskStore::with_notifications(Duration::from_secs(60), Duration::from_secs(5), tx);
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);

        store.cancel_task(&task.task_id, "session1").await.unwrap();

        let notification = tokio::time::timeout(Duration::from_millis(250), rx.recv())
            .await
            .expect("timed out waiting for task status notification")
            .expect("notification channel closed unexpectedly");

        assert_eq!(notification.method, "notifications/tasks/status");
        let params = notification.params.expect("expected notification params");
        assert_eq!(params["status"], json!("cancelled"));
        assert_eq!(params["taskId"], json!(task.task_id));
    }

    #[tokio::test]
    async fn test_task_no_notification_without_tx() {
        let store = TaskStore::new(Duration::from_secs(60), Duration::from_secs(5));
        let req = JsonRpcRequest::new(json!(1), "tools/call", Some(json!({})));
        let (task, _rx) = store.create_task("session1", req, None);

        store
            .update_status(&task.task_id, TaskStatus::Completed, None)
            .await
            .unwrap();
    }
}
