//! Agent Request Tracker
//!
//! This module provides infrastructure for tracking agent-initiated requests
//! that require responses from HTTP clients. It's used to implement bidirectional
//! JSON-RPC communication in the ACP Server.
//!
//! ## Use Case
//!
//! When an agent (like Claude) sends a request to the client (like a permission
//! prompt), the ACP Server needs to:
//! 1. Forward the request to the client via SSE
//! 2. Wait for the client's response via HTTP POST
//! 3. Deliver the response back to the agent
//!
//! The `AgentRequestTracker` manages the pending requests and provides a way
//! to correlate responses with their corresponding requests.
//!
//! ## Example Flow
//!
//! ```rust,ignore
//! use dirigent_acp_api::agent_requests::AgentRequestTracker;
//! use serde_json::json;
//!
//! // Create tracker
//! let tracker = AgentRequestTracker::new();
//!
//! // Agent sends request - register it and get receiver
//! let request_id = json!(0);
//! let client_id = "client-123";
//! let receiver = tracker.register(client_id, request_id.clone());
//!
//! // Forward request to client via SSE...
//!
//! // Client responds via HTTP POST - complete the request
//! let response = json!({"selectedOptionId": "allow"});
//! tracker.complete(client_id, request_id, response)?;
//!
//! // The receiver now gets the response
//! let response_value = receiver.await?;
//! ```

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use anyhow::{anyhow, Result};
use serde_json::Value;
use tokio::sync::oneshot;
use tracing::{debug, warn};

/// Tracks pending agent requests awaiting client responses
///
/// This struct provides thread-safe storage for correlating agent requests
/// with their eventual client responses. It uses oneshot channels to deliver
/// responses to waiting tasks.
///
/// ## Thread Safety
///
/// The tracker is designed to be cloned and shared across async tasks.
/// Internal state is protected by `Arc<Mutex<...>>` for thread-safe access.
#[derive(Debug, Clone)]
pub struct AgentRequestTracker {
    /// Maps (client_id, request_id_string) to oneshot sender for delivering response
    ///
    /// The key is a tuple of client ID and request ID (as string) to uniquely
    /// identify each pending request. The value is a oneshot sender that will
    /// be used to deliver the response when it arrives.
    pending: Arc<Mutex<HashMap<(String, String), oneshot::Sender<Value>>>>,
}

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

