package engine

import (
	"context"
	"fmt"
	"slices"
	"sync"
	"time"
)

const defaultAgingThreshold = 30 * time.Second

type Resources[Self any] interface {
	Fits(Self) bool
	Add(Self) Self
	Sub(Self) Self
}

type ResourceScheduler[R Resources[R]] struct {
	mu             sync.Mutex
	budget         R
	max            R
	used           R
	queue          []*resourceWaiter[R]
	now            func() time.Time // get time now, is a field for mocking
	agingThreshold time.Duration

	OnAcquire  func(req R, allowed bool, reason string)
	OnSnapshot func(budget, max, used R, queueLen int)
}

type resourceWaiter[R Resources[R]] struct {
	req        R
	ready      chan struct{}
	enqueuedAt time.Time
}

type resourceLease[R Resources[R]] struct {
	scheduler *ResourceScheduler[R]
	req       R
	once      sync.Once
}

func NewResourceScheduler[R Resources[R]](budget, max R, agingThreshold time.Duration) *ResourceScheduler[R] {
	if agingThreshold <= 0 {
		agingThreshold = defaultAgingThreshold
	}
	return &ResourceScheduler[R]{
		budget:         budget,
		max:            max,
		now:            time.Now,
		agingThreshold: agingThreshold,
	}
}

// the mill owns the backlog so a NoWait caller must have room immediately or fail
func (s *ResourceScheduler[R]) Acquire(ctx context.Context, req R, mode AcquireMode) (WorkflowSlot, error) {
	if s == nil {
		return NoopSlot{}, nil
	}

	s.mu.Lock()
	if !req.Fits(s.budget) || !req.Fits(s.max) {
		s.mu.Unlock()
		if s.OnAcquire != nil {
			reason := "request_exceeds_budget"
			if !req.Fits(s.max) {
				reason = "request_exceeds_max"
			}
			s.OnAcquire(req, false, reason)
		}
		return nil, fmt.Errorf("%w: request=%v budget=%v max=%v", ErrNoWorkflowSlots, req, s.budget, s.max)
	}
	// ignores the queue because it never blocks
	// only bypasses empty queues to prevent starvation
	if s.used.Add(req).Fits(s.budget) && (mode == NoWait || len(s.queue) == 0) {
		s.used = s.used.Add(req)
		if s.OnAcquire != nil {
			s.OnAcquire(req, true, "allowed_immediately")
		}
		if s.OnSnapshot != nil {
			s.OnSnapshot(s.budget, s.max, s.used, len(s.queue))
		}
		s.mu.Unlock()
		return &resourceLease[R]{scheduler: s, req: req}, nil
	}
	if mode == NoWait {
		used := s.used
		s.mu.Unlock()
		if s.OnAcquire != nil {
			s.OnAcquire(req, false, "no_wait_slot_unavailable")
		}
		return nil, fmt.Errorf("%w: request=%v used=%v budget=%v", ErrNoWorkflowSlots, req, used, s.budget)
	}

	waiter := &resourceWaiter[R]{req: req, ready: make(chan struct{}), enqueuedAt: s.now()}
	s.queue = append(s.queue, waiter)
	if s.OnSnapshot != nil {
		s.OnSnapshot(s.budget, s.max, s.used, len(s.queue))
	}
	s.schedule()
	s.mu.Unlock()

	select {
	case <-waiter.ready:
		if s.OnAcquire != nil {
			s.OnAcquire(req, true, "allowed_from_queue")
		}
		return &resourceLease[R]{scheduler: s, req: req}, nil
	case <-ctx.Done():
		s.mu.Lock()
		select {
		case <-waiter.ready:
			// undo committed resources, schedule already did that
			s.used = s.used.Sub(req)
			if s.OnAcquire != nil {
				s.OnAcquire(req, false, "context_done_after_ready")
			}
		default:
			// still in queue, just remove
			s.remove(waiter)
			if s.OnAcquire != nil {
				s.OnAcquire(req, false, "context_done_in_queue")
			}
		}
		s.schedule()
		if s.OnSnapshot != nil {
			s.OnSnapshot(s.budget, s.max, s.used, len(s.queue))
		}
		s.mu.Unlock()
		return nil, ctx.Err()
	}
}

func (l *resourceLease[R]) Release() {
	if l == nil || l.scheduler == nil {
		return
	}
	l.once.Do(func() {
		l.scheduler.release(l.req)
	})
}

func (s *ResourceScheduler[R]) release(req R) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.used = s.used.Sub(req)
	s.schedule()
	if s.OnSnapshot != nil {
		s.OnSnapshot(s.budget, s.max, s.used, len(s.queue))
	}
}

// start every waiter whose request fits. once a waiter is older than
// agingThreshold, count its request as already used so younger waiters
// stop being scheduled ahead of it
func (s *ResourceScheduler[R]) schedule() {
	var reserved R
	now := s.now()
	i := 0
	changed := false
	for i < len(s.queue) {
		w := s.queue[i]
		if s.used.Add(reserved).Add(w.req).Fits(s.budget) {
			s.queue = slices.Delete(s.queue, i, i+1)
			s.used = s.used.Add(w.req)
			close(w.ready)
			changed = true
			continue
		}
		if now.Sub(w.enqueuedAt) >= s.agingThreshold {
			reserved = reserved.Add(w.req)
		}
		i++
	}
	if changed && s.OnSnapshot != nil {
		s.OnSnapshot(s.budget, s.max, s.used, len(s.queue))
	}
}

func (s *ResourceScheduler[R]) remove(waiter *resourceWaiter[R]) {
	if i := slices.Index(s.queue, waiter); i >= 0 {
		s.queue = slices.Delete(s.queue, i, i+1)
	}
}