//! Prompt turn handling for ACP.
//!
//! This module implements session/prompt requests and responses, including:
//! - Content blocks (text, image, audio, resource, resource_link)
//! - Prompt request handling with timeout
//! - Stop reason interpretation
//! - Content validation against agent capabilities

use crate::acp::connector_state::AgentCapabilities;
use crate::acp::transport::{JsonRpcRequest, JsonRpcResult, Transport, TransportError};
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Request to send a prompt to the agent.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SessionPromptRequest {
    /// Session ID to send prompt to.
    pub session_id: String,
    /// Content blocks that make up the prompt.
    pub prompt: Vec<ContentBlock>,
}

impl SessionPromptRequest {
    /// Create a new session/prompt request.
    pub fn new(session_id: impl Into<String>, prompt: Vec<ContentBlock>) -> Self {
        Self {
            session_id: session_id.into(),
            prompt,
        }
    }

    /// Convert to JSON-RPC request.
    pub fn to_jsonrpc(&self, id: u64) -> JsonRpcRequest {
        JsonRpcRequest::new(
            id,
            "session/prompt",
            Some(serde_json::to_value(self).expect("Failed to serialize SessionPromptRequest")),
        )
    }

    /// Validate content blocks against agent capabilities.
    ///
    /// Returns an error if any content block uses a capability the agent
    /// doesn't support.
    pub fn validate(&self, agent_capabilities: &AgentCapabilities) -> Result<(), String> {
        for block in &self.prompt {
            match block {
                ContentBlock::Text { .. } => {
                    // Text is always supported
                }
                ContentBlock::Image { .. } => {
                    if !agent_capabilities.supports_image() {
                        return Err("Agent does not support image content blocks".to_string());
                    }
                }
                ContentBlock::Audio { .. } => {
                    if !agent_capabilities.supports_audio() {
                        return Err("Agent does not support audio content blocks".to_string());
                    }
                }
                ContentBlock::Resource { .. } => {
                    if !agent_capabilities.supports_embedded_context() {
                        return Err(
                            "Agent does not support embedded context (resource blocks)"
                                .to_string(),
                        );
                    }
                }
                ContentBlock::ResourceLink { .. } => {
                    // ResourceLink is always supported
                }
            }
        }
        Ok(())
    }
}

/// Response from session/prompt.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SessionPromptResponse {
    /// Reason the prompt turn stopped.
    pub stop_reason: StopReason,
}

impl SessionPromptResponse {
    /// Parse from JSON-RPC result.
    pub fn from_jsonrpc(result: &serde_json::Value) -> Result<Self, serde_json::Error> {
        serde_json::from_value(result.clone())
    }
}

/// Reason a prompt turn stopped.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
    /// Normal completion (agent finished its turn).
    EndTurn,
    /// Hit token limit for the response.
    MaxTokens,
    /// Hit maximum number of LLM requests per turn.
    MaxTurnRequests,
    /// Agent refused to continue.
    Refusal,
    /// User cancelled the prompt.
    Cancelled,
}

/// Content block in a prompt or message.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
    /// Plain text content.
    Text {
        text: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        annotations: Option<Annotations>,
    },
    /// Image content (base64-encoded).
    Image {
        /// Base64-encoded image data.
        data: String,
        /// MIME type (e.g., "image/png", "image/jpeg").
        mime_type: String,
        /// Optional URI reference.
        #[serde(skip_serializing_if = "Option::is_none")]
        uri: Option<String>,
        #[serde(skip_serializing_if = "Option::is_none")]
        annotations: Option<Annotations>,
    },
    /// Audio content (base64-encoded).
    Audio {
        /// Base64-encoded audio data.
        data: String,
        /// MIME type (e.g., "audio/wav", "audio/mp3").
        mime_type: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        annotations: Option<Annotations>,
    },
    /// Embedded resource (file content or blob).
    Resource {
        resource: EmbeddedResource,
        #[serde(skip_serializing_if = "Option::is_none")]
        annotations: Option<Annotations>,
    },
    /// Link to a resource (reference without full content).
    ResourceLink {
        /// URI of the resource.
        uri: String,
        /// Name/filename of the resource.
        name: String,
        /// MIME type (if known).
        #[serde(skip_serializing_if = "Option::is_none")]
        mime_type: Option<String>,
        /// Human-readable title.
        #[serde(skip_serializing_if = "Option::is_none")]
        title: Option<String>,
        /// Description of the resource.
        #[serde(skip_serializing_if = "Option::is_none")]
        description: Option<String>,
        /// Size in bytes (if known).
        #[serde(skip_serializing_if = "Option::is_none")]
        size: Option<u64>,
        #[serde(skip_serializing_if = "Option::is_none")]
        annotations: Option<Annotations>,
    },
}

