gochat/internal/service/captain_task_extended_service.go

package service

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

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
	applogger "github.com/gochat/gochat/pkg/logger"
	pkgvalidator "github.com/gochat/gochat/pkg/validator"
)

// CaptainTaskExtendedService provides label suggestion and follow-up features.
// These are GET-based endpoints that return AI-generated suggestions without modifying data.
// Reference: Chatwoot Captain::Llm::ConversationInsightService + M12 PRD
type CaptainTaskExtendedService struct {
	conversationRepo *repository.ConversationRepo
	messageRepo      *repository.MessageRepo
	assistantRepo    *repository.CaptainAssistantRepo
	preferenceRepo   *repository.CaptainPreferenceRepo
	suggestionRepo   *repository.CopilotSuggestionRepo
	llmProvider      llm.Provider
}

func NewCaptainTaskExtendedService(
	conversationRepo *repository.ConversationRepo,
	messageRepo *repository.MessageRepo,
	assistantRepo *repository.CaptainAssistantRepo,
	preferenceRepo *repository.CaptainPreferenceRepo,
	llmProvider llm.Provider,
	suggestionRepos ...*repository.CopilotSuggestionRepo,
) *CaptainTaskExtendedService {
	var suggestionRepo *repository.CopilotSuggestionRepo
	if len(suggestionRepos) > 0 {
		suggestionRepo = suggestionRepos[0]
	}
	return &CaptainTaskExtendedService{
		conversationRepo: conversationRepo,
		messageRepo:      messageRepo,
		assistantRepo:    assistantRepo,
		preferenceRepo:   preferenceRepo,
		suggestionRepo:   suggestionRepo,
		llmProvider:      llmProvider,
	}
}

// --- Request/Response DTOs ---

// LabelSuggestionQuery is the input for label suggestion on multiple conversations.
type LabelSuggestionQuery struct {
	ConversationIDs []uint `json:"conversation_ids" validate:"required,min=1"`
	AssistantID     uint   `json:"assistant_id,omitempty"`
}

// ConversationLabelSuggestion holds label suggestions for a single conversation.
type ConversationLabelSuggestion struct {
	ConversationID uint     `json:"conversation_id"`
	Labels         []string `json:"labels"`
	Priority       string   `json:"priority"` // urgent, high, medium, low
	Reason         string   `json:"reason"`
}

// LabelSuggestionResult holds label suggestions for multiple conversations.
type LabelSuggestionResult struct {
	Suggestions []ConversationLabelSuggestion `json:"suggestions"`
}

// FollowUpQuery is the input for follow-up task suggestions.
type FollowUpQuery struct {
	ConversationIDs []uint `json:"conversation_ids" validate:"required,min=1"`
	AssistantID     uint   `json:"assistant_id,omitempty"`
}

// FollowUpTask represents a suggested follow-up action.
type FollowUpTask struct {
	ConversationID uint   `json:"conversation_id"`
	Title          string `json:"title"`
	Description    string `json:"description"`
	DueDateHint    string `json:"due_date_hint,omitempty"` // e.g. "within 24 hours"
	Priority       string `json:"priority"`                // urgent, high, medium, low
	AssigneeHint   string `json:"assignee_hint,omitempty"` // suggested assignee name
}

// FollowUpResult holds follow-up suggestions for multiple conversations.
type FollowUpResult struct {
	FollowUps []FollowUpTask `json:"follow_ups"`
}

type ChatwootLabelSuggestionRequest struct {
	ConversationDisplayID uint `json:"conversation_display_id"`
}

type ChatwootFollowUpRequest struct {
	FollowUpContext       map[string]interface{} `json:"follow_up_context"`
	Message               string                 `json:"message"`
	ConversationDisplayID uint                   `json:"conversation_display_id"`
}

type ChatwootTaskResult struct {
	Message         *string                `json:"message"`
	FollowUpContext map[string]interface{} `json:"follow_up_context,omitempty"`
}

// --- Helper: fetch conversation context for LLM ---

func (s *CaptainTaskExtendedService) fetchConversationMessages(ctx context.Context, conversationID uint) (string, error) {
	messages, _, err := s.messageRepo.FindByConversation(ctx, conversationID, 0, 50)
	if err != nil {
		return "", fmt.Errorf("fetch messages for conversation %d: %w", conversationID, err)
	}
	if len(messages) == 0 {
		return "", nil
	}

	var parts []string
	for _, m := range messages {
		role := "contact"
		if m.SenderType == "user" || m.SenderType == "agent_bot" {
			role = "agent"
		}
		parts = append(parts, fmt.Sprintf("[%s]: %s", role, m.Content))
	}
	return strings.Join(parts, "\n"), nil
}

