gochat/backend/internal/service/copilot_context_service_test.go

package service

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
)

// ========== Mock LLM Provider for CopilotContext tests ==========

type mockContextLLMProvider struct {
	chatResponse    *llm.ChatResponse
	chatError       error
	embeddingResp   *llm.EmbeddingResponse
	embeddingError  error
	streamChunks    []llm.StreamChunk
	streamError     error
	lastChatRequest *llm.ChatRequest
	lastEmbeddingReq *llm.EmbeddingRequest
	chatCalls       int
	embeddingCalls  int
}

func (m *mockContextLLMProvider) ChatCompletion(ctx context.Context, req llm.ChatRequest) (*llm.ChatResponse, error) {
	m.lastChatRequest = &req
	m.chatCalls++
	return m.chatResponse, m.chatError
}

func (m *mockContextLLMProvider) CreateEmbedding(ctx context.Context, req llm.EmbeddingRequest) (*llm.EmbeddingResponse, error) {
	m.lastEmbeddingReq = &req
	m.embeddingCalls++
	return m.embeddingResp, m.embeddingError
}

func (m *mockContextLLMProvider) ChatCompletionStream(ctx context.Context, req llm.ChatRequest, onChunk func(llm.StreamChunk) error) error {
	for _, chunk := range m.streamChunks {
		if err := onChunk(chunk); err != nil {
			return err
		}
	}
	return m.streamError
}

// ========== Setup ==========

func setupCopilotContextTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	require.NoError(t, err)

	err = db.AutoMigrate(
		&model.Account{},
		&model.Contact{},
		&model.Conversation{},
		&model.Message{},
		&model.CaptainAssistant{},
		&model.Inbox{},
	)
	require.NoError(t, err)

	t.Cleanup(func() {
		sqlDB, _ := db.DB()
		sqlDB.Close()
	})
	return db
}

func setupCopilotContextService(t *testing.T) (*CopilotContextService, *mockContextLLMProvider, *gorm.DB) {
	t.Helper()
	db := setupCopilotContextTestDB(t)
	mockProvider := &mockContextLLMProvider{
		chatResponse: &llm.ChatResponse{
			Choices: []llm.ChatChoice{
				{Message: llm.ChatMessage{Content: "test summary"}, FinishReason: "stop"},
			},
		},
		embeddingResp: &llm.EmbeddingResponse{
			Data: []llm.EmbeddingData{
				{Embedding: []float64{0.1, 0.2, 0.3}},
			},
		},
	}
	messageRepo := repository.NewMessageRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	contactRepo := repository.NewContactRepo(db)

	svc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockProvider)
	require.NotNil(t, svc)
	return svc, mockProvider, db
}

// createContextTestConversation creates a conversation with messages for context tests.
// (Separate from captain_task_service_test.go's createTestConversationWithMessages to avoid signature conflicts.)
func createContextTestConversation(t *testing.T, db *gorm.DB, accountID uint) (*model.Conversation, []*model.Message) {
	t.Helper()
	contact := &model.Contact{
		Name:  "Test Contact",
		Email: "test@example.com",
	}
	require.NoError(t, db.Create(contact).Error)

	inbox := &model.Inbox{
		Name:      "Test Inbox",
		AccountID: accountID,
	}
	require.NoError(t, db.Create(inbox).Error)

	conversation := &model.Conversation{
		AccountID: accountID,
		InboxID:   inbox.ID,
		ContactID: contact.ID,
		Status:    "open",
	}
	require.NoError(t, db.Create(conversation).Error)

	var messages []*model.Message
	for i, content := range []string{
		"Hello, I need help with my order",
		"Sure, let me look into that",
		"Order #12345 hasn't arrived",
		"I'll check the tracking info",
	} {
		msgType := "incoming"
		if i%2 == 1 {
			msgType = "outgoing"
		}
		msg := &model.Message{
			ConversationID: conversation.ID,
			Content:         content,
			MessageType:     msgType,
			ContentType:     "text",
			AccountID:       accountID,
		}
		require.NoError(t, db.Create(msg).Error)
		messages = append(messages, msg)
	}

	return conversation, messages
}

