internal/queue/priority.go

// Package queue provides priority-based request queuing.
package queue

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/sirupsen/logrus"
)

// Priority represents request priority level.
type Priority int

const (
	PriorityCritical Priority = 0 // User typing, streaming
	PriorityHigh     Priority = 1 // Interactive commands
	PriorityNormal   Priority = 2 // Standard requests
	PriorityLow      Priority = 3 // Background analysis
	PriorityBatch    Priority = 4 // Bulk operations
)

// String returns the priority name.
func (p Priority) String() string {
	switch p {
	case PriorityCritical:
		return "critical"
	case PriorityHigh:
		return "high"
	case PriorityNormal:
		return "normal"
	case PriorityLow:
		return "low"
	case PriorityBatch:
		return "batch"
	default:
		return "unknown"
	}
}

// Config holds priority queue configuration.
type Config struct {
	Enabled        bool `yaml:"enabled" json:"enabled"`
	Workers        int  `yaml:"workers" json:"workers"`
	CriticalBuffer int  `yaml:"critical-buffer" json:"critical_buffer"`
	HighBuffer     int  `yaml:"high-buffer" json:"high_buffer"`
	NormalBuffer   int  `yaml:"normal-buffer" json:"normal_buffer"`
	LowBuffer      int  `yaml:"low-buffer" json:"low_buffer"`
	BatchBuffer    int  `yaml:"batch-buffer" json:"batch_buffer"`
}

// DefaultConfig returns default priority queue configuration.
func DefaultConfig() *Config {
	return &Config{
		Enabled:        true,
		Workers:        10,
		CriticalBuffer: 100,
		HighBuffer:     200,
		NormalBuffer:   500,
		LowBuffer:      300,
		BatchBuffer:    1000,
	}
}

// Request represents a queued request.
type Request struct {
	ID         string
	Priority   Priority
	Context    context.Context
	Handler    func() (interface{}, error)
	ResultCh   chan *Result
	EnqueuedAt time.Time
}

// Result represents the result of a processed request.
type Result struct {
	Value    interface{}
	Error    error
	WaitTime time.Duration
}

// PriorityQueue manages priority-based request processing.
type PriorityQueue struct {
	queues  [5]chan *Request
	workers int
	stopCh  chan struct{}
	wg      sync.WaitGroup
	config  *Config
	mu      sync.RWMutex
	closed  bool

	// Stats
	processed [5]int64
	totalWait [5]int64
}

// NewPriorityQueue creates a new priority queue.
func NewPriorityQueue(config *Config) *PriorityQueue {
	if config == nil {
		config = DefaultConfig()
	}

	if !config.Enabled {
		return nil
	}

	if config.Workers <= 0 {
		config.Workers = 10
	}

	pq := &PriorityQueue{
		workers: config.Workers,
		stopCh:  make(chan struct{}),
		config:  config,
	}

	// Initialize queues with configured buffer sizes
	pq.queues[PriorityCritical] = make(chan *Request, config.CriticalBuffer)
	pq.queues[PriorityHigh] = make(chan *Request, config.HighBuffer)
	pq.queues[PriorityNormal] = make(chan *Request, config.NormalBuffer)
	pq.queues[PriorityLow] = make(chan *Request, config.LowBuffer)
	pq.queues[PriorityBatch] = make(chan *Request, config.BatchBuffer)

	// Start workers
	for i := 0; i < config.Workers; i++ {
		pq.wg.Add(1)
		go pq.worker(i)
	}

	log.WithFields(log.Fields{
		"workers":         config.Workers,
		"critical_buffer": config.CriticalBuffer,
		"normal_buffer":   config.NormalBuffer,
	}).Info("priority queue initialized")

	return pq
}

// Enqueue adds a request to the appropriate priority queue.
func (pq *PriorityQueue) Enqueue(req *Request) bool {
	if pq == nil {
		return false
	}

	pq.mu.RLock()
	if pq.closed {
		pq.mu.RUnlock()
		return false
	}
	pq.mu.RUnlock()

	req.EnqueuedAt = time.Now()
	if req.ResultCh == nil {
		req.ResultCh = make(chan *Result, 1)
	}

	select {
	case pq.queues[req.Priority] <- req:
		return true
	default:
		// Queue full, try to process immediately for critical
		if req.Priority == PriorityCritical {
			go pq.processRequest(req)
			return true
		}
		return false
	}
}

