//! Streaming deduplication for assistant messages.

use crate::types::{RawAssistantMessage, RawMessage};

/// Deduplicate streamed assistant messages.
///
/// Claude Code writes multiple JSONL lines for the same assistant message
/// as it streams. Each shares the same `uuid` with progressively more
/// content blocks. We keep only the last entry per uuid.
///
/// Non-assistant messages pass through unchanged.
pub fn dedup_messages(messages: Vec<RawMessage>) -> Vec<RawMessage> {
    let mut result: Vec<RawMessage> = Vec::new();
    let mut buffered_assistant: Option<RawAssistantMessage> = None;

    for msg in messages {
        match msg {
            RawMessage::Assistant(ref asst) => {
                let current_uuid = asst.uuid.as_deref();

                if let Some(ref buffered) = buffered_assistant {
                    let buffered_uuid = buffered.uuid.as_deref();
                    if current_uuid == buffered_uuid {
                        // Same uuid — replace buffer with newer (more complete) version
                        buffered_assistant = Some(asst.clone());
                    } else {
                        // Different uuid — flush old buffer, start new
                        result.push(RawMessage::Assistant(buffered.clone()));
                        buffered_assistant = Some(asst.clone());
                    }
                } else {
                    // No buffer yet — start buffering
                    buffered_assistant = Some(asst.clone());
                }
            }
            _ => {
                // Non-assistant: flush any buffered assistant first, then push this
                if let Some(buffered) = buffered_assistant.take() {
                    result.push(RawMessage::Assistant(buffered));
                }
                result.push(msg);
            }
        }
    }

    // Flush remaining buffer
    if let Some(buffered) = buffered_assistant {
        result.push(RawMessage::Assistant(buffered));
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{AssistantInner, ContentBlock};

    fn make_assistant(uuid: &str, stop_reason: Option<&str>, text: &str) -> RawMessage {
        RawMessage::Assistant(RawAssistantMessage {
            uuid: Some(uuid.to_string()),
            parent_uuid: None,
            timestamp: None,
            session_id: None,
            cwd: None,
            version: None,
            git_branch: None,
            is_sidechain: false,
            request_id: None,
            message: AssistantInner {
                model: None,
                id: None,
                message_type: None,
                role: None,
                content: vec![ContentBlock::Text { text: text.to_string() }],
                stop_reason: stop_reason.map(str::to_string),
                stop_sequence: None,
                usage: None,
            },
        })
    }

    #[test]
    fn dedup_single_streamed_message() {
        let msgs = vec![
            make_assistant("a-1", None, "Part 1"),
            make_assistant("a-1", None, "Part 1 more"),
            make_assistant("a-1", Some("end_turn"), "Part 1 final"),
        ];
        let deduped = dedup_messages(msgs);
        assert_eq!(deduped.len(), 1);
        if let RawMessage::Assistant(a) = &deduped[0] {
            assert_eq!(a.message.stop_reason.as_deref(), Some("end_turn"));
            match &a.message.content[0] {
                ContentBlock::Text { text } => assert_eq!(text, "Part 1 final"),
                _ => panic!("Expected text block"),
            }
        }
    }

    #[test]
    fn dedup_two_distinct_assistants() {
        let msgs = vec![
            make_assistant("a-1", Some("end_turn"), "First"),
            make_assistant("a-2", Some("end_turn"), "Second"),
        ];
        let deduped = dedup_messages(msgs);
        assert_eq!(deduped.len(), 2);
    }

    #[test]
    fn dedup_empty_input() {
        let deduped = dedup_messages(vec![]);
        assert!(deduped.is_empty());
    }
}
