kandev/apps/backend/internal/orchestrator/messagequeue/service_test.go

package messagequeue

import (
	"context"
	"sync"
	"testing"
	"time"

	"github.com/kandev/kandev/internal/common/logger"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func setupService(t *testing.T) *Service {
	log, err := logger.NewLogger(logger.LoggingConfig{
		Level:      "error",
		Format:     "console",
		OutputPath: "stdout",
	})
	require.NoError(t, err)
	return NewService(log)
}

func TestQueueMessage(t *testing.T) {
	t.Run("successfully queues a message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "test content", "model-1", "user-1", false, nil)

		require.NoError(t, err)
		assert.NotEmpty(t, msg.ID)
		assert.Equal(t, "session-1", msg.SessionID)
		assert.Equal(t, "task-1", msg.TaskID)
		assert.Equal(t, "test content", msg.Content)
		assert.Equal(t, "model-1", msg.Model)
		assert.Equal(t, "user-1", msg.QueuedBy)
		assert.False(t, msg.PlanMode)
		assert.NotZero(t, msg.QueuedAt)
	})

	t.Run("replaces existing queued message for same session", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue first message
		msg1, err := svc.QueueMessage(ctx, "session-1", "task-1", "first message", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// Queue second message for same session
		msg2, err := svc.QueueMessage(ctx, "session-1", "task-1", "second message", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// IDs should be different
		assert.NotEqual(t, msg1.ID, msg2.ID)

		// Only the second message should be queued
		status := svc.GetStatus(ctx, "session-1")
		assert.True(t, status.IsQueued)
		assert.Equal(t, msg2.ID, status.Message.ID)
		assert.Equal(t, "second message", status.Message.Content)
	})

	t.Run("queues messages with attachments", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		attachments := []MessageAttachment{
			{Type: "image", Data: "base64data", MimeType: "image/png"},
		}

		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "message with attachment", "model-1", "user-1", false, attachments)

		require.NoError(t, err)
		assert.Len(t, msg.Attachments, 1)
		assert.Equal(t, "image", msg.Attachments[0].Type)
	})

	t.Run("supports plan mode flag", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "plan mode message", "model-1", "user-1", true, nil)

		require.NoError(t, err)
		assert.True(t, msg.PlanMode)
	})
}

func TestTransferSession(t *testing.T) {
	t.Run("moves queued message and pending move from old to new session", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		_, err := svc.QueueMessage(ctx, "old-sess", "task-1", "hand-off prompt", "", "mcp-move-task", false, nil)
		require.NoError(t, err)
		svc.SetPendingMove(ctx, "old-sess", &PendingMove{
			TaskID:         "task-1",
			WorkflowStepID: "step-b",
		})

		svc.TransferSession(ctx, "old-sess", "new-sess")

		// Old session is empty.
		_, exists := svc.TakeQueued(ctx, "old-sess")
		assert.False(t, exists)
		_, exists = svc.TakePendingMove(ctx, "old-sess")
		assert.False(t, exists)

		// New session has both.
		msg, exists := svc.TakeQueued(ctx, "new-sess")
		require.True(t, exists)
		assert.Equal(t, "hand-off prompt", msg.Content)
		assert.Equal(t, "new-sess", msg.SessionID)

		move, exists := svc.TakePendingMove(ctx, "new-sess")
		require.True(t, exists)
		assert.Equal(t, "step-b", move.WorkflowStepID)
	})

	t.Run("no-op when source session is empty", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		svc.TransferSession(ctx, "empty-sess", "new-sess")

		_, exists := svc.TakeQueued(ctx, "new-sess")
		assert.False(t, exists)
	})
}

func TestPendingMove(t *testing.T) {
	t.Run("set then take returns the move and clears it", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		svc.SetPendingMove(ctx, "session-1", &PendingMove{
			TaskID:         "task-1",
			WorkflowID:     "wf-1",
			WorkflowStepID: "step-2",
			Position:       3,
		})

		got, exists := svc.TakePendingMove(ctx, "session-1")
		require.True(t, exists)
		assert.Equal(t, "task-1", got.TaskID)
		assert.Equal(t, "step-2", got.WorkflowStepID)
		assert.Equal(t, 3, got.Position)
		assert.NotZero(t, got.QueuedAt)

		// Idempotent — second take returns nothing.
		_, exists = svc.TakePendingMove(ctx, "session-1")
		assert.False(t, exists)
	})

	t.Run("setting twice replaces the previous move", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		svc.SetPendingMove(ctx, "session-1", &PendingMove{TaskID: "task-1", WorkflowStepID: "step-a"})
		svc.SetPendingMove(ctx, "session-1", &PendingMove{TaskID: "task-1", WorkflowStepID: "step-b"})

		got, exists := svc.TakePendingMove(ctx, "session-1")
		require.True(t, exists)
		assert.Equal(t, "step-b", got.WorkflowStepID)
	})
}

