gochat/internal/handler/api/v1/captain_task_extended_handler_test.go

package v1

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
	"github.com/gochat/gochat/internal/service"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

type mockTaskExtHandlerLLM struct {
	response *llm.ChatResponse
	err      error
}

func (m *mockTaskExtHandlerLLM) ChatCompletion(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) {
	return m.response, m.err
}
func (m *mockTaskExtHandlerLLM) CreateEmbedding(_ context.Context, _ llm.EmbeddingRequest) (*llm.EmbeddingResponse, error) {
	return nil, nil
}
func (m *mockTaskExtHandlerLLM) ChatCompletionStream(_ context.Context, _ llm.ChatRequest, _ func(llm.StreamChunk) error) error {
	return nil
}

func setupTaskExtendedHandlerTest(t *testing.T, mockLLM llm.Provider) (*CaptainTaskExtendedHandler, *gorm.DB) {
	t.Helper()
	gin.SetMode(gin.TestMode)
	dbName := fmt.Sprintf("file:%s?mode=memory&cache=private", t.Name())
	db, err := gorm.Open(sqlite.Open(dbName), &gorm.Config{})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(
		&model.CaptainPreference{},
		&model.CaptainAssistant{},
		&model.Conversation{},
		&model.Message{},
		&model.Inbox{},
		&model.Contact{},
		&model.CopilotSuggestionMessage{},
	))
	convRepo := repository.NewConversationRepo(db)
	msgRepo := repository.NewMessageRepo(db)
	assistantRepo := repository.NewCaptainAssistantRepo(db)
	prefRepo := repository.NewCaptainPreferenceRepo(db)
	suggestionRepo := repository.NewCopilotSuggestionRepo(db)
	svc := service.NewCaptainTaskExtendedService(convRepo, msgRepo, assistantRepo, prefRepo, mockLLM, suggestionRepo)
	handler := NewCaptainTaskExtendedHandler(svc)
	return handler, db
}

func TestCaptainTaskExtendedHandler_LabelSuggestion(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{
		response: &llm.ChatResponse{
			Choices: []llm.ChatChoice{
				{Message: llm.ChatMessage{Role: "assistant", Content: `{"labels": ["billing"], "priority": "high", "reason": "billing issue"}`}},
			},
		},
	}
	handler, db := setupTaskExtendedHandlerTest(t, mockLLM)

	// Seed conversation + message
	inbox := &model.Inbox{AccountID: 1, Name: "Inbox", ChannelType: "web_widget"}
	require.NoError(t, db.Create(inbox).Error)
	conv := &model.Conversation{AccountID: 1, InboxID: inbox.ID, Status: "open"}
	require.NoError(t, db.Create(conv).Error)
	msg := &model.Message{ConversationID: conv.ID, AccountID: 1, InboxID: inbox.ID, SenderType: "contact", Content: "billing help", ContentType: "text", MessageType: "incoming"}
	require.NoError(t, db.Create(msg).Error)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/accounts/1/captain/tasks/label_suggestion?conversation_ids="+fmt.Sprintf("%d", conv.ID), nil)

	handler.LabelSuggestion(c)

	assert.Equal(t, http.StatusOK, w.Code)

	var resp handlerTestResponse
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	assert.True(t, resp.Success)
}

func TestCaptainTaskExtendedHandler_LabelSuggestion_MissingConversationIDs(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{}
	handler, _ := setupTaskExtendedHandlerTest(t, mockLLM)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/accounts/1/captain/tasks/label_suggestion", nil)

	handler.LabelSuggestion(c)

	assert.Equal(t, http.StatusBadRequest, w.Code)
}

