//! Anthropic streaming response implementation.
//!
//! This module provides the streaming response parser for Anthropic's
//! event-based streaming format.

use aix_core::{
    error::{AixError, AixResult},
    streaming::SseParser,
    types::StreamChunk,
};
use bytes::Bytes;
use futures_core::Stream;
use pin_project_lite::pin_project;
use serde::Deserialize;
use std::pin::Pin;
use std::task::{Context, Poll};

pin_project! {
    /// Streaming response parser for Anthropic event format.
    pub struct AnthropicStreamingResponse<S> {
        #[pin]
        byte_stream: S,
        parser: SseParser,
        buffer: Vec<u8>,
        current_id: Option<String>,
    }
}

impl<S> AnthropicStreamingResponse<S>
where
    S: Stream<Item = std::io::Result<Bytes>> + Send,
{
    /// Create a new streaming response parser.
    pub fn new(byte_stream: S) -> Self {
        Self {
            byte_stream,
            parser: SseParser::new(),
            buffer: Vec::new(),
            current_id: None,
        }
    }
}

impl<S> Stream for AnthropicStreamingResponse<S>
where
    S: Stream<Item = std::io::Result<Bytes>> + Send,
{
    type Item = AixResult<StreamChunk>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        loop {
            // Try to parse any buffered data first
            if !this.buffer.is_empty() {
                match this.parser.parse_chunk(&this.buffer) {
                    Ok(events) => {
                        // Clear the buffer since it's been processed
                        this.buffer.clear();
                        
                        for event in &events {
                            // Parse the JSON event
                            match serde_json::from_str::<AnthropicStreamEvent>(&event) {
                                Ok(stream_event) => {
                                    match stream_event {
                                        AnthropicStreamEvent::MessageStart { message } => {
                                            // Store the message ID for subsequent chunks
                                            *this.current_id = Some(message.id);
                                        }
                                        AnthropicStreamEvent::ContentBlockDelta { delta } => {
                                            if let Some(text) = delta.text {
                                                if !text.is_empty() {
                                                    let id = this.current_id.as_ref()
                                                        .unwrap_or(&"unknown".to_string())
                                                        .clone();
                                                    let chunk = StreamChunk::new(id.clone(), text.clone());
                                                    return Poll::Ready(Some(Ok(chunk)));
                                                }
                                            }
                                        }
                                        AnthropicStreamEvent::MessageDelta { delta } => {
                                            if let Some(stop_reason) = delta.stop_reason {
                                                let id = this.current_id.as_ref()
                                                    .unwrap_or(&"unknown".to_string())
                                                    .clone();
                                                let chunk = StreamChunk::with_finish_reason(id.clone(), stop_reason.to_string());
                                                return Poll::Ready(Some(Ok(chunk)));
                                            }
                                        }
                                        AnthropicStreamEvent::MessageStop => {
                                            // End of stream
                                            return Poll::Ready(None);
                                        }
                                        // Ignore other events
                                        _ => {}
                                    }
                                }
                                Err(e) => {
                                    // Try to parse as a regular JSON object (some events might not have the type field)
                                    if let Ok(value) = serde_json::from_str::<serde_json::Value>(&event) {
                                        if let Some(event_type) = value.get("type").and_then(|t| t.as_str()) {
                                            match event_type {
                                                "content_block_delta" => {
                                                    if let Some(delta) = value.get("delta") {
                                                        if let Some(text) = delta.get("text").and_then(|t| t.as_str()) {
                                                            if !text.is_empty() {
                                                                let id = this.current_id.as_ref()
                                                                    .unwrap_or(&"unknown".to_string())
                                                                    .clone();
                                                                let chunk = StreamChunk::new(id.clone(), text.to_string());
                                                                return Poll::Ready(Some(Ok(chunk)));
                                                            }
                                                        }
                                                    }
                                                }
                                                "message_delta" => {
                                                    if let Some(delta) = value.get("delta") {
                                                        if let Some(stop_reason) = delta.get("stop_reason").and_then(|s| s.as_str()) {
                                                            let id = this.current_id.as_ref()
                                                                .unwrap_or(&"unknown".to_string())
                                                                .clone();
                                                            let chunk = StreamChunk::with_finish_reason(id.clone(), stop_reason.to_string());
                                                            return Poll::Ready(Some(Ok(chunk)));
                                                        }
                                                    }
                                                }
                                                "message_stop" => {
                                                    return Poll::Ready(None);
                                                }
                                                _ => {}
                                            }
                                        }
                                    } else {
                                        return Poll::Ready(Some(Err(AixError::serialization(
                                            e.to_string(),
                                            "parsing Anthropic stream event",
                                        ))));
                                    }
                                }
                            }
                        }
                    }
                    Err(e) => {
                        return Poll::Ready(Some(Err(e)));
                    }
                }
            }

            // Need more data from the byte stream
            match futures_core::ready!(this.byte_stream.as_mut().poll_next(cx)) {
                Some(Ok(bytes)) => {
                    if bytes.is_empty() {
                        continue;
                    }
                    this.buffer.extend_from_slice(&bytes);
                }
                Some(Err(e)) => {
                    return Poll::Ready(Some(Err(AixError::stream_with_source(
                        "Error reading from stream",
                        e.to_string(),
                    ))));
                }
                None => {
                    // Stream ended, try to parse any remaining data
                    if !this.buffer.is_empty() {
                        match this.parser.parse_chunk(&this.buffer) {
                            Ok(events) => {
                                this.buffer.clear();
                                for event in &events {
                                    match serde_json::from_str::<AnthropicStreamEvent>(&event) {
                                        Ok(AnthropicStreamEvent::MessageStop) => {
                                            return Poll::Ready(None);
                                        }
                                        _ => {
                                            // Ignore other events at the end of stream
                                            continue;
                                        }
                                    }
                                }
                            }
                            Err(_) => {
                                // Ignore parsing errors at the end of stream
                            }
                        }
                    }
                    return Poll::Ready(None);
                }
            }
        }
    }
}

