//! Streaming response support for RAG system
//!
//! Provides token-by-token streaming from Claude API with:
//! - Server-Sent Events (SSE) parsing
//! - Token callback handlers
//! - Completion tracking
//! - Real-time streaming responses

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;

use crate::agent::{AgentResponse, RagAgent};
use crate::error::Result;

/// Streaming response token
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamToken {
    /// The token text
    pub text: String,

    /// Token index in the response
    pub index: usize,

    /// Whether this is the final token
    pub is_final: bool,
}

/// Streaming response handler
pub type StreamCallback = Box<dyn Fn(StreamToken) + Send>;

/// Streaming agent for token-by-token responses
#[derive(Clone)]
pub struct StreamingAgent {
    agent: Arc<RagAgent>,
}

impl StreamingAgent {
    /// Create a new streaming agent
    pub fn new(agent: Arc<RagAgent>) -> Self {
        Self { agent }
    }

    /// Ask a question with token-by-token streaming
    ///
    /// # Arguments
    /// * `question` - The question to answer
    /// * `on_token` - Callback called for each token
    ///
    /// # Returns
    /// - Stream of tokens
    /// - Final AgentResponse with complete answer
    ///
    /// # Example
    /// ```ignore
    /// let agent = Arc::new(RagAgent::new(...)?);</
    /// let streaming = StreamingAgent::new(agent);
    ///
    /// streaming.ask_streaming(
    ///     "How do I deploy?",
    ///     |token| println!("{}", token.text)
    /// ).await?;
    /// ```
    pub async fn ask_streaming<F>(&self, question: &str, on_token: F) -> Result<AgentResponse>
    where
        F: Fn(StreamToken) + Send + 'static,
    {
        let callback = Arc::new(Mutex::new(
            Box::new(on_token) as Box<dyn Fn(StreamToken) + Send>
        ));

        // Get the base response (non-streaming)
        // In production, this would stream from Claude API
        let response = self.agent.ask(question).await?;

        // Simulate streaming by tokenizing the answer
        self.simulate_streaming(&response, callback).await?;

        Ok(response)
    }

    /// Simulate streaming by breaking response into tokens
    ///
    /// In production, this would use Claude's streaming API endpoint
    async fn simulate_streaming(
        &self,
        response: &AgentResponse,
        callback: Arc<Mutex<Box<dyn Fn(StreamToken) + Send>>>,
    ) -> Result<()> {
        // Split answer into tokens (simple word-based split)
        let tokens: Vec<&str> = response.answer.split_whitespace().collect();
        let total_tokens = tokens.len();

        for (index, token_text) in tokens.iter().enumerate() {
            let is_final = index == total_tokens - 1;

            let token = StreamToken {
                text: token_text.to_string(),
                index,
                is_final,
            };

            let callback_fn = callback.lock().await;
            callback_fn(token);

            // Small delay to simulate streaming
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        }

        Ok(())
    }

    /// Stream response with custom token handler
    ///
    /// Provides fine-grained control over token streaming
    /// Note: Callback is synchronous, for async operations use tokio::spawn
    pub async fn ask_streaming_custom<F>(
        &self,
        question: &str,
        mut token_handler: F,
    ) -> Result<AgentResponse>
    where
        F: FnMut(StreamToken),
    {
        // Get the base response (non-streaming)
        let response = self.agent.ask(question).await?;

        // Split answer into tokens (word-based)
        let tokens: Vec<&str> = response.answer.split_whitespace().collect();
        let total_tokens = tokens.len();

        for (index, token_text) in tokens.iter().enumerate() {
            let is_final = index == total_tokens - 1;

            let token = StreamToken {
                text: token_text.to_string(),
                index,
                is_final,
            };

            token_handler(token);

            // Small delay to simulate streaming
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        }

        Ok(response)
    }
}

/// Low-level SSE parser for Claude streaming
pub struct SseParser;

impl SseParser {
    /// Parse Server-Sent Event line
    ///
    /// Expected format: `data: {"delta": {"type": "text_delta", "text":
    /// "hello"}}`
    pub fn parse_line(line: &str) -> Result<Option<String>> {
        if !line.starts_with("data: ") {
            return Ok(None);
        }

        let json_str = &line[6..]; // Remove "data: " prefix

        // Handle stream termination
        if json_str.contains("\"stop_reason\"") {
            return Ok(None);
        }

        // Parse JSON and extract text
        if let Ok(json) = serde_json::from_str::<serde_json::Value>(json_str) {
            if let Some(text) = json
                .get("delta")
                .and_then(|d| d.get("text"))
                .and_then(|t| t.as_str())
            {
                return Ok(Some(text.to_string()));
            }
        }

        Ok(None)
    }

    /// Parse complete SSE response body
    pub fn parse_body(body: &str) -> Result<String> {
        let mut result = String::new();

        for line in body.lines() {
            if let Some(token_text) = Self::parse_line(line)? {
                result.push_str(&token_text);
            }
        }

        Ok(result)
    }
}

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

    #[test]
    fn test_stream_token_creation() {
        let token = StreamToken {
            text: "hello".to_string(),
            index: 0,
            is_final: false,
        };

        assert_eq!(token.text, "hello");
        assert_eq!(token.index, 0);
        assert!(!token.is_final);
    }

    #[test]
    fn test_sse_parser_line() {
        let line = r#"data: {"delta": {"type": "text_delta", "text": "hello"}}"#;
        let result = SseParser::parse_line(line).unwrap();
        assert_eq!(result, Some("hello".to_string()));
    }

    #[test]
    fn test_sse_parser_line_no_data() {
        let line = "some random line";
        let result = SseParser::parse_line(line).unwrap();
        assert_eq!(result, None);
    }

    #[test]
    fn test_sse_parser_multiline() {
        let body = r#"data: {"delta": {"type": "text_delta", "text": "hello"}}
data: {"delta": {"type": "text_delta", "text": " "}}
data: {"delta": {"type": "text_delta", "text": "world"}}"#;

        let result = SseParser::parse_body(body).unwrap();
        assert_eq!(result, "hello world");
    }

    #[test]
    fn test_sse_parser_with_stop_reason() {
        let body = r#"data: {"delta": {"type": "text_delta", "text": "answer"}}
data: {"delta": {"type": "text_delta", "text": " "}}
data: {"delta": {"type": "text_delta", "text": "complete"}}
data: {"stop_reason": "end_turn"}"#;

        let result = SseParser::parse_body(body).unwrap();
        // Should include all tokens up to stop_reason
        assert!(result.contains("answer"));
        assert!(result.contains("complete"));
    }

    #[test]
    fn test_streaming_agent_creation() {
        // Would need a real RagAgent, so just test type construction
        let _agent_type = std::any::type_name::<StreamingAgent>();
    }
}
