package gitee

import (
	"context"
	"errors"
	"sync"
)

var (
	ErrPendingOverFlow = errors.New("Pending queue overflowed, please try again later")
	ErrPendingTimeout  = errors.New("Pending timeout, please try again later")
)

type TaskFunc func() error

type Limiter struct {
	mux            *sync.Mutex   // lock
	max            int64         // max processing
	current        int64         // current procession
	maxPending     int64         // max pending
	currentPending int64         // current pending
	pendingChannel chan struct{} // wait chan
}

// return a new rate limiter instance.
//
// params:
//   max: Number of processing tasks
//   maxPending: Number of pending tasks after processing tasks is overflowed
func NewLimiter(max, maxPending int64) *Limiter {
	return &Limiter{
		mux:            &sync.Mutex{},
		max:            max,
		maxPending:     maxPending,
		pendingChannel: make(chan struct{}, maxPending),
	}
}

// return number of pending tasks
func (l *Limiter) Pending() int64 {
	l.mux.Lock()
	defer l.mux.Unlock()
	return l.currentPending
}

// return number of processing tasks
func (l *Limiter) Processing() int64 {
	l.mux.Lock()
	defer l.mux.Unlock()
	return l.current
}

// return number of processing and pending tasks
func (l *Limiter) ProcessingAndPending() int64 {
	l.mux.Lock()
	defer l.mux.Unlock()
	return l.current + l.currentPending
}

// Set the maximum number of pending tasks
func (l *Limiter) SetMaxPending(maxPending int64) {
	l.mux.Lock()
	defer l.mux.Unlock()
	l.maxPending = maxPending
}

// Set the maximum number of processing tasks
func (l *Limiter) SetMaxProcessing(maxProcessing int64) {
	l.mux.Lock()
	defer l.mux.Unlock()
	l.max = maxProcessing
}

// Run task with rate limit
func (l *Limiter) Limit(ctx context.Context, task TaskFunc) error {
	if err := l.increment(ctx); err != nil {
		return err
	}
	defer l.decrement()

	return task()
}

func (l *Limiter) canProcessing() (bool, error) {
	l.mux.Lock()
	defer l.mux.Unlock()

	if l.current < l.max {
		l.current++
		return true, nil
	}

	if l.currentPending >= l.maxPending {
		return false, ErrPendingOverFlow
	}

	l.currentPending++
	return false, nil
}

// currentPending -1 after pending timeout
func (l *Limiter) pendingTimeout() error {
	l.mux.Lock()
	defer l.mux.Unlock()
	l.currentPending--
	return ErrPendingTimeout
}

func (l *Limiter) increment(ctx context.Context) error {
	ok, err := l.canProcessing()
	if err != nil {
		return err
	}

	if ok {
		return nil
	}

	select {
	case <-l.pendingChannel:
		return nil
	case <-ctx.Done():
		return l.pendingTimeout()
	}
}

func (l *Limiter) decrement() {
	l.mux.Lock()
	defer l.mux.Unlock()

	if l.currentPending > 0 {
		l.pendingChannel <- struct{}{}
		l.currentPending--
		return
	}

	l.current--
	return
}