package llm

import (
	"context"
	"fmt"
	"sync"
)

// Scheduler bounds concurrent backend LLM calls.
type Scheduler struct {
	maxConcurrency int
	mu             sync.Mutex
	inFlight       int
	queue          []*waiter
}

type waiter struct {
	ready   chan struct{}
	queued  bool
	granted bool
}

// NewScheduler creates a scheduler with a fixed concurrency limit.
func NewScheduler(maxConcurrency int) (*Scheduler, error) {
	if maxConcurrency <= 0 {
		return nil, fmt.Errorf("max concurrency must be greater than zero")
	}
	return &Scheduler{
		maxConcurrency: maxConcurrency,
	}, nil
}

// Acquire blocks until a permit is available or the context is canceled.
// The returned release function is safe to call multiple times.
func (s *Scheduler) Acquire(ctx context.Context) (func(), error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	s.mu.Lock()
	if s.inFlight < s.maxConcurrency && len(s.queue) == 0 {
		s.inFlight++
		s.mu.Unlock()
		return s.releaseFunc(), nil
	}
	w := &waiter{
		ready:  make(chan struct{}),
		queued: true,
	}
	s.queue = append(s.queue, w)
	s.mu.Unlock()

	select {
	case <-w.ready:
		return s.releaseFunc(), nil
	case <-ctx.Done():
		s.mu.Lock()
		if w.queued {
			s.removeQueuedWaiterLocked(w)
			s.mu.Unlock()
			return nil, ctx.Err()
		}
		if w.granted {
			// Permit was granted concurrently with cancellation; release it to
			// avoid leaks and unblock queued callers.
			s.inFlight--
			s.grantQueuedLocked()
		}
		s.mu.Unlock()
		return nil, ctx.Err()
	}
}

func (s *Scheduler) releaseFunc() func() {
	var once sync.Once
	return func() {
		once.Do(func() {
			s.mu.Lock()
			if s.inFlight > 0 {
				s.inFlight--
				s.grantQueuedLocked()
			}
			s.mu.Unlock()
		})
	}
}

func (s *Scheduler) grantQueuedLocked() {
	for s.inFlight < s.maxConcurrency && len(s.queue) > 0 {
		w := s.queue[0]
		s.queue = s.queue[1:]
		w.queued = false
		w.granted = true
		s.inFlight++
		close(w.ready)
	}
}

func (s *Scheduler) removeQueuedWaiterLocked(target *waiter) {
	for i, w := range s.queue {
		if w == target {
			w.queued = false
			s.queue = append(s.queue[:i], s.queue[i+1:]...)
			return
		}
	}
}

// Run acquires a permit, executes fn, and releases the permit.
func (s *Scheduler) Run(ctx context.Context, fn func(context.Context) error) error {
	release, err := s.Acquire(ctx)
	if err != nil {
		return err
	}
	defer release()
	return fn(ctx)
}