/// Embedded resource content.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum EmbeddedResource {
    /// Text resource (e.g., file contents).
    Text {
        uri: String,
        text: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        mime_type: Option<String>,
    },
    /// Binary resource (base64-encoded).
    Blob {
        uri: String,
        /// Base64-encoded binary data.
        blob: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        mime_type: Option<String>,
    },
}

/// Annotations for content blocks.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct Annotations {
    /// Intended audience for this content (e.g., ["assistant"]).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audience: Option<Vec<String>>,
    /// Priority hint (higher = more important).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub priority: Option<f64>,
}

/// Error type for prompt operations.
#[derive(Debug, thiserror::Error)]
pub enum PromptError {
    #[error("Prompt timed out after {0:?}")]
    Timeout(Duration),

    #[error("JSON-RPC error: {0}")]
    JsonRpcError(crate::acp::transport::JsonRpcError),

    #[error("Transport error: {0}")]
    TransportError(#[from] TransportError),

    #[error("Prompt was cancelled")]
    Cancelled,

    #[error("Validation error: {0}")]
    ValidationError(String),
}

impl From<crate::acp::transport::JsonRpcError> for PromptError {
    fn from(err: crate::acp::transport::JsonRpcError) -> Self {
        PromptError::JsonRpcError(err)
    }
}

/// Send a prompt request to the agent.
///
/// # Arguments
///
/// * `transport` - The transport to use for communication
/// * `session_id` - ID of the session to prompt
/// * `prompt` - Content blocks to send
/// * `agent_capabilities` - Agent capabilities (for validation)
///
/// # Returns
///
/// The stop reason for the prompt turn.
pub async fn session_prompt(
    transport: &mut dyn Transport,
    session_id: impl Into<String>,
    prompt: Vec<ContentBlock>,
    agent_capabilities: &AgentCapabilities,
) -> Result<StopReason, PromptError> {
    // Create and validate request
    let request = SessionPromptRequest::new(session_id, prompt);
    request
        .validate(agent_capabilities)
        .map_err(PromptError::ValidationError)?;

    // Convert to JSON-RPC
    let jsonrpc_request = request.to_jsonrpc(10); // Use ID 10+ for prompts

    // Send request
    let result = transport.send_request(jsonrpc_request).await?;

    // Handle response
    match result {
        JsonRpcResult::Success(response) => {
            let prompt_response = SessionPromptResponse::from_jsonrpc(&response.result)
                .map_err(|e| TransportError::SerializationError(e))?;
            Ok(prompt_response.stop_reason)
        }
        JsonRpcResult::Error(error_response) => Err(PromptError::JsonRpcError(error_response.error)),
    }
}

/// Send a prompt request with a timeout.
///
/// This function sends a prompt and waits for a response, but will return
/// a timeout error if no response is received within the specified duration.
pub async fn session_prompt_with_timeout(
    transport: &mut dyn Transport,
    session_id: impl Into<String>,
    prompt: Vec<ContentBlock>,
    agent_capabilities: &AgentCapabilities,
    timeout: Duration,
) -> Result<StopReason, PromptError> {
    let session_id = session_id.into();

    match tokio::time::timeout(
        timeout,
        session_prompt(transport, session_id, prompt, agent_capabilities),
    )
    .await
    {
        Ok(result) => result,
        Err(_) => Err(PromptError::Timeout(timeout)),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::acp::connector_state::PromptCapabilities;

    #[test]
    fn test_session_prompt_request() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Text {
                text: "Hello".to_string(),
                annotations: None,
            }],
        );

