//! Cancellation management for MCP operations
//!
//! Provides a centralized manager for tracking and signaling cancellation of
//! long-running operations. Uses request ID to correlate cancellation requests
//! with active operations.
//!
//! # MCP Protocol
//!
//! When a client sends `notifications/cancelled` with a `requestId`, the server
//! should attempt to cancel the corresponding operation and clean up resources.
//!
//! # Example
//!
//! ```rust,ignore
//! let manager = CancellationManager::new();
//!
//! // Register a cancellable operation
//! let token = manager.register("request-123");
//!
//! // In the operation task
//! tokio::select! {
//!     result = do_work() => { /* completed */ }
//!     _ = token.cancelled() => { /* cancelled */ }
//! }
//!
//! // When client sends cancellation
//! manager.cancel("request-123");
//!
//! // Cleanup when done
//! manager.remove("request-123");
//! ```

use std::sync::Arc;

use dashmap::DashMap;
use serde::Deserialize;
use tokio::sync::watch;

/// Cancellation token for a single operation
///
/// Clone this token to share cancellation state across async tasks.
/// The token is cheap to clone (Arc internally).
#[derive(Clone)]
pub struct CancellationToken {
    /// Watch receiver for cancellation signal
    rx: watch::Receiver<bool>,
}

impl CancellationToken {
    /// Create a new cancellation token with its trigger
    fn new() -> (Self, CancellationTrigger) {
        let (tx, rx) = watch::channel(false);
        (Self { rx }, CancellationTrigger { tx })
    }

    /// Check if cancellation has been requested
    pub fn is_cancelled(&self) -> bool {
        *self.rx.borrow()
    }

    /// Wait until cancellation is requested
    ///
    /// Returns immediately if already cancelled.
    /// Use with `tokio::select!` to race against other operations.
    pub async fn cancelled(&mut self) {
        // If already cancelled, return immediately
        if *self.rx.borrow() {
            return;
        }

        // Wait for change
        loop {
            if self.rx.changed().await.is_err() {
                // Sender dropped - treat as cancelled
                return;
            }
            if *self.rx.borrow() {
                return;
            }
        }
    }
}

/// Internal trigger to signal cancellation
struct CancellationTrigger {
    tx: watch::Sender<bool>,
}

impl CancellationTrigger {
    /// Signal cancellation to all token holders
    fn cancel(&self) {
        let _ = self.tx.send(true);
    }
}

/// Entry tracking a cancellable operation
struct CancellationEntry {
    /// Trigger to signal cancellation
    trigger: CancellationTrigger,
    /// Optional reason for cancellation
    reason: Option<String>,
}

/// Manager for tracking cancellable operations
///
/// Thread-safe and designed for concurrent access from multiple async tasks.
#[derive(Clone)]
pub struct CancellationManager {
    /// Active operations: request_id -> entry
    operations: Arc<DashMap<String, CancellationEntry>>,
}

impl Default for CancellationManager {
    fn default() -> Self {
        Self::new()
    }
}

impl CancellationManager {
    /// Create a new cancellation manager
    pub fn new() -> Self {
        Self {
            operations: Arc::new(DashMap::new()),
        }
    }

    /// Register a new cancellable operation
    ///
    /// Returns a `CancellationToken` that can be polled or awaited.
    /// The request_id should match the JSON-RPC request ID as a string.
    ///
    /// # Arguments
    ///
    /// * `request_id` - The request ID (typically from JSON-RPC id field)
    pub fn register(&self, request_id: impl Into<String>) -> CancellationToken {
        let (token, trigger) = CancellationToken::new();
        let entry = CancellationEntry {
            trigger,
            reason: None,
        };
        self.operations.insert(request_id.into(), entry);
        token
    }

    /// Cancel an operation by request ID
    ///
    /// Signals cancellation to any tasks holding the corresponding token.
    /// No-op if the request_id is not registered.
    ///
    /// # Arguments
    ///
    /// * `request_id` - The request ID to cancel
    /// * `reason` - Optional reason for cancellation
    pub fn cancel(&self, request_id: &str, reason: Option<String>) {
        if let Some(mut entry) = self.operations.get_mut(request_id) {
            entry.reason = reason;
            entry.trigger.cancel();
        }
    }

    /// Check if an operation has been cancelled
    ///
    /// # Arguments
    ///
    /// * `request_id` - The request ID to check
    pub fn is_cancelled(&self, request_id: &str) -> bool {
        self.operations
            .get(request_id)
            .is_some_and(|entry| *entry.trigger.tx.borrow())
    }

    /// Get the cancellation reason if one was provided
    ///
    /// # Arguments
    ///
    /// * `request_id` - The request ID to query
    pub fn get_reason(&self, request_id: &str) -> Option<String> {
        self.operations
            .get(request_id)
            .and_then(|entry| entry.reason.clone())
    }

    /// Remove a tracked operation
    ///
    /// Call this when an operation completes (success, error, or cancelled)
    /// to free resources.
    ///
    /// # Arguments
    ///
    /// * `request_id` - The request ID to remove
    pub fn remove(&self, request_id: &str) {
        self.operations.remove(request_id);
    }

