gochat/backend/internal/service/captain_conversation_service.go

package service

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/worker"
	applogger "github.com/gochat/gochat/pkg/logger"
	"gorm.io/datatypes"
	"gorm.io/gorm"
)

const TaskTypeCaptainConversationResponseBuilder = "captain:conversation_response_builder"

type CaptainConversationResponseBackend interface {
	GenerateCaptainConversationResponse(ctx context.Context, req CaptainConversationResponseRequest) (*CaptainConversationResponse, error)
}

type CaptainConversationResponseRequest struct {
	AccountID    uint
	Conversation *model.Conversation
	Assistant    *model.CaptainAssistant
	Messages     []CaptainConversationMessage
}

type CaptainConversationMessage struct {
	Role      string
	Content   string
	AgentName string
}

type CaptainConversationResponse struct {
	Content           string
	Action            string
	AgentName         string
	HandoffToolCalled bool
}

type CaptainConversationService struct {
	db          *gorm.DB
	llmProvider llm.Provider
	backend     CaptainConversationResponseBackend
	worker      *worker.WorkerPool
	// toolExecSvc enables LLM function calling (tool_call loop). nil = tools disabled.
	toolExecSvc *ToolExecutionService
}

func NewCaptainConversationService(db *gorm.DB, llmProvider llm.Provider) *CaptainConversationService {
	return &CaptainConversationService{db: db, llmProvider: llmProvider}
}

// SetToolExecutionService injects the tool execution service for function calling.
func (s *CaptainConversationService) SetToolExecutionService(svc *ToolExecutionService) {
	s.toolExecSvc = svc
}

func (s *CaptainConversationService) SetResponseBackend(backend CaptainConversationResponseBackend) {
	s.backend = backend
}

func (s *CaptainConversationService) SetWorkerPool(wp *worker.WorkerPool) {
	s.worker = wp
	RegisterCaptainConversationJobs(wp, s)
}

func (s *CaptainConversationService) BuildConversationResponseByAccount(ctx context.Context, accountID, conversationID, assistantID uint) (*model.Message, error) {
	if s == nil || s.db == nil {
		return nil, fmt.Errorf("captain conversation service database is required")
	}
	var conversation model.Conversation
	if err := s.db.WithContext(ctx).Where("account_id = ? AND id = ?", accountID, conversationID).First(&conversation).Error; err != nil {
		return nil, fmt.Errorf("conversation not found: %w", err)
	}
	if conversation.Status != string(model.ConversationStatusPending) {
		return nil, nil
	}
	var assistant model.CaptainAssistant
	if err := s.db.WithContext(ctx).Where("account_id = ? AND id = ?", accountID, assistantID).First(&assistant).Error; err != nil {
		return nil, fmt.Errorf("assistant not found: %w", err)
	}

	history, err := s.collectConversationMessages(ctx, accountID, conversation.ID)
	if err != nil {
		return nil, err
	}
	response, err := s.generateConversationResponse(ctx, accountID, &conversation, &assistant, history)
	if err != nil {
		return nil, err
	}
	if response == nil {
		return nil, fmt.Errorf("empty captain conversation response")
	}
	if response.Action == "handoff" || response.Content == "conversation_handoff" || response.HandoffToolCalled {
		return s.createCaptainHandoffMessage(ctx, &conversation, &assistant)
	}
	if strings.TrimSpace(response.Content) == "" {
		return nil, fmt.Errorf("captain conversation response content is blank")
	}
	return s.createCaptainOutgoingMessage(ctx, &conversation, &assistant, response.Content, response.AgentName)
}

func (s *CaptainConversationService) collectConversationMessages(ctx context.Context, accountID, conversationID uint) ([]CaptainConversationMessage, error) {
	var messages []model.Message
	if err := s.db.WithContext(ctx).
		Where("account_id = ? AND conversation_id = ? AND private = ? AND message_type IN ?", accountID, conversationID, false, []string{string(model.MessageTypeIncoming), string(model.MessageTypeOutgoing)}).
		Order("created_at ASC, id ASC").
		Find(&messages).Error; err != nil {
		return nil, fmt.Errorf("load conversation messages: %w", err)
	}
	history := make([]CaptainConversationMessage, 0, len(messages))
	for _, msg := range messages {
		role := "assistant"
		if msg.MessageType == string(model.MessageTypeIncoming) {
			role = "user"
		}
		history = append(history, CaptainConversationMessage{Role: role, Content: msg.Content, AgentName: messageAgentName(msg.AdditionalAttributes)})
	}
	return history, nil
}