// Submit submits a request and waits for the result.
func (pq *PriorityQueue) Submit(ctx context.Context, priority Priority, handler func() (interface{}, error)) (interface{}, error) {
	if pq == nil {
		// Direct execution if queue is disabled
		return handler()
	}

	req := &Request{
		Priority: priority,
		Context:  ctx,
		Handler:  handler,
		ResultCh: make(chan *Result, 1),
	}

	if !pq.Enqueue(req) {
		// Queue full, execute directly
		return handler()
	}

	select {
	case result := <-req.ResultCh:
		return result.Value, result.Error
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (pq *PriorityQueue) worker(id int) {
	defer pq.wg.Done()

	for {
		select {
		case <-pq.stopCh:
			return
		default:
			req := pq.getNextRequest()
			if req != nil {
				pq.processRequest(req)
			}
		}
	}
}

func (pq *PriorityQueue) getNextRequest() *Request {
	// Check queues in priority order with select
	select {
	case req := <-pq.queues[PriorityCritical]:
		return req
	default:
	}

	select {
	case req := <-pq.queues[PriorityCritical]:
		return req
	case req := <-pq.queues[PriorityHigh]:
		return req
	default:
	}

	select {
	case req := <-pq.queues[PriorityCritical]:
		return req
	case req := <-pq.queues[PriorityHigh]:
		return req
	case req := <-pq.queues[PriorityNormal]:
		return req
	default:
	}

	// Wait for any request with timeout
	select {
	case req := <-pq.queues[PriorityCritical]:
		return req
	case req := <-pq.queues[PriorityHigh]:
		return req
	case req := <-pq.queues[PriorityNormal]:
		return req
	case req := <-pq.queues[PriorityLow]:
		return req
	case req := <-pq.queues[PriorityBatch]:
		return req
	case <-pq.stopCh:
		return nil
	case <-time.After(100 * time.Millisecond):
		return nil
	}
}

func (pq *PriorityQueue) processRequest(req *Request) {
	waitTime := time.Since(req.EnqueuedAt)

	// Check if context is cancelled
	if req.Context != nil {
		select {
		case <-req.Context.Done():
			req.ResultCh <- &Result{Error: req.Context.Err(), WaitTime: waitTime}
			return
		default:
		}
	}

	// Execute handler
	value, err := req.Handler()

	// Update stats
	atomic.AddInt64(&pq.processed[req.Priority], 1)
	atomic.AddInt64(&pq.totalWait[req.Priority], int64(waitTime))

	// Send result
	select {
	case req.ResultCh <- &Result{Value: value, Error: err, WaitTime: waitTime}:
	default:
	}
}

// Stats returns queue statistics.
type Stats struct {
	Processed     [5]int64
	AvgWaitTimeMs [5]float64
	QueueLengths  [5]int
}

// Stats returns current queue statistics.
func (pq *PriorityQueue) Stats() *Stats {
	if pq == nil {
		return nil
	}

	stats := &Stats{}
	for i := 0; i < 5; i++ {
		stats.Processed[i] = atomic.LoadInt64(&pq.processed[i])
		stats.QueueLengths[i] = len(pq.queues[i])

		totalWait := atomic.LoadInt64(&pq.totalWait[i])
		if stats.Processed[i] > 0 {
			stats.AvgWaitTimeMs[i] = float64(totalWait) / float64(stats.Processed[i]) / float64(time.Millisecond)
		}
	}

	return stats
}

// Close shuts down the priority queue.
func (pq *PriorityQueue) Close() {
	if pq == nil {
		return
	}

	pq.mu.Lock()
	if pq.closed {
		pq.mu.Unlock()
		return
	}
	pq.closed = true
	close(pq.stopCh)
	pq.mu.Unlock()

	pq.wg.Wait()

	log.Info("priority queue closed")
}

// IsEnabled returns whether the queue is enabled.
func (pq *PriorityQueue) IsEnabled() bool {
	return pq != nil && !pq.closed
}

// DetectPriority determines request priority based on request characteristics.
func DetectPriority(isStreaming bool, messageCount int, headers map[string]string) Priority {
	// Check explicit header
	if p, ok := headers["X-Request-Priority"]; ok {
		switch p {
		case "critical":
			return PriorityCritical
		case "high":
			return PriorityHigh
		case "low":
			return PriorityLow
		case "batch":
			return PriorityBatch
		}
	}

	// Streaming requests are critical
	if isStreaming {
		return PriorityCritical
	}

	// Batch requests
	if _, ok := headers["X-Batch-Request"]; ok {
		return PriorityBatch
	}

	// More messages = established session = higher priority
	if messageCount > 5 {
		return PriorityHigh
	}

	return PriorityNormal
}