package batch

import (
	"context"
	"maps"
	"sync"

	E "github.com/sagernet/sing/common/exceptions"
)

type Option[T any] func(b *Batch[T])

type Result[T any] struct {
	Value T
	Err   error
}

type Error struct {
	Key string
	Err error
}

func (e *Error) Error() string {
	return E.Cause(e.Err, e.Key).Error()
}

func WithConcurrencyNum[T any](n int) Option[T] {
	return func(b *Batch[T]) {
		q := make(chan struct{}, n)
		for range n {
			q <- struct{}{}
		}
		b.queue = q
	}
}

// Batch similar to errgroup, but can control the maximum number of concurrent
type Batch[T any] struct {
	result map[string]Result[T]
	queue  chan struct{}
	wg     sync.WaitGroup
	mux    sync.Mutex
	err    *Error
	once   sync.Once
	cancel context.CancelCauseFunc
}

func (b *Batch[T]) Go(key string, fn func() (T, error)) {
	b.wg.Go(func() {
		if b.queue != nil {
			<-b.queue
			defer func() {
				b.queue <- struct{}{}
			}()
		}

		value, err := fn()
		if err != nil {
			b.once.Do(func() {
				b.err = &Error{key, err}
				if b.cancel != nil {
					b.cancel(b.err)
				}
			})
		}

		ret := Result[T]{value, err}
		b.mux.Lock()
		defer b.mux.Unlock()
		b.result[key] = ret
	})
}

func (b *Batch[T]) Wait() *Error {
	b.wg.Wait()
	if b.cancel != nil {
		b.cancel(nil)
	}
	return b.err
}

func (b *Batch[T]) WaitAndGetResult() (map[string]Result[T], *Error) {
	err := b.Wait()
	return b.Result(), err
}

func (b *Batch[T]) Result() map[string]Result[T] {
	b.mux.Lock()
	defer b.mux.Unlock()
	result := make(map[string]Result[T], len(b.result))
	maps.Copy(result, b.result)
	return result
}

func New[T any](ctx context.Context, opts ...Option[T]) (*Batch[T], context.Context) {
	ctx, cancel := context.WithCancelCause(ctx)

	b := &Batch[T]{
		result: map[string]Result[T]{},
	}

	for _, o := range opts {
		o(b)
	}

	b.cancel = cancel
	return b, ctx
}
