package context_test

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/yaoapp/yao/agent/context"
)

// =============================================================================
// ChatBuffer Creation Tests
// =============================================================================

func TestBufferNewChatBuffer(t *testing.T) {
	t.Run("CreateWithAllFields", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-123", "req-456", "assistant-789", "", "")

		assert.NotNil(t, buffer)
		assert.Equal(t, "chat-123", buffer.ChatID())
		assert.Equal(t, "req-456", buffer.RequestID())
		assert.Equal(t, "assistant-789", buffer.AssistantID())
		assert.Empty(t, buffer.GetMessages())
		assert.Empty(t, buffer.GetAllSteps())
		assert.Equal(t, 0, buffer.GetMessageCount())
	})

	t.Run("CreateWithEmptyFields", func(t *testing.T) {
		buffer := context.NewChatBuffer("", "", "", "", "")

		assert.NotNil(t, buffer)
		assert.Empty(t, buffer.ChatID())
		assert.Empty(t, buffer.RequestID())
		assert.Empty(t, buffer.AssistantID())
	})
}

// =============================================================================
// Message Buffer Tests
// =============================================================================

func TestBufferAddMessage(t *testing.T) {
	t.Run("AddSingleMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		msg := &context.BufferedMessage{
			Role:  "assistant",
			Type:  "text",
			Props: map[string]interface{}{"content": "Hello"},
		}
		buffer.AddMessage(msg)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "assistant", messages[0].Role)
		assert.Equal(t, "text", messages[0].Type)
		assert.Equal(t, 1, messages[0].Sequence)
		assert.NotEmpty(t, messages[0].MessageID) // Auto-generated
		assert.Equal(t, "chat-1", messages[0].ChatID)
		assert.Equal(t, "req-1", messages[0].RequestID)
		assert.False(t, messages[0].CreatedAt.IsZero())
	})

	t.Run("AddMultipleMessages", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		for i := 0; i < 5; i++ {
			buffer.AddMessage(&context.BufferedMessage{
				Role:  "assistant",
				Type:  "text",
				Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i+1)},
			})
		}

		messages := buffer.GetMessages()
		require.Len(t, messages, 5)

		// Verify sequence numbers
		for i, msg := range messages {
			assert.Equal(t, i+1, msg.Sequence)
		}
	})

	t.Run("AddNilMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
		buffer.AddMessage(nil)

		assert.Equal(t, 0, buffer.GetMessageCount())
	})

	t.Run("AddMessageWithExistingID", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")

		msg := &context.BufferedMessage{
			MessageID: "custom-id-123",
			Role:      "assistant",
			Type:      "text",
		}
		buffer.AddMessage(msg)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "custom-id-123", messages[0].MessageID) // Preserved
	})

	t.Run("AddMessageWithExistingTimestamp", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")

		customTime := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
		msg := &context.BufferedMessage{
			Role:      "assistant",
			Type:      "text",
			CreatedAt: customTime,
		}
		buffer.AddMessage(msg)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, customTime, messages[0].CreatedAt) // Preserved
	})
}

func TestBufferAddUserInput(t *testing.T) {
	t.Run("AddStringContent", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		buffer.AddUserInput("What is the weather?", "")

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "user", messages[0].Role)
		assert.Equal(t, "user_input", messages[0].Type)
		assert.Equal(t, "What is the weather?", messages[0].Props["content"])
		assert.Equal(t, "user", messages[0].Props["role"])
	})

	t.Run("AddUserInputWithName", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
		buffer.AddUserInput("Hello", "John")

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "John", messages[0].Props["name"])
	})

	t.Run("AddComplexContent", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
		complexContent := []map[string]interface{}{
			{"type": "text", "text": "Look at this image"},
			{"type": "image_url", "image_url": map[string]string{"url": "https://example.com/image.jpg"}},
		}
		buffer.AddUserInput(complexContent, "")

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		content, ok := messages[0].Props["content"].([]map[string]interface{})
		require.True(t, ok)
		assert.Len(t, content, 2)
	})
}

func TestBufferAddAssistantMessage(t *testing.T) {
	t.Run("AddTextMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		buffer.AddAssistantMessage(
			"M1",
			"text",
			map[string]interface{}{"content": "Hello, how can I help?"},
			"block-1",
			"thread-1",
			"assistant-1",
			map[string]interface{}{"model": "gpt-4"},
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "M1", messages[0].MessageID)
		assert.Equal(t, "assistant", messages[0].Role)
		assert.Equal(t, "text", messages[0].Type)
		assert.Equal(t, "block-1", messages[0].BlockID)
		assert.Equal(t, "thread-1", messages[0].ThreadID)
		assert.Equal(t, "assistant-1", messages[0].AssistantID)
		assert.Equal(t, "gpt-4", messages[0].Metadata["model"])
	})

	t.Run("SkipEventMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
		buffer.AddAssistantMessage(
			"E1",
			"event",
			map[string]interface{}{"event": "message_start"},
			"", "", "", nil,
		)

		// Event messages should be skipped
		assert.Equal(t, 0, buffer.GetMessageCount())
	})

	t.Run("AddRetrievalMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
		buffer.AddAssistantMessage(
			"M2",
			"retrieval",
			map[string]interface{}{
				"sources": []map[string]interface{}{
					{"title": "Doc 1", "score": 0.95},
					{"title": "Doc 2", "score": 0.87},
				},
			},
			"block-1", "", "assistant-3", nil,
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "retrieval", messages[0].Type)
	})

	t.Run("AddToolCallMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
		buffer.AddAssistantMessage(
			"M3",
			"tool_call",
			map[string]interface{}{
				"name":      "get_weather",
				"arguments": `{"location": "San Francisco"}`,
			},
			"block-1", "", "assistant-4", nil,
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "tool_call", messages[0].Type)
		assert.Equal(t, "get_weather", messages[0].Props["name"])
	})

	t.Run("AddCustomTypeMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")
		buffer.AddAssistantMessage(
			"M4",
			"custom_chart",
			map[string]interface{}{
				"chart_type": "bar",
				"data":       []int{1, 2, 3, 4, 5},
			},
			"block-1", "", "assistant-5", nil,
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "custom_chart", messages[0].Type)
	})
}