        assert_eq!(request.session_id, "session-123");
        assert_eq!(request.prompt.len(), 1);
    }

    #[test]
    fn test_session_prompt_request_serialization() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Text {
                text: "Hello".to_string(),
                annotations: None,
            }],
        );

        let json = serde_json::to_string(&request).unwrap();
        let deserialized: SessionPromptRequest = serde_json::from_str(&json).unwrap();

        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_content_block_text() {
        let block = ContentBlock::Text {
            text: "Test content".to_string(),
            annotations: Some(Annotations {
                audience: Some(vec!["assistant".to_string()]),
                priority: Some(1.0),
            }),
        };

        let json = serde_json::to_string(&block).unwrap();
        assert!(json.contains("\"type\":\"text\""));
        assert!(json.contains("Test content"));
    }

    #[test]
    fn test_content_block_image() {
        let block = ContentBlock::Image {
            data: "base64data".to_string(),
            mime_type: "image/png".to_string(),
            uri: Some("file://test.png".to_string()),
            annotations: None,
        };

        let json = serde_json::to_string(&block).unwrap();
        assert!(json.contains("\"type\":\"image\""));
        assert!(json.contains("base64data"));
        assert!(json.contains("image/png"));
    }

    #[test]
    fn test_content_block_resource_link() {
        let block = ContentBlock::ResourceLink {
            uri: "file:///path/to/file.txt".to_string(),
            name: "file.txt".to_string(),
            mime_type: Some("text/plain".to_string()),
            title: Some("Test File".to_string()),
            description: Some("A test file".to_string()),
            size: Some(1024),
            annotations: None,
        };

        let json = serde_json::to_string(&block).unwrap();
        assert!(json.contains("\"type\":\"resource_link\""));
        assert!(json.contains("file.txt"));
    }

    #[test]
    fn test_embedded_resource_text() {
        let resource = EmbeddedResource::Text {
            uri: "file:///test.txt".to_string(),
            text: "File contents".to_string(),
            mime_type: Some("text/plain".to_string()),
        };

        let json = serde_json::to_string(&resource).unwrap();
        assert!(json.contains("File contents"));
    }

    #[test]
    fn test_stop_reason_serialization() {
        let reasons = vec![
            StopReason::EndTurn,
            StopReason::MaxTokens,
            StopReason::MaxTurnRequests,
            StopReason::Refusal,
            StopReason::Cancelled,
        ];

        for reason in reasons {
            let json = serde_json::to_string(&reason).unwrap();
            let deserialized: StopReason = serde_json::from_str(&json).unwrap();
            assert_eq!(reason, deserialized);
        }
    }

    #[test]
    fn test_validate_text_block() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Text {
                text: "Hello".to_string(),
                annotations: None,
            }],
        );

        let caps = AgentCapabilities {
            load_session: None,
            prompt_capabilities: None,
            mcp: None,
            _meta: None,
        };

        // Text is always supported
        assert!(request.validate(&caps).is_ok());
    }

    #[test]
    fn test_validate_image_unsupported() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Image {
                data: "base64".to_string(),
                mime_type: "image/png".to_string(),
                uri: None,
                annotations: None,
            }],
        );

        let caps = AgentCapabilities {
            load_session: None,
            prompt_capabilities: None,
            mcp: None,
            _meta: None,
        };

        // Image not supported
        assert!(request.validate(&caps).is_err());
    }

    #[test]
    fn test_validate_image_supported() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Image {
                data: "base64".to_string(),
                mime_type: "image/png".to_string(),
                uri: None,
                annotations: None,
            }],
        );

        let caps = AgentCapabilities {
            load_session: None,
            prompt_capabilities: Some(PromptCapabilities {
                image: Some(true),
                audio: None,
                embedded_context: None,
            }),
            mcp: None,
            _meta: None,
        };

        // Image supported
        assert!(request.validate(&caps).is_ok());
    }

    #[test]
    fn test_validate_resource_unsupported() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Resource {
                resource: EmbeddedResource::Text {
                    uri: "file://test".to_string(),
                    text: "content".to_string(),
                    mime_type: None,
                },
                annotations: None,
            }],
        );

        let caps = AgentCapabilities {
            load_session: None,
            prompt_capabilities: None,
            mcp: None,
            _meta: None,
        };

        // Embedded context not supported
        assert!(request.validate(&caps).is_err());
    }

    #[test]
    fn test_validate_resource_supported() {
        let request = SessionPromptRequest::new(
            "session-123",
            vec![ContentBlock::Resource {
                resource: EmbeddedResource::Text {
                    uri: "file://test".to_string(),
                    text: "content".to_string(),
                    mime_type: None,
                },
                annotations: None,
            }],
        );

        let caps = AgentCapabilities {
            load_session: None,
            prompt_capabilities: Some(PromptCapabilities {
                image: None,
                audio: None,
                embedded_context: Some(true),
            }),
            mcp: None,
            _meta: None,
        };

        // Embedded context supported
        assert!(request.validate(&caps).is_ok());
    }
}
