package limiting

import (
	"fmt"

	semaphores "git.frostfs.info/TrueCloudLab/frostfs-qos/limiting/semaphore"
)

type ReleaseFunc func()

type Limiter interface {
	// Acquire attempts to reserve a slot without blocking.
	//
	// Returns a release function and true if successful, otherwise false.
	// The release function must be called exactly once.
	//
	// If the key was not defined in the limiter, no limit is applied.
	Acquire(key string) (ReleaseFunc, bool)
}

type semaphore interface {
	Acquire() bool
	Release()
}

type semaphoreLimiter[T semaphore] struct {
	m map[string]T
}

// KeyLimit defines a concurrency limit for a set of keys.
//
// All keys of one set share the same limit.
// Keys of different sets have separate limits.
//
// Sets must not overlap.
type KeyLimit struct {
	Keys  []string
	Limit int64
}

var NewBurstAtomicSemaphoreLimiter = func(limits []KeyLimit) (*semaphoreLimiter[*semaphores.BurstAtomicSemaphore], error) {
	return newSemaphoreLimiter(limits, semaphores.NewBurstAtomicSemaphore)
}

func newSemaphoreLimiter[T semaphore](limits []KeyLimit, newSemaphore func(size int64) T) (*semaphoreLimiter[T], error) {
	lr := semaphoreLimiter[T]{make(map[string]T)}
	for _, limit := range limits {
		if err := lr.addLimit(&limit, newSemaphore); err != nil {
			return nil, err
		}
	}
	return &lr, nil
}

func (lr *semaphoreLimiter[T]) addLimit(limit *KeyLimit, newSemaphore func(size int64) T) error {
	if limit.Limit < 0 {
		return fmt.Errorf("invalid limit %d", limit.Limit)
	}

	sem := newSemaphore(limit.Limit)
	for _, key := range limit.Keys {
		if _, exists := lr.m[key]; exists {
			return fmt.Errorf("duplicate key %q", key)
		}
		lr.m[key] = sem
	}

	return nil
}

func (lr *semaphoreLimiter[T]) Acquire(key string) (ReleaseFunc, bool) {
	sem, ok := lr.m[key]
	if !ok {
		return func() {}, true
	}

	if ok := sem.Acquire(); ok {
		return func() { sem.Release() }, true
	}
	return nil, false
}