//! Cancellation and disconnect handling for ACP.
//!
//! This module implements:
//! - Cancellation cleanup (tool calls, permission requests, state reset)
//! - Disconnect detection and handling
//! - Timeout handling for cancellation

use crate::acp::connector_state::{ConnectionState, SessionState};
use crate::acp::protocol::streaming::{ToolCallInfo, ToolCallStatus};
use crate::acp::transport::{Transport, TransportError, TransportState};
use std::time::Duration;

/// Errors that can occur during cancellation.
#[derive(Debug, thiserror::Error)]
pub enum CancellationError {
    #[error("Transport error during cancellation: {0}")]
    TransportError(#[from] TransportError),

    #[error("Cancellation timed out after {0:?}")]
    Timeout(Duration),

    #[error("Session not found: {0}")]
    SessionNotFound(String),
}

/// Handle cancellation of a prompt turn.
///
/// This function performs comprehensive cleanup when a user cancels a prompt:
/// 1. Marks all pending tool calls as cancelled
/// 2. Sends session/cancel notification to agent
/// 3. Waits for agent to respond with stopReason: cancelled (with timeout)
///
/// # Arguments
///
/// * `session_state` - The session state to clean up
/// * `transport` - The transport to send cancellation notification
///
/// # Returns
///
/// Ok(()) if cancellation completed successfully, or an error if it failed.
pub async fn handle_cancellation(
    session_state: &mut SessionState,
    transport: &mut dyn Transport,
) -> Result<(), CancellationError> {
    // Mark session as cancelling
    session_state.start_cancellation();

    // Send session/cancel notification
    let notification = crate::acp::protocol::session::SessionCancelNotification::new(
        session_state.session_id.clone(),
    );

    let jsonrpc_notification = crate::acp::transport::JsonRpcNotification {
        jsonrpc: "2.0".to_string(),
        method: "session/cancel".to_string(),
        params: Some(
            serde_json::to_value(&notification)
                .expect("Failed to serialize SessionCancelNotification"),
        ),
    };

    transport.send_notification(jsonrpc_notification).await?;

    // Note: Actual waiting for stopReason: cancelled is done by the prompt
    // response handler. This function just initiates the cancellation.

    Ok(())
}

/// Mark all pending tool calls as cancelled.
///
/// This is called immediately when cancellation is initiated, before
/// waiting for the agent to respond.
pub fn cancel_pending_tool_calls(tool_calls: &mut [ToolCallInfo]) {
    for tool_call in tool_calls.iter_mut() {
        if !tool_call.is_terminal() {
            tool_call.status = ToolCallStatus::Cancelled;
        }
    }
}

/// Handle transport disconnect.
///
/// This function is called when the transport connection is lost. It:
/// 1. Marks all pending operations as failed
/// 2. Updates connection state to Disconnected
/// 3. Cleans up resources
///
/// # Arguments
///
/// * `session_state` - The session state to clean up
/// * `connection_state` - The connection state to update
/// * `tool_calls` - Active tool calls to mark as failed
pub fn handle_disconnect(
    session_state: &mut SessionState,
    connection_state: &mut ConnectionState,
    tool_calls: &mut [ToolCallInfo],
) {
    // Mark all pending tool calls as failed
    for tool_call in tool_calls.iter_mut() {
        if !tool_call.is_terminal() {
            tool_call.status = ToolCallStatus::Failed;
        }
    }

    // Clear session flags
    session_state.complete_prompt();
    session_state.cancelling = false;
    session_state.loading = false;

    // Update connection state
    *connection_state = ConnectionState::Disconnected;
}

/// Check if a transport is connected.
pub fn is_transport_connected(transport: &dyn Transport) -> bool {
    transport.is_connected() && transport.state() == TransportState::Connected
}

/// Detect if a disconnect has occurred.
///
/// This checks the transport state and returns true if the connection
/// has been lost.
pub fn detect_disconnect(transport: &dyn Transport) -> bool {
    !is_transport_connected(transport)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::acp::protocol::streaming::ToolKind;

    #[test]
    fn test_cancel_pending_tool_calls() {
        let mut tool_calls = vec![
            ToolCallInfo::new(
                "tool-1".to_string(),
                "Test 1".to_string(),
                ToolKind::Read,
                ToolCallStatus::Pending,
                None,
                None,
            ),
            ToolCallInfo::new(
                "tool-2".to_string(),
                "Test 2".to_string(),
                ToolKind::Edit,
                ToolCallStatus::InProgress,
                None,
                None,
            ),
            ToolCallInfo::new(
                "tool-3".to_string(),
                "Test 3".to_string(),
                ToolKind::Execute,
                ToolCallStatus::Completed,
                None,
                None,
            ),
        ];

        cancel_pending_tool_calls(&mut tool_calls);

        // First two should be cancelled (not terminal)
        assert_eq!(tool_calls[0].status, ToolCallStatus::Cancelled);
        assert_eq!(tool_calls[1].status, ToolCallStatus::Cancelled);

        // Last one should remain completed (already terminal)
        assert_eq!(tool_calls[2].status, ToolCallStatus::Completed);
    }

    #[test]
    fn test_handle_disconnect() {
        let mut session_state =
            SessionState::new("session-123".to_string(), "/path".to_string(), vec![]);
        session_state.start_prompt();

        let mut connection_state = ConnectionState::Ready;

        let mut tool_calls = vec![
            ToolCallInfo::new(
                "tool-1".to_string(),
                "Test".to_string(),
                ToolKind::Read,
                ToolCallStatus::InProgress,
                None,
                None,
            ),
            ToolCallInfo::new(
                "tool-2".to_string(),
                "Test 2".to_string(),
                ToolKind::Edit,
                ToolCallStatus::Completed,
                None,
                None,
            ),
        ];

        handle_disconnect(&mut session_state, &mut connection_state, &mut tool_calls);

        // In-progress tool call should be marked failed
        assert_eq!(tool_calls[0].status, ToolCallStatus::Failed);

        // Completed tool call should remain completed
        assert_eq!(tool_calls[1].status, ToolCallStatus::Completed);

        // Session state should be cleared
        assert!(!session_state.prompt_in_progress);
        assert!(!session_state.cancelling);
        assert!(!session_state.loading);

        // Connection state should be disconnected
        assert_eq!(connection_state, ConnectionState::Disconnected);
    }

    #[test]
    fn test_session_cancellation_state() {
        let mut session_state =
            SessionState::new("session-123".to_string(), "/path".to_string(), vec![]);

        // Start a prompt
        session_state.start_prompt();
        assert!(session_state.prompt_in_progress);
        assert!(!session_state.is_ready());

        // Start cancellation
        session_state.start_cancellation();
        assert!(session_state.cancelling);
        assert!(!session_state.is_ready());

        // Complete the prompt (after cancel)
        session_state.complete_prompt();
        assert!(!session_state.prompt_in_progress);
        assert!(!session_state.cancelling);
        assert!(session_state.is_ready());
    }
}