    /// Cancel all active operations
    ///
    /// Used during server shutdown to signal all pending operations.
    pub fn cancel_all(&self) {
        for entry in self.operations.iter() {
            entry.trigger.cancel();
        }
    }

    /// Get count of active tracked operations
    pub fn active_count(&self) -> usize {
        self.operations.len()
    }

    /// Clean up completed operations that weren't explicitly removed
    ///
    /// Removes entries where cancellation was already signaled.
    /// Call periodically if you're not explicitly calling `remove()`.
    pub fn cleanup_cancelled(&self) {
        self.operations
            .retain(|_id, entry| !*entry.trigger.tx.borrow());
    }
}

/// Parameters for notifications/cancelled
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CancelledNotificationParams {
    /// The ID of the request to cancel
    pub request_id: serde_json::Value,
    /// Optional reason for cancellation
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
}

impl CancelledNotificationParams {
    /// Get the request ID as a string for lookup
    pub fn request_id_string(&self) -> String {
        match &self.request_id {
            serde_json::Value::String(s) => s.clone(),
            serde_json::Value::Number(n) => n.to_string(),
            v => v.to_string(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn test_register_and_cancel() {
        let manager = CancellationManager::new();
        let token = manager.register("req-1");

        assert!(!token.is_cancelled());
        assert!(!manager.is_cancelled("req-1"));

        manager.cancel("req-1", Some("user request".to_string()));

        assert!(token.is_cancelled());
        assert!(manager.is_cancelled("req-1"));
        assert_eq!(
            manager.get_reason("req-1"),
            Some("user request".to_string())
        );
    }

    #[test]
    fn test_token_clone() {
        let manager = CancellationManager::new();
        let token1 = manager.register("req-1");
        let token2 = token1.clone();

        assert!(!token1.is_cancelled());
        assert!(!token2.is_cancelled());

        manager.cancel("req-1", None);

        assert!(token1.is_cancelled());
        assert!(token2.is_cancelled());
    }

    #[test]
    fn test_remove() {
        let manager = CancellationManager::new();
        let _token = manager.register("req-1");

        assert_eq!(manager.active_count(), 1);

        manager.remove("req-1");

        assert_eq!(manager.active_count(), 0);
        assert!(!manager.is_cancelled("req-1"));
    }

    #[test]
    fn test_cancel_all() {
        let manager = CancellationManager::new();
        let token1 = manager.register("req-1");
        let token2 = manager.register("req-2");
        let token3 = manager.register("req-3");

        assert!(!token1.is_cancelled());
        assert!(!token2.is_cancelled());
        assert!(!token3.is_cancelled());

        manager.cancel_all();

        assert!(token1.is_cancelled());
        assert!(token2.is_cancelled());
        assert!(token3.is_cancelled());
    }

    #[test]
    fn test_cancel_nonexistent() {
        let manager = CancellationManager::new();
        // Should not panic
        manager.cancel("nonexistent", None);
        assert!(!manager.is_cancelled("nonexistent"));
    }

    #[tokio::test]
    async fn test_cancelled_future() {
        let manager = CancellationManager::new();
        let mut token = manager.register("req-1");

        // Spawn cancellation after a delay
        let manager_clone = manager.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(10)).await;
            manager_clone.cancel("req-1", None);
        });

        // Wait for cancellation
        tokio::time::timeout(Duration::from_millis(100), token.cancelled())
            .await
            .expect("should complete within timeout");

        assert!(token.is_cancelled());
    }

    #[tokio::test]
    async fn test_cancelled_future_immediate() {
        let manager = CancellationManager::new();
        let mut token = manager.register("req-1");

        // Cancel before waiting
        manager.cancel("req-1", None);

        // Should return immediately
        tokio::time::timeout(Duration::from_millis(10), token.cancelled())
            .await
            .expect("should return immediately");
    }

    #[tokio::test]
    async fn test_select_with_cancellation() {
        let manager = CancellationManager::new();
        let mut token = manager.register("req-1");

        // Spawn cancellation
        let manager_clone = manager.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(10)).await;
            manager_clone.cancel("req-1", None);
        });

        let result = tokio::select! {
            _ = token.cancelled() => "cancelled",
            _ = tokio::time::sleep(Duration::from_secs(10)) => "completed",
        };

        assert_eq!(result, "cancelled");
    }

    #[test]
    fn test_cleanup_cancelled() {
        let manager = CancellationManager::new();
        let _token1 = manager.register("req-1");
        let _token2 = manager.register("req-2");

        manager.cancel("req-1", None);

        assert_eq!(manager.active_count(), 2);

        manager.cleanup_cancelled();

        assert_eq!(manager.active_count(), 1);
        assert!(!manager.is_cancelled("req-1")); // removed
    }

    #[test]
    fn test_params_parsing() {
        let json = r#"{"requestId": 123, "reason": "user cancelled"}"#;
        let params: CancelledNotificationParams = serde_json::from_str(json).unwrap();
        assert_eq!(params.request_id_string(), "123");
        assert_eq!(params.reason, Some("user cancelled".to_string()));

        let json = r#"{"requestId": "abc-123"}"#;
        let params: CancelledNotificationParams = serde_json::from_str(json).unwrap();
        assert_eq!(params.request_id_string(), "abc-123");
        assert_eq!(params.reason, None);
    }
}