// helper to create stream chunks matching OpenAI SSE format
func makeStreamChunks(contents []string) []llm.StreamChunk {
	var chunks []llm.StreamChunk
	for i, content := range contents {
		finishReason := ""
		if i == len(contents)-1 {
			finishReason = "stop"
		}
		chunks = append(chunks, llm.StreamChunk{
			Choices: []llm.StreamChoice{
				{
					Index:        0,
					Delta:        llm.StreamDelta{Content: content},
					FinishReason: finishReason,
				},
			},
		})
	}
	return chunks
}

// ========== Context Enrichment Tests ==========

func TestCopilotContextService_GetCurrentViewingContext(t *testing.T) {
	svc, _, db := setupCopilotContextService(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	conversation, _ := createContextTestConversation(t, db, accountID)

	ctx := context.Background()
	result, err := svc.GetCurrentViewingContext(ctx, accountID, conversation.ID)

	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.Equal(t, conversation.ID, result.ConversationID)
	assert.Equal(t, "Test Contact", result.ContactName)
	assert.Equal(t, "test@example.com", result.ContactEmail)
	assert.NotEmpty(t, result.Messages)
}

func TestCopilotContextService_GetCurrentViewingContext_NoContact(t *testing.T) {
	svc, _, db := setupCopilotContextService(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	inbox := &model.Inbox{Name: "Test Inbox", AccountID: accountID}
	require.NoError(t, db.Create(inbox).Error)

	conversation := &model.Conversation{
		AccountID: accountID,
		InboxID:   inbox.ID,
		ContactID: 0,
		Status:    "open",
	}
	require.NoError(t, db.Create(conversation).Error)

	ctx := context.Background()
	result, err := svc.GetCurrentViewingContext(ctx, accountID, conversation.ID)

	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.Empty(t, result.ContactName)
}

func TestCopilotContextService_GetCurrentViewingContext_NonexistentConversation(t *testing.T) {
	svc, _, _ := setupCopilotContextService(t)

	ctx := context.Background()
	result, err := svc.GetCurrentViewingContext(ctx, 999, 99999)

	assert.Error(t, err)
	assert.Nil(t, result)
}

func TestCopilotContextService_ContextMessageFormatting(t *testing.T) {
	svc, _, db := setupCopilotContextService(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	conversation, _ := createContextTestConversation(t, db, accountID)

	ctx := context.Background()
	result, err := svc.GetCurrentViewingContext(ctx, accountID, conversation.ID)

	require.NoError(t, err)
	require.NotNil(t, result)

	// Verify customer/agent role assignment
	for _, m := range result.Messages {
		if m.Role == "customer" {
			assert.Contains(t, []string{"Hello, I need help with my order", "Order #12345 hasn't arrived"}, m.Content)
		} else if m.Role == "agent" {
			assert.Contains(t, []string{"Sure, let me look into that", "I'll check the tracking info"}, m.Content)
		}
	}
}

func TestCopilotContextService_EnrichmentIntegration(t *testing.T) {
	db := setupCopilotContextTestDB(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	mockLLM := &mockContextLLMProvider{
		chatResponse: &llm.ChatResponse{
			Choices: []llm.ChatChoice{
				{Message: llm.ChatMessage{Content: "I can help with that order"}, FinishReason: "stop"},
			},
		},
		streamChunks: makeStreamChunks([]string{"I can", " help", " with that"}),
	}

	// Create test data
	contact := &model.Contact{Name: "Alice", Email: "alice@example.com"}
	require.NoError(t, db.Create(contact).Error)
	inbox := &model.Inbox{Name: "Support", AccountID: accountID}
	require.NoError(t, db.Create(inbox).Error)
	conversation := &model.Conversation{
		AccountID: accountID, InboxID: inbox.ID, ContactID: contact.ID,
		Status: "open",
	}
	require.NoError(t, db.Create(conversation).Error)
	msg := &model.Message{
		ConversationID: conversation.ID, Content: "Help me!", MessageType: "incoming",
		ContentType: "text", AccountID: accountID,
	}
	require.NoError(t, db.Create(msg).Error)

	assistantRepo := repository.NewCaptainAssistantRepo(db)
	responseRepo := repository.NewCaptainAssistantResponseRepo(db)
	customToolRepo := repository.NewCaptainCustomToolRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	messageRepo := repository.NewMessageRepo(db)
	contactRepo := repository.NewContactRepo(db)

	copilotContextSvc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockLLM)

	captainTaskSvc := NewCaptainTaskService(
		assistantRepo, responseRepo, customToolRepo,
		conversationRepo, messageRepo,
		mockLLM, copilotContextSvc,
	)

	require.NotNil(t, captainTaskSvc)
	assert.NotNil(t, captainTaskSvc.copilotContextSvc)
}

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

func TestCaptainTaskService_ReplySuggestionStream(t *testing.T) {
	db := setupCopilotContextTestDB(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	streamChunks := makeStreamChunks([]string{"Hello!", " I can", " help you."})

	mockLLM := &mockLLMProvider{
		streamChunks: streamChunks,
	}

	assistantRepo := repository.NewCaptainAssistantRepo(db)
	responseRepo := repository.NewCaptainAssistantResponseRepo(db)
	customToolRepo := repository.NewCaptainCustomToolRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	messageRepo := repository.NewMessageRepo(db)
	contactRepo := repository.NewContactRepo(db)

	copilotContextSvc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockLLM)
	captainTaskSvc := NewCaptainTaskService(
		assistantRepo, responseRepo, customToolRepo,
		conversationRepo, messageRepo,
		mockLLM, copilotContextSvc,
	)

	// Create assistant and conversation
	assistant := &model.CaptainAssistant{
		AccountID: accountID,
		Name:      "Test Assistant",
		Status:    model.AssistantStatusActive,
		Config:    marshalAssistantConfig(model.DefaultAssistantConfig()),
	}
	require.NoError(t, db.Create(assistant).Error)

	conversation := &model.Conversation{
		AccountID: accountID,
		Status:    "open",
	}
	require.NoError(t, db.Create(conversation).Error)

	req := &TaskReplySuggestionRequest{
		ConversationID: conversation.ID,
		AssistantID:    assistant.ID,
	}

	var collectedContents []string
	err := captainTaskSvc.ReplySuggestionStream(context.Background(), accountID, req, func(chunk llm.StreamChunk) error {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			if content != "" {
				collectedContents = append(collectedContents, content)
			}
		}
		return nil
	})

	require.NoError(t, err)
	assert.Equal(t, []string{"Hello!", " I can", " help you."}, collectedContents)
	assert.NotNil(t, mockLLM.lastStreamReq)
}

func TestCaptainTaskService_SummarizeStream(t *testing.T) {
	db := setupCopilotContextTestDB(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	streamChunks := makeStreamChunks([]string{"Customer", " asked about", " order status."})

	mockLLM := &mockLLMProvider{
		streamChunks: streamChunks,
	}

	assistantRepo := repository.NewCaptainAssistantRepo(db)
	responseRepo := repository.NewCaptainAssistantResponseRepo(db)
	customToolRepo := repository.NewCaptainCustomToolRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	messageRepo := repository.NewMessageRepo(db)
	contactRepo := repository.NewContactRepo(db)

	copilotContextSvc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockLLM)
	captainTaskSvc := NewCaptainTaskService(
		assistantRepo, responseRepo, customToolRepo,
		conversationRepo, messageRepo,
		mockLLM, copilotContextSvc,
	)

	conversation := &model.Conversation{
		AccountID: accountID,
		Status:    "open",
	}
	require.NoError(t, db.Create(conversation).Error)

	// Add messages to the conversation
	for _, content := range []string{"I need help", "Let me check"} {
		msg := &model.Message{
			ConversationID: conversation.ID,
			Content:         content,
			MessageType:     "incoming",
			ContentType:     "text",
			AccountID:       accountID,
		}
		require.NoError(t, db.Create(msg).Error)
	}

	req := &TaskSummarizeRequest{
		ConversationID: conversation.ID,
	}

	var collectedContents []string
	err := captainTaskSvc.SummarizeStream(context.Background(), accountID, req, func(chunk llm.StreamChunk) error {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			if content != "" {
				collectedContents = append(collectedContents, content)
			}
		}
		return nil
	})

	require.NoError(t, err)
	assert.Equal(t, []string{"Customer", " asked about", " order status."}, collectedContents)
}

func TestCaptainTaskService_RewriteStream(t *testing.T) {
	db := setupCopilotContextTestDB(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	streamChunks := makeStreamChunks([]string{"Dear", " customer,", " we apologize."})

	mockLLM := &mockLLMProvider{
		streamChunks: streamChunks,
	}

	assistantRepo := repository.NewCaptainAssistantRepo(db)
	responseRepo := repository.NewCaptainAssistantResponseRepo(db)
	customToolRepo := repository.NewCaptainCustomToolRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	messageRepo := repository.NewMessageRepo(db)
	contactRepo := repository.NewContactRepo(db)

	copilotContextSvc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockLLM)
	captainTaskSvc := NewCaptainTaskService(
		assistantRepo, responseRepo, customToolRepo,
		conversationRepo, messageRepo,
		mockLLM, copilotContextSvc,
	)

	req := &TaskRewriteRequest{
		Content: "sorry for delay",
		Tone:    "professional",
	}

	var collectedContents []string
	err := captainTaskSvc.RewriteStream(context.Background(), accountID, req, func(chunk llm.StreamChunk) error {
		if len(chunk.Choices) > 0 {
			content := chunk.Choices[0].Delta.Content
			if content != "" {
				collectedContents = append(collectedContents, content)
			}
		}
		return nil
	})

	require.NoError(t, err)
	assert.Equal(t, []string{"Dear", " customer,", " we apologize."}, collectedContents)
}

func TestCaptainTaskService_StreamErrorHandling(t *testing.T) {
	db := setupCopilotContextTestDB(t)

	account := &model.Account{Name: "Test Account"}
	require.NoError(t, db.Create(account).Error)
	accountID := account.ID

	mockLLM := &mockLLMProvider{
		streamError: context.DeadlineExceeded,
	}

	assistantRepo := repository.NewCaptainAssistantRepo(db)
	responseRepo := repository.NewCaptainAssistantResponseRepo(db)
	customToolRepo := repository.NewCaptainCustomToolRepo(db)
	conversationRepo := repository.NewConversationRepo(db)
	messageRepo := repository.NewMessageRepo(db)
	contactRepo := repository.NewContactRepo(db)

	copilotContextSvc := NewCopilotContextService(messageRepo, conversationRepo, contactRepo, mockLLM)
	captainTaskSvc := NewCaptainTaskService(
		assistantRepo, responseRepo, customToolRepo,
		conversationRepo, messageRepo,
		mockLLM, copilotContextSvc,
	)

	conversation := &model.Conversation{
		AccountID: accountID,
		Status:    "open",
	}
	require.NoError(t, db.Create(conversation).Error)

	req := &TaskReplySuggestionRequest{
		ConversationID: conversation.ID,
	}

	err := captainTaskSvc.ReplySuggestionStream(context.Background(), accountID, req, func(chunk llm.StreamChunk) error {
		return nil
	})

	assert.Error(t, err)
}