func TestBufferGetMessages(t *testing.T) {
	t.Run("GetMessagesReturnsSliceCopy", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		buffer.AddUserInput("Hello", "")

		messages1 := buffer.GetMessages()
		messages2 := buffer.GetMessages()

		// Slices should be different (copy of slice)
		// But pointers point to same underlying objects (shallow copy)
		assert.Len(t, messages1, 1)
		assert.Len(t, messages2, 1)
	})

	t.Run("GetEmptyMessages", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
		messages := buffer.GetMessages()

		assert.NotNil(t, messages)
		assert.Empty(t, messages)
	})
}

func TestBufferGetMessageCount(t *testing.T) {
	buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
	assert.Equal(t, 0, buffer.GetMessageCount())

	buffer.AddUserInput("Message 1", "")
	assert.Equal(t, 1, buffer.GetMessageCount())

	buffer.AddAssistantMessage("M1", "text", map[string]interface{}{"content": "Reply"}, "", "", "", nil)
	assert.Equal(t, 2, buffer.GetMessageCount())
}

// =============================================================================
// Step Buffer Tests (for Resume)
// =============================================================================

func TestBufferBeginStep(t *testing.T) {
	t.Run("BeginStepWithStack", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		stack := &context.Stack{
			ID:       "stack-123",
			ParentID: "stack-parent-456",
			Depth:    2,
		}

		step := buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"prompt": "Hello"}, stack)

		require.NotNil(t, step)
		assert.NotEmpty(t, step.ResumeID)
		assert.Equal(t, "chat-1", step.ChatID)
		assert.Equal(t, "req-1", step.RequestID)
		assert.Equal(t, "assistant-1", step.AssistantID)
		assert.Equal(t, "stack-123", step.StackID)
		assert.Equal(t, "stack-parent-456", step.StackParentID)
		assert.Equal(t, 2, step.StackDepth)
		assert.Equal(t, context.StepTypeLLM, step.Type)
		assert.Equal(t, context.StepStatusRunning, step.Status)
		assert.Equal(t, 1, step.Sequence)
		assert.Equal(t, "Hello", step.Input["prompt"])
		assert.False(t, step.CreatedAt.IsZero())
	})

	t.Run("BeginStepWithNilStack", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		step := buffer.BeginStep(context.StepTypeInput, nil, nil)

		require.NotNil(t, step)
		assert.Empty(t, step.StackID)
		assert.Empty(t, step.StackParentID)
		assert.Equal(t, 0, step.StackDepth)
	})

	t.Run("BeginMultipleSteps", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")

		step1 := buffer.BeginStep(context.StepTypeInput, nil, nil)
		step2 := buffer.BeginStep(context.StepTypeHookCreate, nil, nil)
		step3 := buffer.BeginStep(context.StepTypeLLM, nil, nil)

		assert.Equal(t, 1, step1.Sequence)
		assert.Equal(t, 2, step2.Sequence)
		assert.Equal(t, 3, step3.Sequence)

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 3)
	})

	t.Run("BeginStepWithSpaceSnapshot", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")

		// Set space snapshot before beginning step
		buffer.SetSpaceSnapshot(map[string]interface{}{
			"key1": "value1",
			"key2": 42,
		})

		step := buffer.BeginStep(context.StepTypeLLM, nil, nil)

		require.NotNil(t, step.SpaceSnapshot)
		assert.Equal(t, "value1", step.SpaceSnapshot["key1"])
		assert.Equal(t, 42, step.SpaceSnapshot["key2"])
	})
}