func (s *CaptainConversationService) generateConversationResponse(ctx context.Context, accountID uint, conversation *model.Conversation, assistant *model.CaptainAssistant, history []CaptainConversationMessage) (*CaptainConversationResponse, error) {
	if s.backend != nil {
		return s.backend.GenerateCaptainConversationResponse(ctx, CaptainConversationResponseRequest{AccountID: accountID, Conversation: conversation, Assistant: assistant, Messages: history})
	}
	if s.llmProvider == nil {
		return nil, fmt.Errorf("captain conversation response generation disabled")
	}

	// Build system prompt from assistant config (not hardcoded)
	cfg, _ := assistant.GetConfig()
	systemPrompt := fmt.Sprintf("You are %s, a customer support assistant.", assistant.Name)
	if cfg.ProductName != "" {
		systemPrompt += fmt.Sprintf(" You represent the product: %s.", cfg.ProductName)
	}
	if cfg.Instructions != "" {
		systemPrompt += "\nInstructions: " + cfg.Instructions
	}
	systemPrompt += "\nReply to the latest customer message. Return only the reply text."

	messages := []llm.ChatMessage{{Role: "system", Content: systemPrompt}}
	for _, msg := range history {
		if strings.TrimSpace(msg.Content) == "" {
			continue
		}
		messages = append(messages, llm.ChatMessage{Role: msg.Role, Content: msg.Content})
	}

	modelName := cfg.Model
	if modelName == "" {
		modelName = "gpt-4"
	}
	temperature := cfg.Temperature
	if temperature == 0 {
		temperature = 0.7
	}

	// If tool execution service is available, run the full tool_call loop
	if s.toolExecSvc != nil {
		content, err := s.toolExecSvc.RunToolCallLoop(ctx, accountID, messages, modelName, temperature, 1024, 5)
		if err != nil {
			applogger.L().Warnf("Tool call loop failed, falling back to plain LLM: %v", err)
			// Fall through to plain LLM call below
		} else if strings.TrimSpace(content) != "" {
			return &CaptainConversationResponse{Content: content}, nil
		}
	}

	resp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{Model: modelName, Messages: messages, Temperature: temperature, MaxTokens: 1024})
	if err != nil {
		return nil, fmt.Errorf("generate captain conversation response: %w", err)
	}
	if len(resp.Choices) == 0 || strings.TrimSpace(resp.Choices[0].Message.Content) == "" {
		return nil, fmt.Errorf("empty captain conversation response")
	}
	return &CaptainConversationResponse{Content: resp.Choices[0].Message.Content}, nil
}

func (s *CaptainConversationService) createCaptainOutgoingMessage(ctx context.Context, conversation *model.Conversation, assistant *model.CaptainAssistant, content, agentName string) (*model.Message, error) {
	senderID := assistant.ID
	attrs := datatypes.JSON([]byte(`{}`))
	if strings.TrimSpace(agentName) != "" {
		raw, _ := json.Marshal(map[string]any{"agent_name": strings.TrimSpace(agentName)})
		attrs = datatypes.JSON(raw)
	}
	message := &model.Message{
		AccountID:            conversation.AccountID,
		ConversationID:       conversation.ID,
		InboxID:              conversation.InboxID,
		SenderID:             &senderID,
		SenderType:           "CaptainAssistant",
		Content:              content,
		ContentType:          string(model.MessageContentTypeText),
		MessageType:          string(model.MessageTypeOutgoing),
		Status:               string(model.MessageStatusSent),
		AdditionalAttributes: attrs,
	}
	if err := s.db.WithContext(ctx).Create(message).Error; err != nil {
		return nil, fmt.Errorf("create captain outgoing message: %w", err)
	}
	if s.worker != nil {
		if _, err := EnqueueSendReply(ctx, s.worker, message.ID); err != nil {
			return message, err
		}
	}
	return message, nil
}

func (s *CaptainConversationService) createCaptainHandoffMessage(ctx context.Context, conversation *model.Conversation, assistant *model.CaptainAssistant) (*model.Message, error) {
	content := assistantConfigString(assistant.Config, "handoff_message")
	if content == "" {
		content = "Transferring to another agent for further assistance."
	}
	message, err := s.createCaptainOutgoingMessage(ctx, conversation, assistant, content, "")
	if err != nil {
		return nil, err
	}
	if err := s.db.WithContext(ctx).Model(&model.Conversation{}).Where("account_id = ? AND id = ? AND status = ?", conversation.AccountID, conversation.ID, string(model.ConversationStatusPending)).Update("status", string(model.ConversationStatusOpen)).Error; err != nil {
		return message, fmt.Errorf("handoff captain conversation: %w", err)
	}
	return message, nil
}

