ai-agent/internal/services/conversation_queue_service.go

package services

import (
	"math"
	"slices"
	"sync"
	"time"

	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"

	"github.com/mlogclub/simple/sqls"
)

var ConversationQueueService = newConversationQueueService()

const (
	queueEscalationInterval = 5 * time.Minute
	queueEscalationMaxLevel = 6
	queueAverageHandleTime  = 8 * time.Minute
)

type ConversationQueueSnapshot struct {
	Queued               bool
	EnteredAt            *time.Time
	Position             int
	AheadCount           int
	WaitingCount         int
	WaitSeconds          int64
	EstimatedWaitSeconds int64
	EscalationLevel      int
	EffectivePriority    int
	ServiceOnline        bool
}

type queueSnapshotCacheEntry struct {
	expiresAt time.Time
	snapshots map[int64]ConversationQueueSnapshot
}

type conversationQueueService struct {
	mu    sync.Mutex
	cache map[int64]queueSnapshotCacheEntry
}

func newConversationQueueService() *conversationQueueService {
	return &conversationQueueService{cache: make(map[int64]queueSnapshotCacheEntry)}
}

func (s *conversationQueueService) GetSnapshot(conversation *models.Conversation) ConversationQueueSnapshot {
	if conversation == nil || conversation.Status != enums.IMConversationStatusPending || conversation.CurrentAssigneeID > 0 {
		return ConversationQueueSnapshot{}
	}
	now := time.Now()
	s.mu.Lock()
	entry, found := s.cache[conversation.CurrentTeamID]
	s.mu.Unlock()
	if found && now.Before(entry.expiresAt) {
		return entry.snapshots[conversation.ID]
	}
	snapshots := s.buildPoolSnapshotsAt(conversation.CurrentTeamID, now)
	s.mu.Lock()
	s.cache[conversation.CurrentTeamID] = queueSnapshotCacheEntry{
		expiresAt: now.Add(time.Second),
		snapshots: snapshots,
	}
	s.mu.Unlock()
	return snapshots[conversation.ID]
}

func (s *conversationQueueService) GetSnapshotAt(conversation *models.Conversation, now time.Time) ConversationQueueSnapshot {
	if conversation == nil || conversation.Status != enums.IMConversationStatusPending || conversation.CurrentAssigneeID > 0 {
		return ConversationQueueSnapshot{}
	}
	return s.buildPoolSnapshotsAt(conversation.CurrentTeamID, now)[conversation.ID]
}

func (s *conversationQueueService) buildPoolSnapshotsAt(teamID int64, now time.Time) map[int64]ConversationQueueSnapshot {
	queue := s.findPoolQueue(teamID)
	s.Sort(queue, now)
	capacity, freeSlots := s.poolCapacity(teamID, now)
	snapshots := make(map[int64]ConversationQueueSnapshot, len(queue))
	for index := range queue {
		conversation := &queue[index]
		snapshot := ConversationQueueSnapshot{
			Queued:            true,
			EnteredAt:         queueEnteredAt(conversation),
			Position:          index + 1,
			AheadCount:        index,
			WaitingCount:      len(queue),
			EffectivePriority: s.EffectivePriority(conversation, now),
			EscalationLevel:   s.EscalationLevel(conversation, now),
			ServiceOnline:     capacity > 0,
		}
		if snapshot.EnteredAt != nil && now.After(*snapshot.EnteredAt) {
			snapshot.WaitSeconds = int64(now.Sub(*snapshot.EnteredAt) / time.Second)
		}
		if capacity > 0 {
			remainingBeforeService := snapshot.AheadCount - freeSlots
			if remainingBeforeService >= 0 {
				waves := int64(math.Ceil(float64(remainingBeforeService+1) / float64(capacity)))
				snapshot.EstimatedWaitSeconds = waves * int64(queueAverageHandleTime/time.Second)
			}
		}
		snapshots[conversation.ID] = snapshot
	}
	return snapshots
}