func TestBufferCompleteStep(t *testing.T) {
	t.Run("CompleteCurrentStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"prompt": "Hello"}, nil)
		buffer.CompleteStep(map[string]interface{}{"response": "Hi there!"})

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 1)
		assert.Equal(t, context.StepStatusCompleted, steps[0].Status)
		assert.Equal(t, "Hi there!", steps[0].Output["response"])
		assert.Nil(t, buffer.GetCurrentStep()) // Current step cleared
	})

	t.Run("CompleteWithNoCurrentStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		// Should not panic
		buffer.CompleteStep(map[string]interface{}{"response": "test"})
		assert.Nil(t, buffer.GetCurrentStep())
	})

	t.Run("CompleteMultipleStepsSequentially", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")

		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(map[string]interface{}{"done": true})

		buffer.BeginStep(context.StepTypeHookCreate, nil, nil)
		buffer.CompleteStep(map[string]interface{}{"hook_result": "ok"})

		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.CompleteStep(map[string]interface{}{"llm_response": "hello"})

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 3)
		for _, step := range steps {
			assert.Equal(t, context.StepStatusCompleted, step.Status)
		}
	})
}

func TestBufferFailCurrentStep(t *testing.T) {
	t.Run("FailWithError", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("API error: rate limit exceeded"))

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 1)
		assert.Equal(t, context.ResumeStatusFailed, steps[0].Status)
		assert.Equal(t, "API error: rate limit exceeded", steps[0].Error)
	})

	t.Run("FailWithInterrupted", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.FailCurrentStep(context.ResumeStatusInterrupted, nil)

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 1)
		assert.Equal(t, context.ResumeStatusInterrupted, steps[0].Status)
		assert.Empty(t, steps[0].Error)
	})

	t.Run("FailAlreadyCompletedStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")

		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.CompleteStep(map[string]interface{}{"done": true})

		// Try to fail completed step (should be no-op since currentStep is nil)
		buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("late error"))

		steps := buffer.GetAllSteps()
		require.Len(t, steps, 1)
		assert.Equal(t, context.StepStatusCompleted, steps[0].Status) // Still completed
	})

	t.Run("FailWithNoCurrentStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")

		// Should not panic
		buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("error"))
	})
}

func TestBufferGetCurrentStep(t *testing.T) {
	t.Run("NoCurrentStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		assert.Nil(t, buffer.GetCurrentStep())
	})

	t.Run("HasCurrentStep", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
		buffer.BeginStep(context.StepTypeLLM, nil, nil)

		current := buffer.GetCurrentStep()
		require.NotNil(t, current)
		assert.Equal(t, context.StepTypeLLM, current.Type)
	})

	t.Run("CurrentStepClearedAfterComplete", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.CompleteStep(nil)

		assert.Nil(t, buffer.GetCurrentStep())
	})
}

func TestBufferGetStepsForResume(t *testing.T) {
	t.Run("CompletedSuccessfully", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(nil)
		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.CompleteStep(nil)

		// Completed successfully - no steps need to be saved
		steps := buffer.GetStepsForResume(context.StepStatusCompleted)
		assert.Nil(t, steps)
	})

	t.Run("FailedRequest", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(nil)
		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		// Step still running when failure occurs

		steps := buffer.GetStepsForResume(context.ResumeStatusFailed)
		require.NotNil(t, steps)
		assert.Len(t, steps, 1)

		// Only the failed step should be returned
		assert.Equal(t, context.ResumeStatusFailed, steps[0].Status)
	})

	t.Run("InterruptedRequest", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")

		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(nil)
		buffer.BeginStep(context.StepTypeHookCreate, nil, nil)
		buffer.CompleteStep(nil)
		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		// Interrupted during LLM

		steps := buffer.GetStepsForResume(context.ResumeStatusInterrupted)
		require.NotNil(t, steps)
		assert.Len(t, steps, 1)
		assert.Equal(t, context.ResumeStatusInterrupted, steps[0].Status)
	})
}

func TestBufferGetAllSteps(t *testing.T) {
	t.Run("GetStepsReturnsSliceCopy", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		buffer.BeginStep(context.StepTypeLLM, nil, nil)

		steps1 := buffer.GetAllSteps()
		steps2 := buffer.GetAllSteps()

		// Slices should be different (copy of slice)
		assert.Len(t, steps1, 1)
		assert.Len(t, steps2, 1)
	})

	t.Run("GetEmptySteps", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
		steps := buffer.GetAllSteps()

		assert.NotNil(t, steps)
		assert.Empty(t, steps)
	})
}

// =============================================================================
// Space Snapshot Tests
// =============================================================================

func TestBufferSpaceSnapshot(t *testing.T) {
	t.Run("SetAndGetSnapshot", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")

		snapshot := map[string]interface{}{
			"user_id":   "user-123",
			"session":   map[string]interface{}{"token": "abc"},
			"counter":   42,
			"is_active": true,
		}
		buffer.SetSpaceSnapshot(snapshot)

		retrieved := buffer.GetSpaceSnapshot()
		assert.Equal(t, "user-123", retrieved["user_id"])
		assert.Equal(t, 42, retrieved["counter"])
		assert.Equal(t, true, retrieved["is_active"])
	})

	t.Run("SnapshotIsCopy", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")

		original := map[string]interface{}{"key": "original"}
		buffer.SetSpaceSnapshot(original)

		// Modify original
		original["key"] = "modified"

		// Buffer should have original value
		retrieved := buffer.GetSpaceSnapshot()
		assert.Equal(t, "original", retrieved["key"])
	})

	t.Run("GetSnapshotReturnsCopy", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
		buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"})

		retrieved1 := buffer.GetSpaceSnapshot()
		retrieved1["key"] = "modified"

		retrieved2 := buffer.GetSpaceSnapshot()
		assert.Equal(t, "value", retrieved2["key"]) // Original unchanged
	})

	t.Run("GetNilSnapshot", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
		snapshot := buffer.GetSpaceSnapshot()
		assert.Nil(t, snapshot)
	})

	t.Run("SetNilSnapshot", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")
		buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"})
		buffer.SetSpaceSnapshot(nil)

		snapshot := buffer.GetSpaceSnapshot()
		assert.Nil(t, snapshot)
	})
}