/// Anthropic streaming response format.
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum AnthropicStreamEvent {
    MessageStart {
        message: AnthropicMessageStart,
    },
    ContentBlockStart {
        content_block: AnthropicContentBlockStart,
    },
    ContentBlockDelta {
        delta: AnthropicDelta,
    },
    ContentBlockStop,
    MessageDelta {
        delta: AnthropicMessageDelta,
    },
    MessageStop,
}

/// Anthropic message start event.
#[derive(Debug, Deserialize)]
struct AnthropicMessageStart {
    id: String,
    type_: String,
    role: String,
    content: Vec<AnthropicResponseContent>,
    model: String,
    usage: AnthropicUsage,
}

/// Anthropic content block start event.
#[derive(Debug, Deserialize)]
struct AnthropicContentBlockStart {
    type_: String,
    text: String,
}

/// Anthropic delta event.
#[derive(Debug, Deserialize)]
struct AnthropicDelta {
    type_: String,
    text: Option<String>,
}

/// Anthropic message delta event.
#[derive(Debug, Deserialize)]
struct AnthropicMessageDelta {
    stop_reason: Option<String>,
    stop_sequence: Option<String>,
}

/// Anthropic response content format.
#[derive(Debug, Deserialize)]
struct AnthropicResponseContent {
    type_: String,
    text: String,
}

/// Anthropic usage format.
#[derive(Debug, Deserialize)]
struct AnthropicUsage {
    input_tokens: u32,
    output_tokens: u32,
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::StreamExt as FuturesStreamExt;