impl AgentRequestTracker {
    /// Create a new agent request tracker
    ///
    /// Returns an empty tracker ready to register pending requests.
    pub fn new() -> Self {
        Self {
            pending: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// Register a pending agent request and return a receiver for the response
    ///
    /// This method creates a oneshot channel, stores the sender in the pending
    /// requests map, and returns the receiver. The caller can await on the
    /// receiver to get the client's response.
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client that should respond to this request
    /// - `request_id`: The request ID from the agent (JSON-RPC id field)
    ///
    /// # Returns
    ///
    /// A oneshot receiver that will receive the client's response when
    /// `complete()` is called with matching client_id and request_id.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let receiver = tracker.register("client-123", json!(0));
    ///
    /// // Later, when client responds...
    /// tracker.complete("client-123", json!(0), response)?;
    ///
    /// // The receiver gets the response
    /// let response = receiver.await?;
    /// ```
    pub fn register(&self, client_id: &str, request_id: Value) -> oneshot::Receiver<Value> {
        let (tx, rx) = oneshot::channel();
        let key = (client_id.to_string(), request_id.to_string());

        let mut pending = self.pending.lock().expect("Lock poisoned");
        pending.insert(key.clone(), tx);

        debug!(
            client_id = %client_id,
            request_id = %request_id,
            "Registered pending agent request"
        );

        rx
    }

    /// Complete a pending agent request with the client's response
    ///
    /// This method looks up the pending request, sends the response through
    /// the oneshot channel, and removes it from the pending map.
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client sending the response
    /// - `request_id`: The request ID from the original agent request
    /// - `response`: The client's response (JSON-RPC response object)
    ///
    /// # Returns
    ///
    /// - `Ok(())` if the request was found and the response was delivered
    /// - `Err` if the request_id was not found in pending requests
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The request_id is not found in pending requests (may have timed out)
    /// - The receiver has been dropped (unlikely but possible)
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// // Client POSTs response to /acp/agent_response
    /// let response = json!({
    ///     "jsonrpc": "2.0",
    ///     "id": 0,
    ///     "result": {"selectedOptionId": "allow"}
    /// });
    ///
    /// tracker.complete("client-123", json!(0), response)?;
    /// ```
    pub fn complete(&self, client_id: &str, request_id: Value, response: Value) -> Result<()> {
        let key = (client_id.to_string(), request_id.to_string());

        let mut pending = self.pending.lock().expect("Lock poisoned");

        if let Some(sender) = pending.remove(&key) {
            debug!(
                client_id = %client_id,
                request_id = %request_id,
                "Completing pending agent request"
            );

            // Send the response through the oneshot channel
            sender.send(response).map_err(|_| {
                anyhow!(
                    "Failed to send response for request {}: receiver dropped",
                    request_id
                )
            })?;

            Ok(())
        } else {
            warn!(
                client_id = %client_id,
                request_id = %request_id,
                "Attempted to complete non-existent agent request (may have timed out)"
            );

            Err(anyhow!(
                "Request ID {} not found for client {}",
                request_id,
                client_id
            ))
        }
    }

    /// Timeout a pending agent request
    ///
    /// This method removes a pending request from the map and logs a timeout
    /// warning. It should be called when a request has been pending for too
    /// long (e.g., 30 seconds) without a client response.
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client that was supposed to respond
    /// - `request_id`: The request ID that timed out
    ///
    /// # Note
    ///
    /// This method does not send any error response through the channel.
    /// The receiver will get a `RecvError` when it tries to receive, which
    /// the caller should interpret as a timeout.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use tokio::time::{timeout, Duration};
    ///
    /// let receiver = tracker.register("client-123", json!(0));
    ///
    /// // Wait up to 30 seconds for response
    /// match timeout(Duration::from_secs(30), receiver).await {
    ///     Ok(Ok(response)) => {
    ///         // Got response
    ///     }
    ///     Ok(Err(_)) => {
    ///         // Channel closed (timeout was called)
    ///         tracker.timeout("client-123", json!(0));
    ///     }
    ///     Err(_) => {
    ///         // Timeout elapsed
    ///         tracker.timeout("client-123", json!(0));
    ///     }
    /// }
    /// ```
    pub fn timeout(&self, client_id: &str, request_id: Value) {
        let key = (client_id.to_string(), request_id.to_string());

        let mut pending = self.pending.lock().expect("Lock poisoned");

        if pending.remove(&key).is_some() {
            warn!(
                client_id = %client_id,
                request_id = %request_id,
                "Agent request timed out (30s elapsed without client response)"
            );
        }
    }

    /// Get the number of pending requests
    ///
    /// This method returns the total number of requests currently awaiting
    /// client responses across all clients.
    pub fn pending_count(&self) -> usize {
        let pending = self.pending.lock().expect("Lock poisoned");
        pending.len()
    }

    /// Get the number of pending requests for a specific client
    ///
    /// # Parameters
    ///
    /// - `client_id`: The ID of the client to query
    pub fn client_pending_count(&self, client_id: &str) -> usize {
        let pending = self.pending.lock().expect("Lock poisoned");
        pending
            .keys()
            .filter(|(cid, _)| cid == client_id)
            .count()
    }

    /// Clear all pending requests (used when shutting down or on client disconnect)
    ///
    /// This method removes all pending requests from the tracker. The oneshot
    /// senders are dropped, which will cause their receivers to get `RecvError`.
    ///
    /// # Parameters
    ///
    /// - `client_id`: Optional client ID to clear only that client's pending requests.
    ///   If None, clears all pending requests.
    pub fn clear(&self, client_id: Option<&str>) {
        let mut pending = self.pending.lock().expect("Lock poisoned");

        match client_id {
            Some(id) => {
                let keys_to_remove: Vec<_> = pending
                    .keys()
                    .filter(|(cid, _)| cid == id)
                    .cloned()
                    .collect();

                for key in keys_to_remove {
                    pending.remove(&key);
                }

                debug!(
                    client_id = %id,
                    "Cleared all pending agent requests for client"
                );
            }
            None => {
                let count = pending.len();
                pending.clear();

                debug!(
                    count = count,
                    "Cleared all pending agent requests"
                );
            }
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[tokio::test]
    async fn test_register_and_complete() {
        let tracker = AgentRequestTracker::new();

        let client_id = "client-123";
        let request_id = json!(0);
        let response = json!({"result": "success"});

        // Register request
        let receiver = tracker.register(client_id, request_id.clone());
        assert_eq!(tracker.pending_count(), 1);

        // Complete request
        let result = tracker.complete(client_id, request_id, response.clone());
        assert!(result.is_ok());
        assert_eq!(tracker.pending_count(), 0);

        // Receiver should get the response
        let received = receiver.await.unwrap();
        assert_eq!(received, response);
    }

    #[tokio::test]
    async fn test_complete_non_existent_request() {
        let tracker = AgentRequestTracker::new();

        let client_id = "client-123";
        let request_id = json!(999);
        let response = json!({"result": "success"});

        // Try to complete non-existent request
        let result = tracker.complete(client_id, request_id, response);
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_timeout() {
        let tracker = AgentRequestTracker::new();

        let client_id = "client-123";
        let request_id = json!(0);

        // Register request
        let receiver = tracker.register(client_id, request_id.clone());
        assert_eq!(tracker.pending_count(), 1);

        // Timeout the request
        tracker.timeout(client_id, request_id);
        assert_eq!(tracker.pending_count(), 0);

        // Receiver should get error (channel closed)
        let result = receiver.await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_multiple_pending_requests() {
        let tracker = AgentRequestTracker::new();

        let client1 = "client-1";
        let client2 = "client-2";

        // Register multiple requests
        let _rx1 = tracker.register(client1, json!(0));
        let _rx2 = tracker.register(client1, json!(1));
        let _rx3 = tracker.register(client2, json!(0));

        assert_eq!(tracker.pending_count(), 3);
        assert_eq!(tracker.client_pending_count(client1), 2);
        assert_eq!(tracker.client_pending_count(client2), 1);
    }

    #[tokio::test]
    async fn test_clear_all() {
        let tracker = AgentRequestTracker::new();

        let _rx1 = tracker.register("client-1", json!(0));
        let _rx2 = tracker.register("client-2", json!(0));

        assert_eq!(tracker.pending_count(), 2);

        tracker.clear(None);
        assert_eq!(tracker.pending_count(), 0);
    }

    #[tokio::test]
    async fn test_clear_client() {
        let tracker = AgentRequestTracker::new();

        let client1 = "client-1";
        let client2 = "client-2";

        let _rx1 = tracker.register(client1, json!(0));
        let _rx2 = tracker.register(client1, json!(1));
        let _rx3 = tracker.register(client2, json!(0));

        assert_eq!(tracker.pending_count(), 3);

        // Clear only client1's requests
        tracker.clear(Some(client1));
        assert_eq!(tracker.pending_count(), 1);
        assert_eq!(tracker.client_pending_count(client1), 0);
        assert_eq!(tracker.client_pending_count(client2), 1);
    }

    #[test]
    fn test_tracker_clone() {
        let tracker = AgentRequestTracker::new();
        let tracker_clone = tracker.clone();

        // Both should point to same underlying state
        let _rx = tracker.register("client-1", json!(0));
        assert_eq!(tracker_clone.pending_count(), 1);
    }
}
