package llm

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

// Scheduler bounds concurrent backend LLM calls.
type Scheduler struct {
	permits chan struct{}
}

// 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{
		permits: make(chan struct{}, 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) {
	select {
	case s.permits <- struct{}{}:
		var once sync.Once
		return func() {
			once.Do(func() {
				<-s.permits
			})
		}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// 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)
}