//! Message accumulator for incremental message assembly.
//!
//! This is a thin wrapper around [`dirigent_protocol::accumulator::MessageAccumulator`]
//! that delegates chunk/tool/thinking operations to the protocol accumulator and
//! converts [`AccumulatedMessage`] to [`MessageRecord`] on `finalize()`.
//!
//! The accumulator preserves the order of content parts (text, thinking, tool calls)
//! as they arrive in the event stream, enabling inline tool rendering in the UI.

use chrono::{DateTime, Utc};
use dirigent_protocol::accumulator::{
    AccumulatedMessage, AccumulatedPart,
    MessageAccumulator as ProtocolAccumulator,
};
#[cfg(test)]
use dirigent_protocol::MessagePart;
use dirigent_protocol::ContentBlock;
use serde_json::Value;
use std::collections::HashMap;
use uuid::Uuid;

use crate::error::Result;
use crate::types::MessageRecord;

// Re-export ToolCallData from the protocol for backward compatibility.
pub use dirigent_protocol::accumulator::ToolCallData;

/// Accumulator for assembling streaming message deltas into [`MessageRecord`]s.
///
/// Wraps the protocol-level [`ProtocolAccumulator`] and adds archivist-specific
/// concerns: per-message metadata, UUID parsing, and markdown generation.
#[derive(Debug, Default)]
pub struct MessageAccumulator {
    inner: ProtocolAccumulator,
    /// Per-message metadata not tracked by the protocol accumulator.
    metadata: HashMap<String, Value>,
}

impl MessageAccumulator {
    /// Create a new message accumulator
    pub fn new() -> Result<Self> {
        Ok(Self {
            inner: ProtocolAccumulator::new(),
            metadata: HashMap::new(),
        })
    }

    /// Add a content chunk to the message buffer
    pub fn add_chunk(
        &mut self,
        message_id: String,
        session_id: String,
        connector_id: String,
        role: String,
        content: ContentBlock,
    ) {
        self.inner
            .add_chunk(&message_id, &session_id, &connector_id, &role, content);
    }

    /// Add thinking content to the message buffer
    pub fn add_thinking(
        &mut self,
        message_id: String,
        session_id: String,
        connector_id: String,
        content: String,
    ) {
        self.inner.add_thinking(&message_id, &session_id, &connector_id, &content);
    }

    /// Add or update a tool call in the message buffer
    ///
    /// This method handles both initial ToolCall events and ToolCallUpdate events.
    /// If a tool call with the given ID already exists, it updates the existing entry.
    /// Otherwise, it adds a new entry.
    ///
    /// This ensures that each tool_call_id appears exactly ONCE in the final message,
    /// with the most recent input/output data.
    pub fn add_or_update_tool_call(&mut self, message_id: String, tool_call: ToolCallData) {
        self.inner.add_or_update_tool_call(&message_id, tool_call);
    }

    /// Add a tool call to the message buffer (DEPRECATED - use add_or_update_tool_call)
    #[deprecated(note = "Use add_or_update_tool_call instead to avoid duplicates")]
    pub fn add_tool_call(&mut self, message_id: String, tool_call: ToolCallData) {
        self.add_or_update_tool_call(message_id, tool_call);
    }

    /// Update an existing tool call in the message buffer
    ///
    /// Finds the tool call by ID and updates its input/output with non-empty values
    /// from the update. If no matching tool call is found, this is a no-op (the
    /// update arrived before the initial ToolCall).
    pub fn update_tool_call(
        &mut self,
        message_id: String,
        tool_call_id: &str,
        input: Option<Value>,
        output: Option<Value>,
    ) {
        // Construct a ToolCallData and delegate to add_or_update_tool_call.
        // We need the tool_name but don't have it here; use an empty string
        // since add_or_update_tool_call only updates existing entries when the
        // id matches. However, if there's no existing entry, this would create
        // a new one with empty tool_name - so we need to check first.
        //
        // Instead, we use the protocol accumulator's update semantics directly:
        // build a ToolCallData with the values we have.
        let tool_call = ToolCallData {
            id: tool_call_id.to_string(),
            tool_name: String::new(), // Will be overwritten by existing entry's name
            input: input.unwrap_or(Value::Null),
            output,
        };

        // Only delegate if a buffer exists for this message (matching original behavior).
        if self.inner.has_buffer(&message_id) {
            self.inner.add_or_update_tool_call(&message_id, tool_call);
        }
    }