func messageAgentName(raw datatypes.JSON) string {
	if len(raw) == 0 {
		return ""
	}
	var attrs map[string]any
	if err := json.Unmarshal(raw, &attrs); err != nil {
		return ""
	}
	if value, ok := attrs["agent_name"].(string); ok {
		return value
	}
	return ""
}

func assistantConfigString(raw []byte, key string) string {
	if len(raw) == 0 {
		return ""
	}
	var cfg map[string]any
	if err := json.Unmarshal(raw, &cfg); err != nil {
		return ""
	}
	if value, ok := cfg[key].(string); ok {
		return strings.TrimSpace(value)
	}
	return ""
}

type captainConversationResponseBuilderJob struct {
	AccountID      uint `json:"account_id"`
	ConversationID uint `json:"conversation_id"`
	AssistantID    uint `json:"assistant_id"`
	MessageID      uint `json:"message_id,omitempty"`
}

var captainConversationRegistrations sync.Map

func RegisterCaptainConversationJobs(wp *worker.WorkerPool, svc *CaptainConversationService) {
	if wp == nil || svc == nil {
		return
	}
	if _, loaded := captainConversationRegistrations.LoadOrStore(wp, struct{}{}); loaded {
		return
	}
	wp.Register(TaskTypeCaptainConversationResponseBuilder, svc.performConversationResponseBuilderJob)
}

func EnqueueCaptainConversationResponseForMessage(ctx context.Context, wp *worker.WorkerPool, db *gorm.DB, messageID uint) (*model.BackgroundJob, error) {
	if wp == nil || db == nil || messageID == 0 || !db.Migrator().HasTable(&model.CaptainInbox{}) {
		return nil, nil
	}
	var message model.Message
	if err := db.WithContext(ctx).First(&message, messageID).Error; err != nil {
		return nil, err
	}
	if message.MessageType != string(model.MessageTypeIncoming) || message.Private {
		return nil, nil
	}
	var conversation model.Conversation
	if err := db.WithContext(ctx).Where("account_id = ? AND id = ?", message.AccountID, message.ConversationID).First(&conversation).Error; err != nil {
		return nil, err
	}
	if conversation.Status != string(model.ConversationStatusPending) {
		return nil, nil
	}
	var ci model.CaptainInbox
	if err := db.WithContext(ctx).Where("account_id = ? AND inbox_id = ?", conversation.AccountID, conversation.InboxID).First(&ci).Error; err != nil {
		if err == gorm.ErrRecordNotFound {
			return nil, nil
		}
		return nil, err
	}
	options := []worker.EnqueueOption{
		worker.WithMaxAttempts(3),
		worker.WithIdempotencyKey(fmt.Sprintf("captain:conversation_response_builder:message:%d", message.ID)),
	}
	if scheduledAt := captainConversationResponseScheduledAt(ctx, db, message.ID); !scheduledAt.IsZero() {
		options = append(options, worker.WithScheduledAt(scheduledAt))
	}
	return wp.Enqueue(ctx, TaskTypeCaptainConversationResponseBuilder, captainConversationResponseBuilderJob{AccountID: conversation.AccountID, ConversationID: conversation.ID, AssistantID: ci.AssistantID, MessageID: message.ID}, options...)
}

func captainConversationResponseScheduledAt(ctx context.Context, db *gorm.DB, messageID uint) time.Time {
	if !db.Migrator().HasTable(&model.Attachment{}) {
		return time.Time{}
	}
	var count int64
	if err := db.WithContext(ctx).Model(&model.Attachment{}).Where("message_id = ?", messageID).Count(&count).Error; err != nil || count == 0 {
		return time.Time{}
	}
	wait := time.Second + time.Duration(minInt64(count, 4))*time.Second
	return time.Now().Add(wait)
}

func minInt64(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

func (s *CaptainConversationService) performConversationResponseBuilderJob(ctx context.Context, job *model.BackgroundJob) error {
	var payload captainConversationResponseBuilderJob
	if err := json.Unmarshal(job.Payload, &payload); err != nil {
		return fmt.Errorf("unmarshal captain conversation response builder job: %w", err)
	}
	if payload.AccountID == 0 || payload.ConversationID == 0 || payload.AssistantID == 0 {
		return fmt.Errorf("invalid captain conversation response builder job payload: %#v", payload)
	}
	_, err := s.BuildConversationResponseByAccount(ctx, payload.AccountID, payload.ConversationID, payload.AssistantID)
	return err
}