func (s *CaptainTaskExtendedService) getPreferencePromptSuffix(ctx context.Context, accountID uint) string {
	pref, err := s.preferenceRepo.GetByAccountID(ctx, accountID)
	if err != nil || pref == nil {
		return ""
	}
	return pref.CustomPromptSuffix
}

func (s *CaptainTaskExtendedService) buildAssistantContext(ctx context.Context, assistantID uint) string {
	if assistantID == 0 {
		return ""
	}
	assistant, err := s.assistantRepo.GetByID(ctx, assistantID)
	if err != nil || assistant == nil {
		return ""
	}
	cfg, _ := assistant.GetConfig()
	if cfg != nil && cfg.Instructions != "" {
		return "\nAssistant instructions: " + cfg.Instructions
	}
	return ""
}

func (s *CaptainTaskExtendedService) LabelSuggestion(ctx context.Context, accountID uint, req *ChatwootLabelSuggestionRequest) (*ChatwootTaskResult, error) {
	if s.llmProvider == nil {
		return nil, taskError(422, "Captain is disabled")
	}
	conversation, err := s.conversationRepo.FindByAccountAndDisplayIDOrID(ctx, accountID, req.ConversationDisplayID)
	if err != nil {
		return nil, taskError(422, "Conversation not found")
	}
	msgs, err := s.fetchConversationMessages(ctx, conversation.ID)
	if err != nil {
		return nil, err
	}
	if strings.TrimSpace(msgs) == "" {
		return &ChatwootTaskResult{Message: nil}, nil
	}

	llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
		Messages: []llm.ChatMessage{
			{Role: "system", Content: `Suggest concise labels for this customer support conversation. Return only a comma-separated label list.`},
			{Role: "user", Content: msgs},
		},
		Temperature: 0.2,
		MaxTokens:   256,
	})
	if err != nil {
		return nil, err
	}
	if llmResp == nil || len(llmResp.Choices) == 0 {
		return &ChatwootTaskResult{Message: nil}, nil
	}
	message := strings.TrimSpace(llmResp.Choices[0].Message.Content)
	if message == "" {
		return &ChatwootTaskResult{Message: nil}, nil
	}
	s.persistChatwootTaskSuggestion(ctx, accountID, conversation.ID, message, model.CopilotSuggestionTypeSuggestion)
	return &ChatwootTaskResult{
		Message:         &message,
		FollowUpContext: buildTaskFollowUpContext("label_suggestion", msgs, message, conversation.ChannelType),
	}, nil
}

func (s *CaptainTaskExtendedService) FollowUp(ctx context.Context, accountID uint, req *ChatwootFollowUpRequest) (*ChatwootTaskResult, error) {
	if s.llmProvider == nil {
		return nil, taskError(422, "Captain is disabled")
	}
	if !validTaskFollowUpContext(req.FollowUpContext) {
		return nil, taskError(422, "Follow-up context missing")
	}
	conversation, err := s.conversationRepo.FindByAccountAndDisplayIDOrID(ctx, accountID, req.ConversationDisplayID)
	if err != nil {
		return nil, taskError(422, "Conversation not found")
	}
	messages := []llm.ChatMessage{
		{Role: "system", Content: "You just performed a customer support assistant task. Refine the previous result based on the agent's feedback. Output only the refined reply."},
		{Role: "user", Content: fmt.Sprint(req.FollowUpContext["original_context"])},
		{Role: "assistant", Content: fmt.Sprint(req.FollowUpContext["last_response"])},
	}
	for _, historyMessage := range followUpHistory(req.FollowUpContext) {
		messages = append(messages, historyMessage)
	}
	messages = append(messages, llm.ChatMessage{Role: "user", Content: req.Message})

	llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
		Messages:    messages,
		Temperature: 0.3,
		MaxTokens:   512,
	})
	if err != nil {
		return nil, err
	}
	if llmResp == nil || len(llmResp.Choices) == 0 {
		return nil, taskError(422, "No response from LLM")
	}
	message := llmResp.Choices[0].Message.Content
	updatedContext := updateTaskFollowUpContext(req.FollowUpContext, req.Message, message, conversation.ChannelType)
	s.persistChatwootTaskSuggestion(ctx, accountID, conversation.ID, message, model.CopilotSuggestionTypeSuggestion)
	return &ChatwootTaskResult{Message: &message, FollowUpContext: updatedContext}, nil
}

// --- Label Suggestion (GET) ---
// Reference: Chatwoot Captain::ConversationInsightController#suggest_labels