func TestTakeQueued(t *testing.T) {
	t.Run("retrieves and removes queued message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue a message
		original, err := svc.QueueMessage(ctx, "session-1", "task-1", "test content", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// Take the message
		msg, exists := svc.TakeQueued(ctx, "session-1")

		assert.True(t, exists)
		assert.Equal(t, original.ID, msg.ID)
		assert.Equal(t, "test content", msg.Content)

		// Should be removed after taking
		status := svc.GetStatus(ctx, "session-1")
		assert.False(t, status.IsQueued)
		assert.Nil(t, status.Message)
	})

	t.Run("returns false when no message queued", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		msg, exists := svc.TakeQueued(ctx, "session-1")

		assert.False(t, exists)
		assert.Nil(t, msg)
	})

	t.Run("is idempotent - second take returns false", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue and take once
		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "test", "model-1", "user-1", false, nil)
		require.NoError(t, err)
		_, exists := svc.TakeQueued(ctx, "session-1")
		require.True(t, exists)

		// Second take should return false
		_, exists = svc.TakeQueued(ctx, "session-1")
		assert.False(t, exists)
	})
}

func TestCancelQueued(t *testing.T) {
	t.Run("cancels and removes queued message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue a message
		original, err := svc.QueueMessage(ctx, "session-1", "task-1", "test content", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// Cancel it
		cancelled, err := svc.CancelQueued(ctx, "session-1")

		require.NoError(t, err)
		assert.Equal(t, original.ID, cancelled.ID)

		// Should be removed
		status := svc.GetStatus(ctx, "session-1")
		assert.False(t, status.IsQueued)
	})

	t.Run("returns error when no message queued", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		_, err := svc.CancelQueued(ctx, "session-1")

		assert.Error(t, err)
		assert.Contains(t, err.Error(), "no queued message")
	})
}

func TestGetStatus(t *testing.T) {
	t.Run("returns queued status with message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "test", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		status := svc.GetStatus(ctx, "session-1")

		assert.True(t, status.IsQueued)
		assert.NotNil(t, status.Message)
		assert.Equal(t, msg.ID, status.Message.ID)
	})

	t.Run("returns not queued status when no message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		status := svc.GetStatus(ctx, "session-1")

		assert.False(t, status.IsQueued)
		assert.Nil(t, status.Message)
	})
}

func TestAppendContent(t *testing.T) {
	t.Run("appends content to existing queued message", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "original content", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		err = svc.AppendContent(ctx, "session-1", "appended content")
		require.NoError(t, err)

		status := svc.GetStatus(ctx, "session-1")
		assert.True(t, status.IsQueued)
		assert.Equal(t, "original content\n\n---\n\nappended content", status.Message.Content)
	})

	t.Run("appends multiple times", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "first", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		require.NoError(t, svc.AppendContent(ctx, "session-1", "second"))
		require.NoError(t, svc.AppendContent(ctx, "session-1", "third"))

		status := svc.GetStatus(ctx, "session-1")
		assert.Equal(t, "first\n\n---\n\nsecond\n\n---\n\nthird", status.Message.Content)
	})

	t.Run("returns error when no message queued", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		err := svc.AppendContent(ctx, "session-1", "content")
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "no queued message")
	})

	t.Run("preserves message ID after append", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "original", "model-1", "user-1", false, nil)
		require.NoError(t, err)
		originalID := msg.ID

		require.NoError(t, svc.AppendContent(ctx, "session-1", "extra"))

		status := svc.GetStatus(ctx, "session-1")
		assert.Equal(t, originalID, status.Message.ID)
	})
}

func TestUpdateMessage(t *testing.T) {
	t.Run("updates existing queued message content", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue a message
		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "original content", "model-1", "user-1", false, nil)
		require.NoError(t, err)
		originalID := msg.ID

		// Update it
		err = svc.UpdateMessage(ctx, "session-1", "updated content")

		require.NoError(t, err)

		// Verify update
		status := svc.GetStatus(ctx, "session-1")
		assert.True(t, status.IsQueued)
		assert.Equal(t, originalID, status.Message.ID)
		assert.Equal(t, "updated content", status.Message.Content)
	})

	t.Run("returns error when no message queued", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		err := svc.UpdateMessage(ctx, "session-1", "new content")

		assert.Error(t, err)
		assert.Contains(t, err.Error(), "no queued message")
	})

	t.Run("allows empty content updates", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue a message
		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "original", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// Update to empty (validation happens at handler level)
		err = svc.UpdateMessage(ctx, "session-1", "")

		require.NoError(t, err)
		status := svc.GetStatus(ctx, "session-1")
		assert.Equal(t, "", status.Message.Content)
	})
}