// =============================================================================
// Identity Methods Tests
// =============================================================================

func TestBufferIdentityMethods(t *testing.T) {
	t.Run("SetAssistantID", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-original", "", "")

		assert.Equal(t, "assistant-original", buffer.AssistantID())

		buffer.SetAssistantID("assistant-new")
		assert.Equal(t, "assistant-new", buffer.AssistantID())
	})

	t.Run("ChatID", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
		assert.Equal(t, "chat-test", buffer.ChatID())
	})

	t.Run("RequestID", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
		assert.Equal(t, "req-test", buffer.RequestID())
	})

	t.Run("Connector", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "openai", "")
		assert.Equal(t, "openai", buffer.Connector())
	})

	t.Run("SetConnector", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
		assert.Equal(t, "openai", buffer.Connector())

		// Simulate user switching connector mid-conversation
		buffer.SetConnector("anthropic")
		assert.Equal(t, "anthropic", buffer.Connector())
	})

	t.Run("EmptyConnector", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
		assert.Equal(t, "", buffer.Connector())
	})
}

func TestBufferConnectorInMessages(t *testing.T) {
	t.Run("MessageInheritsConnector", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add assistant message - should inherit connector from buffer
		buffer.AddAssistantMessage(
			"M1",
			"text",
			map[string]interface{}{"content": "Hello"},
			"block-1", "thread-1", "assistant-1", nil,
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Equal(t, "openai", messages[0].Connector)
	})

	t.Run("MessageConnectorUpdatesWithBuffer", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// First message with openai
		buffer.AddAssistantMessage(
			"M1",
			"text",
			map[string]interface{}{"content": "Using OpenAI"},
			"", "", "assistant-1", nil,
		)

		// User switches connector
		buffer.SetConnector("anthropic")

		// Second message with anthropic
		buffer.AddAssistantMessage(
			"M2",
			"text",
			map[string]interface{}{"content": "Now using Claude"},
			"", "", "assistant-1", nil,
		)

		messages := buffer.GetMessages()
		require.Len(t, messages, 2)
		assert.Equal(t, "openai", messages[0].Connector, "First message should use openai")
		assert.Equal(t, "anthropic", messages[1].Connector, "Second message should use anthropic")
	})

	t.Run("UserInputDoesNotSetConnector", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// User input doesn't have connector (it's set by the system based on which model processes it)
		buffer.AddUserInput("Hello", "")

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		// User input messages don't have connector field set by AddUserInput
		// Connector is only set for assistant messages
		assert.Equal(t, "", messages[0].Connector)
	})

	t.Run("MultipleConnectorSwitches", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Simulate a conversation with multiple connector switches
		connectors := []string{"openai", "anthropic", "openai", "google"}
		for i, conn := range connectors {
			buffer.SetConnector(conn)
			buffer.AddAssistantMessage(
				fmt.Sprintf("M%d", i+1),
				"text",
				map[string]interface{}{"content": fmt.Sprintf("Message %d", i+1)},
				"", "", "assistant-1", nil,
			)
		}

		messages := buffer.GetMessages()
		require.Len(t, messages, 4)

		for i, msg := range messages {
			assert.Equal(t, connectors[i], msg.Connector, "Message %d should have connector %s", i+1, connectors[i])
		}
	})
}

// =============================================================================
// Concurrency Tests
// =============================================================================

func TestBufferConcurrentMessageOperations(t *testing.T) {
	buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "", "")

	var wg sync.WaitGroup
	numGoroutines := 100

	// Concurrent writes
	for i := 0; i < numGoroutines; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			buffer.AddMessage(&context.BufferedMessage{
				Role:  "assistant",
				Type:  "text",
				Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", idx)},
			})
		}(i)
	}

	wg.Wait()

	// Verify all messages were added
	messages := buffer.GetMessages()
	assert.Len(t, messages, numGoroutines)

	// Verify sequences are unique
	sequences := make(map[int]bool)
	for _, msg := range messages {
		assert.False(t, sequences[msg.Sequence], "Duplicate sequence found: %d", msg.Sequence)
		sequences[msg.Sequence] = true
	}
}

func TestBufferConcurrentStepOperations(t *testing.T) {
	buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "", "")

	var wg sync.WaitGroup
	numGoroutines := 50

	// Concurrent step operations
	for i := 0; i < numGoroutines; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"idx": idx}, nil)
			time.Sleep(time.Millisecond) // Simulate some work
			buffer.CompleteStep(map[string]interface{}{"result": idx})
		}(i)
	}

	wg.Wait()

	// Verify all steps were recorded
	steps := buffer.GetAllSteps()
	assert.Len(t, steps, numGoroutines)
}