func (s *conversationQueueService) Sort(conversations []models.Conversation, now time.Time) {
	slices.SortFunc(conversations, func(a, b models.Conversation) int {
		aPriority := s.EffectivePriority(&a, now)
		bPriority := s.EffectivePriority(&b, now)
		switch {
		case aPriority > bPriority:
			return -1
		case aPriority < bPriority:
			return 1
		}

		aEnteredAt := queueEnteredAtValue(&a)
		bEnteredAt := queueEnteredAtValue(&b)
		switch {
		case aEnteredAt.Before(bEnteredAt):
			return -1
		case aEnteredAt.After(bEnteredAt):
			return 1
		case a.ID < b.ID:
			return -1
		case a.ID > b.ID:
			return 1
		default:
			return 0
		}
	})
}

func (s *conversationQueueService) EffectivePriority(conversation *models.Conversation, now time.Time) int {
	if conversation == nil {
		return 0
	}
	return conversation.Priority + s.EscalationLevel(conversation, now)
}

func (s *conversationQueueService) EscalationLevel(conversation *models.Conversation, now time.Time) int {
	enteredAt := queueEnteredAt(conversation)
	if enteredAt == nil || !now.After(*enteredAt) {
		return 0
	}
	level := int(now.Sub(*enteredAt) / queueEscalationInterval)
	if level > queueEscalationMaxLevel {
		return queueEscalationMaxLevel
	}
	return level
}

func (s *conversationQueueService) PublishPoolUpdates(teamID int64) {
	s.mu.Lock()
	delete(s.cache, teamID)
	s.mu.Unlock()
	queue := s.findPoolQueue(teamID)
	for index := range queue {
		conversation := queue[index]
		WsService.PublishConversationChanged(&conversation, enums.IMRealtimeEventConversationQueueUpdated)
	}
}

func (s *conversationQueueService) findPoolQueue(teamID int64) []models.Conversation {
	return ConversationService.Find(sqls.NewCnd().
		Eq("status", enums.IMConversationStatusPending).
		Eq("current_assignee_id", 0).
		Eq("current_team_id", teamID).
		Asc("id"))
}

func (s *conversationQueueService) poolCapacity(teamID int64, now time.Time) (int, int) {
	if !sqls.DB().Migrator().HasTable(&models.AgentTeam{}) ||
		!sqls.DB().Migrator().HasTable(&models.AgentTeamSchedule{}) ||
		!sqls.DB().Migrator().HasTable(&models.AgentProfile{}) {
		return 0, 0
	}
	teamIDs := []int64{teamID}
	if teamID <= 0 {
		teamIDs = ConversationDispatchService.findAllActiveScheduleTeamIDs(now)
	}
	activeTeamIDs := ConversationDispatchService.findActiveScheduleTeamIDs(teamIDs, now)
	if len(activeTeamIDs) == 0 {
		return 0, 0
	}
	profiles := AgentProfileService.GetDispatchAgents(activeTeamIDs)
	profiles, userIDs, _ := ConversationDispatchService.filterEnabledDispatchProfiles(profiles)
	if len(profiles) == 0 {
		return 0, 0
	}
	activeCounts, err := ConversationDispatchService.findActiveConversationCountMap(userIDs)
	if err != nil {
		return 0, 0
	}

	totalCapacity := 0
	freeSlots := 0
	for _, profile := range profiles {
		capacity := profile.MaxConcurrentCount
		if capacity <= 0 {
			capacity = 1
		}
		totalCapacity += capacity
		available := capacity - activeCounts[profile.UserID]
		if available > 0 {
			freeSlots += available
		}
	}
	return totalCapacity, freeSlots
}

func queueEnteredAt(conversation *models.Conversation) *time.Time {
	if conversation == nil {
		return nil
	}
	if conversation.QueueEnteredAt != nil {
		return conversation.QueueEnteredAt
	}
	if conversation.HandoffAt != nil {
		return conversation.HandoffAt
	}
	if !conversation.CreatedAt.IsZero() {
		return &conversation.CreatedAt
	}
	return nil
}

func queueEnteredAtValue(conversation *models.Conversation) time.Time {
	if enteredAt := queueEnteredAt(conversation); enteredAt != nil {
		return *enteredAt
	}
	return time.Time{}
}

func queueEnteredAtForTransition(conversation *models.Conversation, now time.Time) time.Time {
	if conversation != nil && conversation.Status == enums.IMConversationStatusPending && conversation.QueueEnteredAt != nil {
		return *conversation.QueueEnteredAt
	}
	return now
}