    /// Get all message IDs for a given session that have active buffers
    pub fn get_message_ids_for_session(&self, session_id: &str) -> Vec<String> {
        self.inner.message_ids_for_session(session_id)
    }

    /// Get message IDs for buffers that have been inactive longer than the threshold
    pub fn get_stale_message_ids(
        &self,
        _now: DateTime<Utc>,
        threshold: std::time::Duration,
    ) -> Vec<String> {
        self.inner.stale_message_ids(threshold)
    }

    /// Get all message IDs that have active buffers
    pub fn get_all_message_ids(&self) -> Vec<String> {
        self.inner.active_message_ids()
    }

    /// Finalize a message and produce a complete `(MessageRecord, connector_id, native_session_id)`.
    ///
    /// Returns `None` if no buffer exists for the given `message_id`.
    /// The `connector_id` and `native_session_id` in the tuple are the raw values
    /// that were passed into `add_chunk`/`add_thinking` — callers in Task 5 will use
    /// these to resolve the canonical scroll_id.
    pub fn finalize(&mut self, message_id: &str) -> Option<(MessageRecord, String, String)> {
        let accumulated = self.inner.finalize(message_id)?;

        let connector_id = accumulated.connector_id.clone();
        let native_session_id = accumulated.session_id.clone();

        // Take stored metadata for this message (if any).
        let metadata = self
            .metadata
            .remove(message_id)
            .unwrap_or(Value::Null);

        let record = accumulated_to_record(accumulated, metadata);
        Some((record, connector_id, native_session_id))
    }
}

// ---------------------------------------------------------------------------
// Conversion helpers
// ---------------------------------------------------------------------------