func TestBufferConcurrentReadWrite(t *testing.T) {
	buffer := context.NewChatBuffer("chat-rw", "req-rw", "assistant-rw", "", "")

	var wg sync.WaitGroup
	done := make(chan bool)

	// Writer goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		for i := 0; i < 100; i++ {
			buffer.AddMessage(&context.BufferedMessage{
				Role:  "assistant",
				Type:  "text",
				Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i)},
			})
			time.Sleep(time.Microsecond)
		}
	}()

	// Reader goroutine
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case <-done:
				return
			default:
				_ = buffer.GetMessages()
				_ = buffer.GetMessageCount()
				time.Sleep(time.Microsecond)
			}
		}
	}()

	// Let it run for a bit
	time.Sleep(50 * time.Millisecond)
	close(done)
	wg.Wait()

	// Should complete without race conditions
	assert.Equal(t, 100, buffer.GetMessageCount())
}

// =============================================================================
// Step Type Constants Tests
// =============================================================================

func TestBufferStepTypeConstants(t *testing.T) {
	// Verify all step types are defined
	assert.Equal(t, "input", context.StepTypeInput)
	assert.Equal(t, "hook_create", context.StepTypeHookCreate)
	assert.Equal(t, "llm", context.StepTypeLLM)
	assert.Equal(t, "tool", context.StepTypeTool)
	assert.Equal(t, "hook_next", context.StepTypeHookNext)
	assert.Equal(t, "delegate", context.StepTypeDelegate)
}

func TestBufferResumeStatusConstants(t *testing.T) {
	assert.Equal(t, "failed", context.ResumeStatusFailed)
	assert.Equal(t, "interrupted", context.ResumeStatusInterrupted)
}

func TestBufferStepStatusConstants(t *testing.T) {
	assert.Equal(t, "running", context.StepStatusRunning)
	assert.Equal(t, "completed", context.StepStatusCompleted)
}

// =============================================================================
// Edge Cases and Error Handling Tests
// =============================================================================

func TestBufferEdgeCases(t *testing.T) {
	t.Run("LargeNumberOfMessages", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-large", "req-large", "assistant-large", "", "")

		// Add 10000 messages
		for i := 0; i < 10000; i++ {
			buffer.AddMessage(&context.BufferedMessage{
				Role:  "assistant",
				Type:  "text",
				Props: map[string]interface{}{"content": fmt.Sprintf("Message %d", i)},
			})
		}

		assert.Equal(t, 10000, buffer.GetMessageCount())
		messages := buffer.GetMessages()
		assert.Len(t, messages, 10000)
	})

	t.Run("MessageWithEmptyProps", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-empty", "req-empty", "assistant-empty", "", "")

		buffer.AddMessage(&context.BufferedMessage{
			Role:  "assistant",
			Type:  "text",
			Props: nil,
		})

		messages := buffer.GetMessages()
		require.Len(t, messages, 1)
		assert.Nil(t, messages[0].Props)
	})

	t.Run("StepWithEmptyInput", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-step", "req-step", "assistant-step", "", "")

		step := buffer.BeginStep(context.StepTypeLLM, nil, nil)
		assert.Nil(t, step.Input)

		buffer.CompleteStep(nil)
		steps := buffer.GetAllSteps()
		assert.Nil(t, steps[0].Output)
	})

	t.Run("AllMessageTypes", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-types", "req-types", "assistant-types", "", "")

		messageTypes := []string{
			"text", "image", "loading", "tool_call", "tool_result",
			"retrieval", "thinking", "action", "chart", "table",
			"custom_type_1", "custom_type_2",
		}

		for i, msgType := range messageTypes {
			buffer.AddAssistantMessage(fmt.Sprintf("M%d", i+1), msgType, map[string]interface{}{"type": msgType}, "", "", "", nil)
		}

		assert.Equal(t, len(messageTypes), buffer.GetMessageCount())
	})

	t.Run("AllStepTypes", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-step-types", "req-step-types", "assistant-step-types", "", "")

		stepTypes := []string{
			context.StepTypeInput, context.StepTypeHookCreate, context.StepTypeLLM,
			context.StepTypeTool, context.StepTypeHookNext, context.StepTypeDelegate,
		}

		for _, stepType := range stepTypes {
			buffer.BeginStep(stepType, nil, nil)
			buffer.CompleteStep(nil)
		}

		steps := buffer.GetAllSteps()
		assert.Len(t, steps, len(stepTypes))
	})
}

// =============================================================================
// Integration-like Tests (Simulating Real Workflow)
// =============================================================================

