package scheduler

import (
	"container/heap"
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/edgeai/gateway/internal/config"
	"github.com/edgeai/gateway/internal/observability"
	"github.com/edgeai/gateway/internal/task"
)

// Scheduler manages task queuing and execution with priority-based scheduling.
type Scheduler struct {
	mu              sync.Mutex
	queue           *priorityQueue
	running         map[string]*task.Task
	maxRunning      int
	maxQueued       int
	notifyCh        chan struct{}
	logger          *observability.Logger
	ctx             context.Context
	cancel          context.CancelFunc
}

// NewScheduler creates a new scheduler.
func NewScheduler(cfg *config.SchedulerConfig, logger *observability.Logger) *Scheduler {
	ctx, cancel := context.WithCancel(context.Background())
	s := &Scheduler{
		queue:      &priorityQueue{},
		running:    make(map[string]*task.Task),
		maxRunning: cfg.MaxRunningTasks,
		maxQueued:  cfg.MaxQueuedTasks,
		notifyCh:   make(chan struct{}, 1),
		logger:     logger,
		ctx:        ctx,
		cancel:     cancel,
	}
	heap.Init(s.queue)
	return s
}

// Submit adds a task to the queue. Returns error if queue is full.
func (s *Scheduler) Submit(t *task.Task) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.queue.Len() >= s.maxQueued {
		return fmt.Errorf("queue full")
	}

	heap.Push(s.queue, t)
	s.logger.Info("task queued",
		observability.F().
			Event("task_queued").
			TaskID(t.ID).
			Set("priority", config.PriorityName(int(t.Priority))).
			Set("queue_length", s.queue.Len()))

	// Notify the scheduler loop
	select {
	case s.notifyCh <- struct{}{}:
	default:
	}
	return nil
}

// GetNext retrieves the next task to execute (blocking until one is available).
func (s *Scheduler) GetNext(ctx context.Context) (*task.Task, error) {
	for {
		s.mu.Lock()
		if s.queue.Len() > 0 && len(s.running) < s.maxRunning {
			t := heap.Pop(s.queue).(*task.Task)
			s.running[t.ID] = t
			s.mu.Unlock()
			return t, nil
		}
		s.mu.Unlock()

		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-s.notifyCh:
		case <-time.After(100 * time.Millisecond):
		}
	}
}

// Complete marks a task as completed and removes it from running.
func (s *Scheduler) Complete(taskID string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.running, taskID)

	select {
	case s.notifyCh <- struct{}{}:
	default:
	}
}

// QueueLength returns the current queue length.
func (s *Scheduler) QueueLength() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.queue.Len()
}

// RunningCount returns the number of running tasks.
func (s *Scheduler) RunningCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.running)
}

// Stop shuts down the scheduler.
func (s *Scheduler) Stop() {
	s.cancel()
}

// priorityQueue implements heap.Interface for priority-based task scheduling.
type priorityQueue []*task.Task

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	// Lower priority value = higher priority (P0 > P1 > P2...)
	if pq[i].Priority != pq[j].Priority {
		return pq[i].Priority < pq[j].Priority
	}
	// Same priority: FIFO by creation time
	return pq[i].CreatedAt.Before(pq[j].CreatedAt)
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
}

func (pq *priorityQueue) Push(x any) {
	t := x.(*task.Task)
	*pq = append(*pq, t)
}

func (pq *priorityQueue) Pop() any {
	old := *pq
	n := len(old)
	t := old[n-1]
	old[n-1] = nil
	*pq = old[:n-1]
	return t
}