/// Convert an [`AccumulatedMessage`] into a [`MessageRecord`] for archival.
fn accumulated_to_record(accumulated: AccumulatedMessage, metadata: Value) -> MessageRecord {
    // Build content_md by iterating parts in order
    let mut content_md = String::new();

    for part in &accumulated.parts {
        match part {
            AccumulatedPart::Text { text } => {
                content_md.push_str(text);
            }
            AccumulatedPart::Thinking { text } => {
                content_md.push_str("\n\n<thinking>\n");
                content_md.push_str(text);
                content_md.push_str("\n</thinking>");
            }
            AccumulatedPart::Tool { data } => {
                content_md.push_str(&format!(
                    "\n\n**Tool**: {}\n```json\n{}\n```",
                    data.tool_name,
                    serde_json::to_string_pretty(&data.input)
                        .unwrap_or_else(|_| "{}".to_string())
                ));
            }
        }
    }

    // Convert accumulated parts to protocol MessageParts for rich rendering
    let message_parts = accumulated.to_message_parts();

    // Serialize content_parts for storage (None if empty to save space)
    let content_parts = if message_parts.is_empty() {
        None
    } else {
        serde_json::to_value(&message_parts).ok()
    };

    // Parse UUIDs from strings
    // Strip "msg-" prefix if present (ACP connectors use this format)
    let message_id_str = accumulated
        .message_id
        .strip_prefix("msg-")
        .unwrap_or(&accumulated.message_id);

    if message_id_str != accumulated.message_id.as_str() {
        tracing::debug!(
            "Stripped 'msg-' prefix from message_id: {} -> {}",
            accumulated.message_id,
            message_id_str
        );
    }

    let message_uuid = match Uuid::parse_str(message_id_str) {
        Ok(uuid) => uuid,
        Err(_) => {
            tracing::warn!(
                "Failed to parse message_id as UUID: {}",
                accumulated.message_id
            );
            Uuid::now_v7()
        }
    };

    // Strip "msg-" prefix from session_id if present (for consistency)
    let session_id_str = accumulated
        .session_id
        .strip_prefix("msg-")
        .unwrap_or(&accumulated.session_id);

    if session_id_str != accumulated.session_id.as_str() {
        tracing::debug!(
            "Stripped 'msg-' prefix from session_id: {} -> {}",
            accumulated.session_id,
            session_id_str
        );
    }

    let session_uuid = match Uuid::parse_str(session_id_str) {
        Ok(uuid) => uuid,
        Err(_) => {
            tracing::warn!(
                "Failed to parse session_id as UUID: {}",
                accumulated.session_id
            );
            Uuid::now_v7()
        }
    };

    MessageRecord {
        version: 1,
        message_id: message_uuid,
        session: session_uuid,
        parent_id: None,
        ts: accumulated.created_at.unwrap_or_else(Utc::now),
        role: accumulated.role,
        author: None,
        content_md,
        content_parts,
        attachments: Vec::new(),
        metadata,
    }
}

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

    #[test]
    fn test_accumulator_creation() {
        let acc = MessageAccumulator::new().unwrap();
        assert_eq!(acc.get_all_message_ids().len(), 0);
    }

    #[test]
    fn test_add_text_chunk() {
        let mut acc = MessageAccumulator::new().unwrap();

        acc.add_chunk(
            "msg_1".to_string(),
            "session_1".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Hello, ".to_string(),
            },
        );

        acc.add_chunk(
            "msg_1".to_string(),
            "session_1".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "world!".to_string(),
            },
        );

        // Two consecutive text chunks should be coalesced.
        // Finalize and check the result.
        let (record, _, _) = acc.finalize("msg_1").unwrap();
        assert_eq!(record.content_md, "Hello, world!");
    }

    #[test]
    fn test_add_thinking_chunk() {
        let mut acc = MessageAccumulator::new().unwrap();

        acc.add_thinking(
            "msg_2".to_string(),
            "session_2".to_string(),
            "connector_1".to_string(),
            "Let me think... ".to_string(),
        );

        acc.add_thinking(
            "msg_2".to_string(),
            "session_2".to_string(),
            "connector_1".to_string(),
            "I need to analyze this.".to_string(),
        );

        // Finalize and verify thinking was coalesced
        let (record, _, _) = acc.finalize("msg_2").unwrap();
        assert!(record.content_md.contains("Let me think... I need to analyze this."));
        assert_eq!(record.role, "assistant");
    }

    #[test]
    fn test_finalize_text_only() {
        let mut acc = MessageAccumulator::new().unwrap();

        acc.add_chunk(
            "01936e8f-e5a7-7000-8000-000000000001".to_string(),
            "01936e8f-e5a7-7000-8000-000000000002".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Hello, ".to_string(),
            },
        );

        acc.add_chunk(
            "01936e8f-e5a7-7000-8000-000000000001".to_string(),
            "01936e8f-e5a7-7000-8000-000000000002".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "world!".to_string(),
            },
        );

        let (record, _, _) = acc
            .finalize("01936e8f-e5a7-7000-8000-000000000001")
            .unwrap();

        assert_eq!(record.content_md, "Hello, world!");
        assert_eq!(record.role, "user");
        assert!(record.ts <= Utc::now());
    }

    #[test]
    fn test_finalize_with_thinking() {
        let mut acc = MessageAccumulator::new().unwrap();

        acc.add_chunk(
            "01936e8f-e5a7-7000-8000-000000000003".to_string(),
            "01936e8f-e5a7-7000-8000-000000000004".to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Here's my response.".to_string(),
            },
        );

        acc.add_thinking(
            "01936e8f-e5a7-7000-8000-000000000003".to_string(),
            "01936e8f-e5a7-7000-8000-000000000004".to_string(),
            "connector_1".to_string(),
            "Let me analyze this carefully.".to_string(),
        );

        let (record, _, _) = acc
            .finalize("01936e8f-e5a7-7000-8000-000000000003")
            .unwrap();

        assert!(record.content_md.contains("Here's my response."));
        assert!(record.content_md.contains("<thinking>"));
        assert!(record.content_md.contains("Let me analyze this carefully."));
        assert!(record.content_md.contains("</thinking>"));
    }

    #[test]
    fn test_finalize_nonexistent_message() {
        let mut acc = MessageAccumulator::new().unwrap();
        let result = acc.finalize("nonexistent");
        assert!(result.is_none());
    }

    #[test]
    fn test_add_tool_call() {
        let mut acc = MessageAccumulator::new().unwrap();

        // First add a text chunk to create the buffer
        acc.add_chunk(
            "msg_tool".to_string(),
            "session_tool".to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "I'll use a tool.".to_string(),
            },
        );

        // Add a tool call
        let tool_call = ToolCallData {
            id: "call_123".to_string(),
            tool_name: "search".to_string(),
            input: serde_json::json!({"query": "test"}),
            output: Some(serde_json::json!({"results": ["a", "b"]})),
        };

        #[allow(deprecated)]
        acc.add_tool_call("msg_tool".to_string(), tool_call);

        // Finalize and verify
        let (record, _, _) = acc.finalize("msg_tool").unwrap();
        let parts =
            serde_json::from_value::<Vec<MessagePart>>(record.content_parts.unwrap()).unwrap();
        assert_eq!(parts.len(), 2); // One Text, one Tool
        assert!(matches!(parts[1], MessagePart::Tool { .. }));
        if let MessagePart::Tool { tool, .. } = &parts[1] {
            assert_eq!(tool, "search");
        }
    }

    #[test]
    fn test_finalize_with_tool_calls() {
        let mut acc = MessageAccumulator::new().unwrap();

        acc.add_chunk(
            "01936e8f-e5a7-7000-8000-000000000005".to_string(),
            "01936e8f-e5a7-7000-8000-000000000006".to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Let me search for that.".to_string(),
            },
        );

        let tool_call = ToolCallData {
            id: "call_456".to_string(),
            tool_name: "web_search".to_string(),
            input: serde_json::json!({"query": "Rust async"}),
            output: None,
        };

        #[allow(deprecated)]
        acc.add_tool_call(
            "01936e8f-e5a7-7000-8000-000000000005".to_string(),
            tool_call,
        );

        let (record, _, _) = acc
            .finalize("01936e8f-e5a7-7000-8000-000000000005")
            .unwrap();

        assert!(record.content_md.contains("Let me search for that."));
        assert!(record.content_md.contains("**Tool**: web_search"));
        assert!(record.content_md.contains("Rust async"));
    }

    #[test]
    fn test_concurrent_messages() {
        let mut acc = MessageAccumulator::new().unwrap();

        // Add chunks for two different messages
        acc.add_chunk(
            "msg_a".to_string(),
            "session_1".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Message A".to_string(),
            },
        );

        acc.add_chunk(
            "msg_b".to_string(),
            "session_1".to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Message B".to_string(),
            },
        );

        acc.add_chunk(
            "msg_a".to_string(),
            "session_1".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: " continued".to_string(),
            },
        );

        // Both messages should be buffered
        assert_eq!(acc.get_all_message_ids().len(), 2);

        // Finalize and check
        let (record_a, _, _) = acc.finalize("msg_a").unwrap();
        assert_eq!(record_a.content_md, "Message A continued");

        let (record_b, _, _) = acc.finalize("msg_b").unwrap();
        assert_eq!(record_b.content_md, "Message B");
    }

    #[test]
    fn test_get_message_ids_for_session() {
        let mut acc = MessageAccumulator::new().unwrap();

        // Add messages to different sessions
        acc.add_chunk(
            "msg_1".to_string(),
            "session_a".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Message 1".to_string(),
            },
        );

        acc.add_chunk(
            "msg_2".to_string(),
            "session_a".to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Message 2".to_string(),
            },
        );

        acc.add_chunk(
            "msg_3".to_string(),
            "session_b".to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Message 3".to_string(),
            },
        );

        // Get message IDs for session_a
        let mut session_a_ids = acc.get_message_ids_for_session("session_a");
        session_a_ids.sort();
        assert_eq!(session_a_ids, vec!["msg_1", "msg_2"]);

        // Get message IDs for session_b
        let session_b_ids = acc.get_message_ids_for_session("session_b");
        assert_eq!(session_b_ids, vec!["msg_3"]);

        // Get message IDs for non-existent session
        let empty_ids = acc.get_message_ids_for_session("session_c");
        assert!(empty_ids.is_empty());
    }

    #[test]
    fn test_finalize_with_msg_prefix() {
        let mut acc = MessageAccumulator::new().unwrap();

        // Use message_id and session_id with "msg-" prefix (ACP format)
        let uuid_str = "01936e8f-e5a7-7000-8000-000000000007";
        let session_uuid_str = "01936e8f-e5a7-7000-8000-000000000008";

        acc.add_chunk(
            format!("msg-{}", uuid_str),
            format!("msg-{}", session_uuid_str),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Testing msg- prefix handling.".to_string(),
            },
        );

        let (record, _, _) = acc.finalize(&format!("msg-{}", uuid_str)).unwrap();

        // Verify that the UUID was correctly parsed (not regenerated)
        assert_eq!(record.message_id.to_string(), uuid_str);
        assert_eq!(record.session.to_string(), session_uuid_str);
        assert_eq!(record.content_md, "Testing msg- prefix handling.");
    }

    #[test]
    fn test_finalize_without_msg_prefix() {
        let mut acc = MessageAccumulator::new().unwrap();

        // Use message_id and session_id without "msg-" prefix
        let uuid_str = "01936e8f-e5a7-7000-8000-000000000009";
        let session_uuid_str = "01936e8f-e5a7-7000-8000-00000000000a";

        acc.add_chunk(
            uuid_str.to_string(),
            session_uuid_str.to_string(),
            "connector_1".to_string(),
            "user".to_string(),
            ContentBlock::Text {
                text: "Testing without prefix.".to_string(),
            },
        );

        let (record, _, _) = acc.finalize(uuid_str).unwrap();

        // Verify that the UUID was correctly parsed
        assert_eq!(record.message_id.to_string(), uuid_str);
        assert_eq!(record.session.to_string(), session_uuid_str);
        assert_eq!(record.content_md, "Testing without prefix.");
    }

    #[test]
    fn test_interleaved_tool_calls() {
        let mut acc = MessageAccumulator::new().unwrap();
        let msg_id = "01936e8f-e5a7-7000-8000-000000000010";
        let session_id = "01936e8f-e5a7-7000-8000-000000000011";

        // Text chunk 1
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Let me search for that. ".to_string(),
            },
        );

        // Tool call 1
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_1".to_string(),
                tool_name: "search".to_string(),
                input: serde_json::json!({"query": "rust"}),
                output: None,
            },
        );

        // Text chunk 2
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Now let me check the documentation. ".to_string(),
            },
        );

        // Tool call 2
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_2".to_string(),
                tool_name: "read_docs".to_string(),
                input: serde_json::json!({"path": "README.md"}),
                output: None,
            },
        );

        // Text chunk 3
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Based on my research...".to_string(),
            },
        );

        let (record, _, _) = acc.finalize(msg_id).unwrap();

        // Verify content_md has correct order (text1, tool1, text2, tool2, text3)
        let content = &record.content_md;
        let search_pos = content.find("**Tool**: search").expect("search tool not found");
        let docs_pos = content
            .find("**Tool**: read_docs")
            .expect("read_docs tool not found");
        let text1_pos = content.find("Let me search").expect("text1 not found");
        let text2_pos = content.find("Now let me check").expect("text2 not found");
        let text3_pos = content.find("Based on my research").expect("text3 not found");

        // Verify order: text1 < search < text2 < read_docs < text3
        assert!(
            text1_pos < search_pos,
            "text1 should come before search tool"
        );
        assert!(
            search_pos < text2_pos,
            "search tool should come before text2"
        );
        assert!(
            text2_pos < docs_pos,
            "text2 should come before read_docs tool"
        );
        assert!(
            docs_pos < text3_pos,
            "read_docs tool should come before text3"
        );

        // Verify content_parts structure
        let parts =
            serde_json::from_value::<Vec<MessagePart>>(record.content_parts.unwrap()).unwrap();
        assert_eq!(
            parts.len(),
            5,
            "Should have 5 parts: text, tool, text, tool, text"
        );

        // Verify each part type in order
        assert!(matches!(parts[0], MessagePart::Text { .. }));
        assert!(matches!(parts[1], MessagePart::Tool { .. }));
        assert!(matches!(parts[2], MessagePart::Text { .. }));
        assert!(matches!(parts[3], MessagePart::Tool { .. }));
        assert!(matches!(parts[4], MessagePart::Text { .. }));
    }

    #[test]
    fn test_text_coalescing_with_tool_separation() {
        let mut acc = MessageAccumulator::new().unwrap();
        let msg_id = "msg1";
        let session_id = "session1";

        // Two consecutive text chunks (should coalesce)
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Hello ".to_string(),
            },
        );
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "world. ".to_string(),
            },
        );

        // Tool call (separates text)
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_1".to_string(),
                tool_name: "search".to_string(),
                input: serde_json::json!({}),
                output: None,
            },
        );

        // Two more consecutive text chunks (should coalesce separately)
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "More ".to_string(),
            },
        );
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "text.".to_string(),
            },
        );

        let (record, _, _) = acc.finalize(msg_id).unwrap();

        // Should have 3 parts: coalesced text1, tool, coalesced text2
        let parts =
            serde_json::from_value::<Vec<MessagePart>>(record.content_parts.unwrap()).unwrap();
        assert_eq!(parts.len(), 3);

        // Verify first text part is coalesced
        if let MessagePart::Text { text } = &parts[0] {
            assert_eq!(text, "Hello world. ");
        } else {
            panic!("Expected Text part");
        }

        // Verify tool part
        assert!(matches!(parts[1], MessagePart::Tool { .. }));

        // Verify second text part is coalesced
        if let MessagePart::Text { text } = &parts[2] {
            assert_eq!(text, "More text.");
        } else {
            panic!("Expected Text part");
        }
    }

    #[test]
    fn test_tool_call_progressive_updates() {
        let mut acc = MessageAccumulator::new().unwrap();
        let msg_id = "msg1";
        let session_id = "session1";

        // Create buffer with initial text chunk
        acc.add_chunk(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "assistant".to_string(),
            ContentBlock::Text {
                text: "Using grep... ".to_string(),
            },
        );

        // Initial tool call (empty input, no output)
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_1".to_string(),
                tool_name: "grep".to_string(),
                input: serde_json::json!({}),
                output: None,
            },
        );

        // Update with actual input
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_1".to_string(),
                tool_name: "grep".to_string(),
                input: serde_json::json!({"pattern": "rust"}),
                output: None,
            },
        );

        // Update with output
        acc.add_or_update_tool_call(
            msg_id.to_string(),
            ToolCallData {
                id: "call_1".to_string(),
                tool_name: "grep".to_string(),
                input: serde_json::json!({}), // Empty, should not overwrite
                output: Some(serde_json::json!({"results": ["match1", "match2"]})),
            },
        );

        let (record, _, _) = acc.finalize(msg_id).unwrap();

        // Should have 2 parts: text and tool (tool merged from 3 updates)
        let parts =
            serde_json::from_value::<Vec<MessagePart>>(record.content_parts.unwrap()).unwrap();
        assert_eq!(parts.len(), 2);

        // Verify first part is text
        assert!(matches!(parts[0], MessagePart::Text { .. }));

        // Verify second part is tool with merged data
        if let MessagePart::Tool {
            tool, input, output, ..
        } = &parts[1]
        {
            assert_eq!(tool, "grep");
            assert_eq!(input, &serde_json::json!({"pattern": "rust"})); // Input preserved
            assert!(output.is_some()); // Output added
        } else {
            panic!("Expected Tool part");
        }
    }

    #[test]
    fn test_thinking_coalescing() {
        let mut acc = MessageAccumulator::new().unwrap();
        let msg_id = "msg1";
        let session_id = "session1";

        // Add multiple thinking chunks
        acc.add_thinking(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "First thought. ".to_string(),
        );
        acc.add_thinking(
            msg_id.to_string(),
            session_id.to_string(),
            "connector_1".to_string(),
            "Second thought.".to_string(),
        );

        let (record, _, _) = acc.finalize(msg_id).unwrap();

        // Should have 1 thinking part (coalesced)
        let parts =
            serde_json::from_value::<Vec<MessagePart>>(record.content_parts.unwrap()).unwrap();
        assert_eq!(parts.len(), 1);

        // Verify it's coalesced thinking
        if let MessagePart::Thinking { text } = &parts[0] {
            assert_eq!(text, "First thought. Second thought.");
        } else {
            panic!("Expected Thinking part");
        }
    }
}