    #[tokio::test]
    async fn test_streaming_response_parsing() {
        // Mock Anthropic streaming data
        let sse_data = vec![
            b"event: message_start\ndata: {\"type\": \"message_start\", \"message\": {\"id\": \"msg_123\", \"type\": \"message\", \"role\": \"assistant\", \"content\": [], \"model\": \"claude-3-sonnet\", \"usage\": {\"input_tokens\": 10, \"output_tokens\": 0}}}\n\n",
            b"event: content_block_start\ndata: {\"type\": \"content_block_start\", \"content_block\": {\"type\": \"text\", \"text\": \"\"}}\n\n",
            b"event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"delta\": {\"type\": \"text_delta\", \"text\": \"Hello\"}}\n\n",
            b"event: content_block_delta\ndata: {\"type\": \"content_block_delta\", \"delta\": {\"type\": \"text_delta\", \"text\": \" world\"}}\n\n",
            b"event: content_block_stop\ndata: {\"type\": \"content_block_stop\"}\n\n",
            b"event: message_delta\ndata: {\"type\": \"message_delta\", \"delta\": {\"stop_reason\": \"end_turn\", \"stop_sequence\": null}}\n\n",
            b"event: message_stop\ndata: {\"type\": \"message_stop\"}\n\n",
        ];

        let byte_stream = futures_util::stream::iter(
            sse_data.into_iter().map(|chunk| Ok(Bytes::from_static(chunk)))
        );

        let mut streaming_response = AnthropicStreamingResponse::new(byte_stream);
        let mut chunks = Vec::new();

        while let Some(chunk) = streaming_response.next().await {
            chunks.push(chunk.unwrap());
        }

        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0].delta, "Hello");
        assert_eq!(chunks[1].delta, " world");
        assert_eq!(chunks[2].delta, ""); // Empty delta with finish reason
        assert_eq!(chunks[2].finish_reason, Some("end_turn".to_string()));
    }

    #[tokio::test]
    async fn test_streaming_response_with_simplified_events() {
        // Test with simplified JSON events (without event: prefix)
        let sse_data = vec![
            b"data: {\"type\": \"message_start\", \"message\": {\"id\": \"msg_123\", \"type\": \"message\", \"role\": \"assistant\", \"content\": [], \"model\": \"claude-3-sonnet\", \"usage\": {\"input_tokens\": 10, \"output_tokens\": 0}}}\n\n",
            b"data: {\"type\": \"content_block_delta\", \"delta\": {\"type\": \"text_delta\", \"text\": \"Hello\"}}\n\n",
            b"data: {\"type\": \"content_block_delta\", \"delta\": {\"type\": \"text_delta\", \"text\": \" world\"}}\n\n",
            b"data: {\"type\": \"message_delta\", \"delta\": {\"stop_reason\": \"end_turn\"}}\n\n",
            b"data: {\"type\": \"message_stop\"}\n\n",
        ];

        let byte_stream = futures_util::stream::iter(
            sse_data.into_iter().map(|chunk| Ok(Bytes::from_static(chunk)))
        );

        let mut streaming_response = AnthropicStreamingResponse::new(byte_stream);
        let mut chunks = Vec::new();

        while let Some(chunk) = streaming_response.next().await {
            chunks.push(chunk.unwrap());
        }

        assert_eq!(chunks.len(), 3);
        assert_eq!(chunks[0].delta, "Hello");
        assert_eq!(chunks[1].delta, " world");
        assert_eq!(chunks[2].finish_reason, Some("end_turn".to_string()));
    }

    #[test]
    fn test_anthropic_stream_event_deserialization() {
        let json = r#"{
            "type": "content_block_delta",
            "delta": {
                "type": "text_delta",
                "text": "Hello"
            }
        }"#;

        let event: AnthropicStreamEvent = serde_json::from_str(json).unwrap();
        match event {
            AnthropicStreamEvent::ContentBlockDelta { delta } => {
                assert_eq!(delta.text, Some("Hello".to_string()));
            }
            _ => panic!("Expected ContentBlockDelta event"),
        }
    }

    #[test]
    fn test_anthropic_message_start_event() {
        let json = r#"{
            "type": "message_start",
            "message": {
                "id": "msg_123",
                "type": "message",
                "role": "assistant",
                "content": [],
                "model": "claude-3-sonnet",
                "usage": {
                    "input_tokens": 10,
                    "output_tokens": 0
                }
            }
        }"#;

        let event: AnthropicStreamEvent = serde_json::from_str(json).unwrap();
        match event {
            AnthropicStreamEvent::MessageStart { message } => {
                assert_eq!(message.id, "msg_123");
                assert_eq!(message.model, "claude-3-sonnet");
                assert_eq!(message.usage.input_tokens, 10);
            }
            _ => panic!("Expected MessageStart event"),
        }
    }

    #[test]
    fn test_anthropic_message_delta_event() {
        let json = r#"{
            "type": "message_delta",
            "delta": {
                "stop_reason": "end_turn",
                "stop_sequence": null
            }
        }"#;

        let event: AnthropicStreamEvent = serde_json::from_str(json).unwrap();
        match event {
            AnthropicStreamEvent::MessageDelta { delta } => {
                assert_eq!(delta.stop_reason, Some("end_turn".to_string()));
                assert_eq!(delta.stop_sequence, None);
            }
            _ => panic!("Expected MessageDelta event"),
        }
    }
}