func TestBufferCompleteWorkflow(t *testing.T) {
	t.Run("SuccessfulChatFlow", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-workflow", "req-workflow", "assistant-main", "", "")

		// 1. User input
		buffer.AddUserInput("What's the weather in San Francisco?", "John")
		buffer.BeginStep(context.StepTypeInput, map[string]interface{}{"content": "What's the weather in San Francisco?"}, nil)
		buffer.CompleteStep(nil)

		// 2. Create hook
		buffer.BeginStep(context.StepTypeHookCreate, nil, nil)
		buffer.AddAssistantMessage("M1", "thinking", map[string]interface{}{"content": "Processing your request..."}, "block-1", "", "assistant-main", nil)
		buffer.CompleteStep(nil)

		// 3. LLM call with tool
		buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"model": "gpt-4"}, nil)
		buffer.AddAssistantMessage("M2", "tool_call", map[string]interface{}{
			"name":      "get_weather",
			"arguments": `{"location":"San Francisco"}`,
		}, "block-2", "", "assistant-main", nil)
		buffer.CompleteStep(map[string]interface{}{"tool_calls": 1})

		// 4. Tool execution
		buffer.BeginStep(context.StepTypeTool, map[string]interface{}{"tool": "get_weather"}, nil)
		buffer.AddAssistantMessage("M3", "tool_result", map[string]interface{}{
			"result": "72°F, Sunny",
		}, "block-2", "", "assistant-main", nil)
		buffer.CompleteStep(map[string]interface{}{"result": "72°F, Sunny"})

		// 5. Final LLM response
		buffer.BeginStep(context.StepTypeLLM, nil, nil)
		buffer.AddAssistantMessage("M4", "text", map[string]interface{}{
			"content": "The weather in San Francisco is currently 72°F and sunny.",
		}, "block-3", "", "assistant-main", nil)
		buffer.CompleteStep(nil)

		// Verify: 1 user_input + 4 assistant messages (thinking, tool_call, tool_result, text)
		assert.Equal(t, 5, buffer.GetMessageCount())
		assert.Len(t, buffer.GetAllSteps(), 5) // 5 steps (no hook_next in this flow)

		// All steps should be completed
		steps := buffer.GetStepsForResume(context.StepStatusCompleted)
		assert.Nil(t, steps)
	})

	t.Run("InterruptedChatFlow", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-interrupted", "req-interrupted", "assistant-main", "", "")

		// Set space snapshot
		buffer.SetSpaceSnapshot(map[string]interface{}{
			"user_context": "previous conversation",
			"preferences":  map[string]interface{}{"language": "en"},
		})

		// 1. User input
		buffer.AddUserInput("Generate a long story", "")
		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(nil)

		// 2. LLM starts generating
		buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"model": "gpt-4"}, nil)
		buffer.AddAssistantMessage("M1", "text", map[string]interface{}{"content": "Once upon a time..."}, "block-1", "", "assistant-main", nil)
		// User interrupts here!

		// Get steps for resume
		steps := buffer.GetStepsForResume(context.ResumeStatusInterrupted)
		require.NotNil(t, steps)
		assert.Len(t, steps, 1)

		// Only the interrupted step should be returned
		lastStep := steps[0]
		assert.Equal(t, context.ResumeStatusInterrupted, lastStep.Status)
		assert.NotNil(t, lastStep.SpaceSnapshot)
		assert.Equal(t, "previous conversation", lastStep.SpaceSnapshot["user_context"])
	})

	t.Run("A2ACallWithDelegation", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-a2a", "req-a2a", "assistant-main", "", "")

		mainStack := &context.Stack{ID: "stack-main", Depth: 0}
		childStack := &context.Stack{ID: "stack-child", ParentID: "stack-main", Depth: 1}

		// Main assistant starts
		buffer.BeginStep(context.StepTypeInput, nil, mainStack)
		buffer.CompleteStep(nil)

		// Delegate to child assistant
		buffer.SetAssistantID("assistant-child")
		buffer.BeginStep(context.StepTypeDelegate, map[string]interface{}{"delegate_to": "assistant-child"}, childStack)

		// Child assistant messages
		buffer.AddAssistantMessage("M1", "text", map[string]interface{}{"content": "Child assistant responding"}, "block-child", "", "assistant-child", nil)
		buffer.CompleteStep(map[string]interface{}{"delegate_result": "success"})

		// Return to main assistant
		buffer.SetAssistantID("assistant-main")
		buffer.BeginStep(context.StepTypeLLM, nil, mainStack)
		buffer.AddAssistantMessage("M2", "text", map[string]interface{}{"content": "Main assistant continuing"}, "block-main", "", "assistant-main", nil)
		buffer.CompleteStep(nil)

		// Verify
		messages := buffer.GetMessages()
		assert.Len(t, messages, 2)
		assert.Equal(t, "assistant-child", messages[0].AssistantID)
		assert.Equal(t, "assistant-main", messages[1].AssistantID)

		steps := buffer.GetAllSteps()
		assert.Len(t, steps, 3)
		assert.Equal(t, "stack-child", steps[1].StackID)
		assert.Equal(t, "stack-main", steps[1].StackParentID)
	})

	t.Run("ConcurrentAgentCalls", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-concurrent-a2a", "req-concurrent-a2a", "assistant-main", "", "")

		// Main assistant spawns multiple concurrent calls
		buffer.BeginStep(context.StepTypeInput, nil, nil)
		buffer.CompleteStep(nil)

		// Simulate concurrent responses with thread IDs
		var wg sync.WaitGroup
		for i := 0; i < 3; i++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				threadID := fmt.Sprintf("thread-%d", idx)
				buffer.AddAssistantMessage(
					fmt.Sprintf("M%d", idx),
					"text",
					map[string]interface{}{"content": fmt.Sprintf("Response from thread %d", idx)},
					"block-concurrent",
					threadID,
					fmt.Sprintf("assistant-%d", idx),
					nil,
				)
			}(i)
		}
		wg.Wait()

		messages := buffer.GetMessages()
		assert.Len(t, messages, 3)

		// Verify all have same block ID but different thread IDs
		threadIDs := make(map[string]bool)
		for _, msg := range messages {
			assert.Equal(t, "block-concurrent", msg.BlockID)
			assert.False(t, threadIDs[msg.ThreadID], "Duplicate thread ID")
			threadIDs[msg.ThreadID] = true
		}
	})
}