func (s *CaptainTaskExtendedService) SuggestLabels(ctx context.Context, accountID uint, query *LabelSuggestionQuery) (*LabelSuggestionResult, error) {
	// Validate request
	if err := pkgvalidator.ValidateStruct(query); err != nil {
		return nil, fmt.Errorf("validation error: %w", err)
	}

	// Verify all conversation IDs belong to this account
	for _, convID := range query.ConversationIDs {
		conv, err := s.conversationRepo.FindByID(ctx, convID)
		if err != nil || conv == nil {
			return nil, fmt.Errorf("conversation %d not found", convID)
		}
		if conv.AccountID != accountID {
			return nil, fmt.Errorf("conversation %d does not belong to account %d", convID, accountID)
		}
	}

	promptSuffix := s.getPreferencePromptSuffix(ctx, accountID)
	assistantCtx := s.buildAssistantContext(ctx, query.AssistantID)

	var suggestions []ConversationLabelSuggestion

	for _, convID := range query.ConversationIDs {
		msgs, err := s.fetchConversationMessages(ctx, convID)
		if err != nil {
			applogger.L().Warnf("Skip label suggestion for conversation %d: %v", convID, err)
			continue
		}
		if msgs == "" {
			continue
		}

		systemPrompt := `You are a conversation categorization AI. Analyze the conversation and suggest:
1. Appropriate labels/tags (3-5 concise labels that categorize the conversation topic, type, and urgency)
2. Priority level: urgent, high, medium, or low
3. A brief reason explaining why these labels and priority were chosen

Respond in JSON format:
{
  "labels": ["label1", "label2", ...],
  "priority": "medium",
  "reason": "brief explanation"
}` + assistantCtx + promptSuffix

		llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
			Model: "",
			Messages: []llm.ChatMessage{
				{Role: "system", Content: systemPrompt},
				{Role: "user", Content: msgs},
			},
			Temperature: 0.2,
			MaxTokens:   512,
		})
		if err != nil {
			applogger.L().Warnf("LLM label suggestion failed for conversation %d: %v", convID, err)
			continue
		}

		if len(llmResp.Choices) == 0 {
			continue
		}

		content := llmResp.Choices[0].Message.Content
		var parsed struct {
			Labels   []string `json:"labels"`
			Priority string   `json:"priority"`
			Reason   string   `json:"reason"`
		}
		if err := json.Unmarshal([]byte(content), &parsed); err != nil {
			// Fallback: try to extract from plain text
			parsed.Labels = []string{"uncategorized"}
			parsed.Priority = "medium"
			parsed.Reason = content
		}

		suggestions = append(suggestions, ConversationLabelSuggestion{
			ConversationID: convID,
			Labels:         parsed.Labels,
			Priority:       parsed.Priority,
			Reason:         parsed.Reason,
		})
	}

	return &LabelSuggestionResult{Suggestions: suggestions}, nil
}

// --- Follow Up (GET) ---
// Reference: Chatwoot Captain::ConversationInsightController + M12 PRD §Follow Up Suggestions

func (s *CaptainTaskExtendedService) SuggestFollowUp(ctx context.Context, accountID uint, query *FollowUpQuery) (*FollowUpResult, error) {
	// Validate request
	if err := pkgvalidator.ValidateStruct(query); err != nil {
		return nil, fmt.Errorf("validation error: %w", err)
	}

	// Verify all conversation IDs belong to this account
	for _, convID := range query.ConversationIDs {
		conv, err := s.conversationRepo.FindByID(ctx, convID)
		if err != nil || conv == nil {
			return nil, fmt.Errorf("conversation %d not found", convID)
		}
		if conv.AccountID != accountID {
			return nil, fmt.Errorf("conversation %d does not belong to account %d", convID, accountID)
		}
	}

	promptSuffix := s.getPreferencePromptSuffix(ctx, accountID)
	assistantCtx := s.buildAssistantContext(ctx, query.AssistantID)

	var followUps []FollowUpTask

	for _, convID := range query.ConversationIDs {
		msgs, err := s.fetchConversationMessages(ctx, convID)
		if err != nil {
			applogger.L().Warnf("Skip follow-up suggestion for conversation %d: %v", convID, err)
			continue
		}
		if msgs == "" {
			continue
		}

		conversation, err := s.conversationRepo.FindByID(ctx, convID)
		if err != nil {
			applogger.L().Warnf("Skip follow-up for conversation %d: %v", convID, err)
			continue
		}

		statusHint := ""
		if conversation != nil {
			statusHint = fmt.Sprintf("\nConversation status: %s", conversation.Status)
		}

		systemPrompt := `You are a customer support follow-up AI. Analyze the conversation and suggest:
1. A follow-up task title (concise, actionable)
2. A description of what needs to be done next
3. A due date hint (e.g., "within 24 hours", "within 3 business days")
4. Priority level: urgent, high, medium, or low
5. If applicable, suggest who should handle this follow-up

Respond in JSON format:
{
  "title": "Follow up with customer about shipping",
  "description": "Customer asked about delivery status, need to check tracking number",
  "due_date_hint": "within 24 hours",
  "priority": "high",
  "assignee_hint": "shipping team"
}` + statusHint + assistantCtx + promptSuffix

		llmResp, err := s.llmProvider.ChatCompletion(ctx, llm.ChatRequest{
			Model: "",
			Messages: []llm.ChatMessage{
				{Role: "system", Content: systemPrompt},
				{Role: "user", Content: msgs},
			},
			Temperature: 0.2,
			MaxTokens:   512,
		})
		if err != nil {
			applogger.L().Warnf("LLM follow-up suggestion failed for conversation %d: %v", convID, err)
			continue
		}

		if len(llmResp.Choices) == 0 {
			continue
		}

		content := llmResp.Choices[0].Message.Content
		var parsed struct {
			Title        string `json:"title"`
			Description  string `json:"description"`
			DueDateHint  string `json:"due_date_hint"`
			Priority     string `json:"priority"`
			AssigneeHint string `json:"assignee_hint"`
		}
		if err := json.Unmarshal([]byte(content), &parsed); err != nil {
			parsed.Title = "Review conversation"
			parsed.Description = content
			parsed.Priority = "medium"
		}

		followUps = append(followUps, FollowUpTask{
			ConversationID: convID,
			Title:          parsed.Title,
			Description:    parsed.Description,
			DueDateHint:    parsed.DueDateHint,
			Priority:       parsed.Priority,
			AssigneeHint:   parsed.AssigneeHint,
		})
	}

	return &FollowUpResult{FollowUps: followUps}, nil
}