func TestConcurrentAccess(t *testing.T) {
	t.Run("handles concurrent queue operations safely", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		var wg sync.WaitGroup
		numGoroutines := 100

		// Concurrently queue messages to different sessions
		for i := 0; i < numGoroutines; i++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				sessionID := "session-" + string(rune('A'+idx%26))
				_, err := svc.QueueMessage(ctx, sessionID, "task-1", "content", "model-1", "user-1", false, nil)
				assert.NoError(t, err)
			}(i)
		}

		wg.Wait()

		// Verify all sessions have messages
		sessionIDs := make(map[string]bool)
		for i := 0; i < numGoroutines; i++ {
			sessionID := "session-" + string(rune('A'+i%26))
			sessionIDs[sessionID] = true
		}

		for sessionID := range sessionIDs {
			status := svc.GetStatus(ctx, sessionID)
			assert.True(t, status.IsQueued, "session %s should be queued", sessionID)
		}
	})

	t.Run("handles concurrent take operations safely", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue a message
		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "test", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		var wg sync.WaitGroup
		successCount := 0
		var mu sync.Mutex

		// Try to take the message from multiple goroutines
		for i := 0; i < 10; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				_, exists := svc.TakeQueued(ctx, "session-1")
				if exists {
					mu.Lock()
					successCount++
					mu.Unlock()
				}
			}()
		}

		wg.Wait()

		// Only one should succeed
		assert.Equal(t, 1, successCount, "only one goroutine should successfully take the message")

		// Message should be gone
		status := svc.GetStatus(ctx, "session-1")
		assert.False(t, status.IsQueued)
	})

	t.Run("handles concurrent updates and reads safely", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue initial message
		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "initial", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		var wg sync.WaitGroup
		numUpdates := 50
		numReads := 50

		// Concurrent updates
		for i := 0; i < numUpdates; i++ {
			wg.Add(1)
			go func(idx int) {
				defer wg.Done()
				_ = svc.UpdateMessage(ctx, "session-1", "update-"+string(rune('0'+idx%10)))
			}(i)
		}

		// Concurrent reads
		for i := 0; i < numReads; i++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				status := svc.GetStatus(ctx, "session-1")
				assert.True(t, status.IsQueued)
			}()
		}

		wg.Wait()

		// Final state should be consistent
		status := svc.GetStatus(ctx, "session-1")
		assert.True(t, status.IsQueued)
		assert.NotEmpty(t, status.Message.Content)
	})
}

func TestMultipleSessionsIsolation(t *testing.T) {
	t.Run("sessions are isolated from each other", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		// Queue messages for multiple sessions
		_, err := svc.QueueMessage(ctx, "session-1", "task-1", "message-1", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		_, err = svc.QueueMessage(ctx, "session-2", "task-2", "message-2", "model-1", "user-1", false, nil)
		require.NoError(t, err)

		// Take from session-1
		msg1, exists := svc.TakeQueued(ctx, "session-1")
		require.True(t, exists)
		assert.Equal(t, "message-1", msg1.Content)

		// session-2 should still have its message
		status2 := svc.GetStatus(ctx, "session-2")
		assert.True(t, status2.IsQueued)
		assert.Equal(t, "message-2", status2.Message.Content)

		// session-1 should be empty
		status1 := svc.GetStatus(ctx, "session-1")
		assert.False(t, status1.IsQueued)
	})
}

func TestQueuedMessageTimestamp(t *testing.T) {
	t.Run("sets queued timestamp correctly", func(t *testing.T) {
		svc := setupService(t)
		ctx := context.Background()

		before := time.Now()
		msg, err := svc.QueueMessage(ctx, "session-1", "task-1", "test", "model-1", "user-1", false, nil)
		after := time.Now()

		require.NoError(t, err)
		assert.True(t, msg.QueuedAt.After(before) || msg.QueuedAt.Equal(before))
		assert.True(t, msg.QueuedAt.Before(after) || msg.QueuedAt.Equal(after))
	})
}