// =============================================================================
// Message Sequence Tests
// =============================================================================

func TestBufferMessageSequence(t *testing.T) {
	t.Run("SequenceAutoIncrement", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-seq", "req-seq", "assistant-seq", "", "")

		for i := 0; i < 10; i++ {
			buffer.AddMessage(&context.BufferedMessage{
				Role: "assistant",
				Type: "text",
			})
		}

		messages := buffer.GetMessages()
		for i, msg := range messages {
			assert.Equal(t, i+1, msg.Sequence)
		}
	})

	t.Run("MixedMessageTypes", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-mixed", "req-mixed", "assistant-mixed", "", "")

		buffer.AddUserInput("Hello", "")
		buffer.AddAssistantMessage("M1", "text", nil, "", "", "", nil)
		buffer.AddUserInput("Follow up", "")
		buffer.AddAssistantMessage("M2", "tool_call", nil, "", "", "", nil)

		messages := buffer.GetMessages()
		assert.Len(t, messages, 4)
		for i, msg := range messages {
			assert.Equal(t, i+1, msg.Sequence)
		}
	})
}

// =============================================================================
// Step Sequence Tests
// =============================================================================

func TestBufferStepSequence(t *testing.T) {
	t.Run("SequenceAutoIncrement", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-step-seq", "req-step-seq", "assistant-step-seq", "", "")

		for i := 0; i < 5; i++ {
			buffer.BeginStep(context.StepTypeLLM, nil, nil)
			buffer.CompleteStep(nil)
		}

		steps := buffer.GetAllSteps()
		for i, step := range steps {
			assert.Equal(t, i+1, step.Sequence)
		}
	})
}

// =============================================================================
// Buffer Reset/Clear Tests (if needed in future)
// =============================================================================

func TestBufferMultipleRequests(t *testing.T) {
	t.Run("NewBufferPerRequest", func(t *testing.T) {
		// Simulate multiple requests with separate buffers
		buffer1 := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
		buffer1.AddUserInput("Request 1", "")

		buffer2 := context.NewChatBuffer("chat-1", "req-2", "assistant-1", "", "")
		buffer2.AddUserInput("Request 2", "")

		// Buffers should be independent
		assert.Equal(t, 1, buffer1.GetMessageCount())
		assert.Equal(t, 1, buffer2.GetMessageCount())

		msg1 := buffer1.GetMessages()[0]
		msg2 := buffer2.GetMessages()[0]

		assert.Equal(t, "req-1", msg1.RequestID)
		assert.Equal(t, "req-2", msg2.RequestID)
	})
}

// =============================================================================
// Streaming Message Tests
// =============================================================================