func TestCaptainTaskExtendedHandler_LabelSuggestion_ChatwootPostRawPayload(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{
		response: &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Role: "assistant", Content: "billing, urgent"}}}},
	}
	handler, db := setupTaskExtendedHandlerTest(t, mockLLM)
	displayID := uint(77)
	inbox := &model.Inbox{AccountID: 1, Name: "Inbox", ChannelType: "web_widget"}
	require.NoError(t, db.Create(inbox).Error)
	conv := &model.Conversation{AccountID: 1, InboxID: inbox.ID, DisplayID: &displayID, Status: "open", ChannelType: "web_widget", Channel: "web_widget"}
	require.NoError(t, db.Create(conv).Error)
	msg := &model.Message{ConversationID: conv.ID, AccountID: 1, InboxID: inbox.ID, SenderType: "contact", Content: "billing help", ContentType: "text", MessageType: "incoming"}
	require.NoError(t, db.Create(msg).Error)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/label_suggestion", bytes.NewReader([]byte(`{"conversation_display_id":77}`)))
	c.Request.Header.Set("Content-Type", "application/json")

	handler.LabelSuggestion(c)

	require.Equal(t, http.StatusOK, w.Code)
	var resp map[string]interface{}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	assert.Equal(t, "billing, urgent", resp["message"])
	assert.NotContains(t, resp, "success")
	require.Contains(t, resp, "follow_up_context")

	var stored []model.CopilotSuggestionMessage
	require.NoError(t, db.Find(&stored).Error)
	require.Len(t, stored, 1)
	assert.Equal(t, "billing, urgent", stored[0].Content)
}

func TestCaptainTaskExtendedHandler_FollowUp(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{
		response: &llm.ChatResponse{
			Choices: []llm.ChatChoice{
				{Message: llm.ChatMessage{Role: "assistant", Content: `{"follow_ups": [{"title": "Follow up", "description": "desc", "priority": "medium", "due_date_hint": "24h"}]}`}},
			},
		},
	}
	handler, db := setupTaskExtendedHandlerTest(t, mockLLM)

	inbox := &model.Inbox{AccountID: 1, Name: "Inbox", ChannelType: "web_widget"}
	require.NoError(t, db.Create(inbox).Error)
	conv := &model.Conversation{AccountID: 1, InboxID: inbox.ID, Status: "open"}
	require.NoError(t, db.Create(conv).Error)
	msg := &model.Message{ConversationID: conv.ID, AccountID: 1, InboxID: inbox.ID, SenderType: "contact", Content: "need follow up", ContentType: "text", MessageType: "incoming"}
	require.NoError(t, db.Create(msg).Error)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/accounts/1/captain/tasks/follow_up?conversation_ids="+fmt.Sprintf("%d", conv.ID), nil)

	handler.FollowUp(c)

	assert.Equal(t, http.StatusOK, w.Code)
}

func TestCaptainTaskExtendedHandler_FollowUp_MissingConversationIDs(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{}
	handler, _ := setupTaskExtendedHandlerTest(t, mockLLM)

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/accounts/1/captain/tasks/follow_up", nil)

	handler.FollowUp(c)

	assert.Equal(t, http.StatusBadRequest, w.Code)
}

func TestCaptainTaskExtendedHandler_FollowUp_ChatwootPostUpdatesContext(t *testing.T) {
	mockLLM := &mockTaskExtHandlerLLM{
		response: &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Role: "assistant", Content: "Refined answer"}}}},
	}
	handler, db := setupTaskExtendedHandlerTest(t, mockLLM)
	displayID := uint(88)
	inbox := &model.Inbox{AccountID: 1, Name: "Inbox", ChannelType: "web_widget"}
	require.NoError(t, db.Create(inbox).Error)
	conv := &model.Conversation{AccountID: 1, InboxID: inbox.ID, DisplayID: &displayID, Status: "open", ChannelType: "web_widget", Channel: "web_widget"}
	require.NoError(t, db.Create(conv).Error)

	body := []byte(`{
    "conversation_display_id": 88,
    "message": "Make it warmer",
    "follow_up_context": {
      "event_name": "professional",
      "original_context": "Original draft",
      "last_response": "Previous answer",
      "conversation_history": []
    }
  }`)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Params = gin.Params{{Key: "id", Value: "1"}}
	c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/follow_up", bytes.NewReader(body))
	c.Request.Header.Set("Content-Type", "application/json")

	handler.FollowUp(c)

	require.Equal(t, http.StatusOK, w.Code)
	var resp map[string]interface{}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
	assert.Equal(t, "Refined answer", resp["message"])
	ctx := resp["follow_up_context"].(map[string]interface{})
	assert.Equal(t, "Refined answer", ctx["last_response"])
	history := ctx["conversation_history"].([]interface{})
	require.Len(t, history, 2)
}