func buildTaskFollowUpContext(eventName, originalContext, lastResponse, channelType string) map[string]interface{} {
	ctx := map[string]interface{}{
		"event_name":           eventName,
		"original_context":     originalContext,
		"last_response":        lastResponse,
		"conversation_history": []map[string]string{},
	}
	if channelType != "" {
		ctx["channel_type"] = channelType
	}
	return ctx
}

func validTaskFollowUpContext(ctx map[string]interface{}) bool {
	if ctx == nil {
		return false
	}
	eventName, _ := ctx["event_name"].(string)
	if !allowedTaskFollowUpEvent(eventName) {
		return false
	}
	return strings.TrimSpace(fmt.Sprint(ctx["original_context"])) != "" && strings.TrimSpace(fmt.Sprint(ctx["last_response"])) != ""
}

func allowedTaskFollowUpEvent(eventName string) bool {
	switch eventName {
	case "professional", "casual", "friendly", "confident", "straightforward", "fix_spelling_grammar", "improve", "summarize", "reply_suggestion", "label_suggestion":
		return true
	default:
		return false
	}
}

func followUpHistory(ctx map[string]interface{}) []llm.ChatMessage {
	raw, ok := ctx["conversation_history"].([]interface{})
	if !ok {
		return nil
	}
	messages := make([]llm.ChatMessage, 0, len(raw))
	for _, item := range raw {
		m, ok := item.(map[string]interface{})
		if !ok {
			continue
		}
		role := fmt.Sprint(m["role"])
		content := fmt.Sprint(m["content"])
		if role == "" || content == "" {
			continue
		}
		messages = append(messages, llm.ChatMessage{Role: role, Content: content})
	}
	return messages
}

func updateTaskFollowUpContext(ctx map[string]interface{}, userMessage, assistantMessage, channelType string) map[string]interface{} {
	history := make([]map[string]string, 0)
	if raw, ok := ctx["conversation_history"].([]interface{}); ok {
		for _, item := range raw {
			m, ok := item.(map[string]interface{})
			if !ok {
				continue
			}
			history = append(history, map[string]string{"role": fmt.Sprint(m["role"]), "content": fmt.Sprint(m["content"])})
		}
	}
	history = append(history,
		map[string]string{"role": "user", "content": userMessage},
		map[string]string{"role": "assistant", "content": assistantMessage},
	)
	updated := map[string]interface{}{
		"event_name":           ctx["event_name"],
		"original_context":     ctx["original_context"],
		"last_response":        assistantMessage,
		"conversation_history": history,
	}
	if channelType != "" {
		updated["channel_type"] = channelType
	} else if value, ok := ctx["channel_type"]; ok {
		updated["channel_type"] = value
	}
	return updated
}

func (s *CaptainTaskExtendedService) persistChatwootTaskSuggestion(ctx context.Context, accountID, conversationID uint, content string, suggestionType model.CopilotSuggestionType) {
	if s.suggestionRepo == nil || strings.TrimSpace(content) == "" {
		return
	}
	msg := &model.CopilotSuggestionMessage{
		AccountID:      accountID,
		ConversationID: conversationID,
		Content:        content,
		SuggestionType: suggestionType,
		Status:         model.CopilotSuggestionStatusPending,
	}
	if err := s.suggestionRepo.Create(ctx, msg); err != nil {
		applogger.L().Warnf("Captain extended task suggestion persistence failed: %v", err)
	}
}