func TestBufferStreamingMessage(t *testing.T) {
	t.Run("AddStreamingMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		buffer.AddStreamingMessage(
			"msg-stream-1",
			"text",
			map[string]interface{}{"content": "# Title\n\n"},
			"block-1",
			"thread-1",
			"assistant-1",
			nil,
		)

		assert.Equal(t, 1, buffer.GetMessageCount())

		// Verify streaming message is added
		msg := buffer.GetStreamingMessage("msg-stream-1")
		assert.NotNil(t, msg)
		assert.Equal(t, "msg-stream-1", msg.MessageID)
		assert.Equal(t, "text", msg.Type)
		assert.Equal(t, "# Title\n\n", msg.Props["content"])
		assert.True(t, msg.IsStreaming)
	})

	t.Run("AppendMessageContent", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add streaming message
		buffer.AddStreamingMessage(
			"msg-stream-2",
			"text",
			map[string]interface{}{"content": "Initial "},
			"", "", "", nil,
		)

		// Append content
		ok := buffer.AppendMessageContent("msg-stream-2", "Line 1\n")
		assert.True(t, ok)

		ok = buffer.AppendMessageContent("msg-stream-2", "Line 2\n")
		assert.True(t, ok)

		// Verify accumulated content
		msg := buffer.GetStreamingMessage("msg-stream-2")
		assert.NotNil(t, msg)
		assert.Equal(t, "Initial Line 1\nLine 2\n", msg.Props["content"])
	})

	t.Run("AppendToNonExistentMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Try to append to non-existent message
		ok := buffer.AppendMessageContent("non-existent", "content")
		assert.False(t, ok)
	})

	t.Run("AppendToCompletedMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add and complete streaming message
		buffer.AddStreamingMessage(
			"msg-stream-3",
			"text",
			map[string]interface{}{"content": "Initial"},
			"", "", "", nil,
		)
		buffer.CompleteStreamingMessage("msg-stream-3")

		// Try to append to completed message (should fail)
		ok := buffer.AppendMessageContent("msg-stream-3", " more")
		assert.False(t, ok)
	})

	t.Run("CompleteStreamingMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add streaming message
		buffer.AddStreamingMessage(
			"msg-stream-4",
			"text",
			map[string]interface{}{"content": "Hello "},
			"", "", "", nil,
		)

		// Append content
		buffer.AppendMessageContent("msg-stream-4", "World!")

		// Complete the message
		content, ok := buffer.CompleteStreamingMessage("msg-stream-4")
		assert.True(t, ok)
		assert.Equal(t, "Hello World!", content)

		// Message should no longer be streaming
		msg := buffer.GetStreamingMessage("msg-stream-4")
		assert.Nil(t, msg)

		// But should still exist in messages
		messages := buffer.GetMessages()
		assert.Equal(t, 1, len(messages))
		assert.False(t, messages[0].IsStreaming)
	})

	t.Run("CompleteNonExistentMessage", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		content, ok := buffer.CompleteStreamingMessage("non-existent")
		assert.False(t, ok)
		assert.Empty(t, content)
	})

	t.Run("StreamingMessageWorkflow", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "deepseek", "")

		// Simulate a typical streaming workflow:
		// 1. SendStream sends initial content
		buffer.AddStreamingMessage(
			"msg-workflow",
			"text",
			map[string]interface{}{"content": "# Available Tests\n\n"},
			"block-main",
			"",
			"assistant-1",
			nil,
		)

		// 2. Multiple Append calls add content
		buffer.AppendMessageContent("msg-workflow", "Send one of these keywords:\n\n")
		buffer.AppendMessageContent("msg-workflow", "- **basic** - Basic tests\n")
		buffer.AppendMessageContent("msg-workflow", "- **advanced** - Advanced tests\n")

		// 3. End completes the message
		finalContent, ok := buffer.CompleteStreamingMessage("msg-workflow")
		assert.True(t, ok)

		expectedContent := "# Available Tests\n\nSend one of these keywords:\n\n- **basic** - Basic tests\n- **advanced** - Advanced tests\n"
		assert.Equal(t, expectedContent, finalContent)

		// Verify final message state
		messages := buffer.GetMessages()
		assert.Equal(t, 1, len(messages))
		assert.Equal(t, "msg-workflow", messages[0].MessageID)
		assert.Equal(t, "deepseek", messages[0].Connector) // Connector should be set
		assert.False(t, messages[0].IsStreaming)
	})

	t.Run("MixedStreamingAndRegularMessages", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add user input (regular)
		buffer.AddUserInput("Hello", "user1")

		// Add streaming assistant message
		buffer.AddStreamingMessage(
			"msg-stream",
			"text",
			map[string]interface{}{"content": "Hi "},
			"", "", "", nil,
		)
		buffer.AppendMessageContent("msg-stream", "there!")
		buffer.CompleteStreamingMessage("msg-stream")

		// Add regular assistant message
		buffer.AddAssistantMessage("M3", "text", map[string]interface{}{"content": "How can I help?"}, "", "", "", nil)

		// Verify all messages
		messages := buffer.GetMessages()
		assert.Equal(t, 3, len(messages))

		// Check sequence
		assert.Equal(t, 1, messages[0].Sequence)
		assert.Equal(t, 2, messages[1].Sequence)
		assert.Equal(t, 3, messages[2].Sequence)

		// Check content
		assert.Equal(t, "user", messages[0].Role)
		assert.Equal(t, "Hi there!", messages[1].Props["content"])
		assert.Equal(t, "How can I help?", messages[2].Props["content"])
	})

	t.Run("StreamingMessageWithEmptyInitialContent", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add streaming message with nil props
		buffer.AddStreamingMessage(
			"msg-empty",
			"text",
			nil,
			"", "", "", nil,
		)

		// Append content
		buffer.AppendMessageContent("msg-empty", "Content")

		// Complete
		content, ok := buffer.CompleteStreamingMessage("msg-empty")
		assert.True(t, ok)
		assert.Equal(t, "Content", content)
	})

	t.Run("ConcurrentStreamingOperations", func(t *testing.T) {
		buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")

		// Add streaming message
		buffer.AddStreamingMessage(
			"msg-concurrent",
			"text",
			map[string]interface{}{"content": ""},
			"", "", "", nil,
		)

		// Concurrent appends with fixed-length content
		var wg sync.WaitGroup
		for i := 0; i < 100; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				buffer.AppendMessageContent("msg-concurrent", "x")
			}()
		}
		wg.Wait()

		// Complete
		content, ok := buffer.CompleteStreamingMessage("msg-concurrent")
		assert.True(t, ok)

		// Content should have 100 'x' characters
		assert.Equal(t, 100, len(content))
	})
}