package maxtimes

import (
	"context"
	"sync"

	"github.com/go-logr/logr"
)

// m := NewMaxTimes()
// m.Register("aa", 3)
// m.Register("bb", 3)
// m.Register("cc", 3)
//
// ctx, cancel := context.WithCancel(context.Background())
// go m.Run(ctx, func() { cancel() })
// ...
// go loop(ctx, "aa") // m.Do("aa")
// go loop(ctx, "bb")
// go loop(ctx, "cc")
// <- ctx.Done()
type MaxTimes interface {
	Register(times int, keys ...string)
	// do once, return false if this item has been done
	Do(key string) bool
	Done(key string)
	KeysExceedMaxTimes() []string

	Run(ctx context.Context, f func())
}

type maxTimes struct {
	lock sync.Mutex
	m    map[string]*Count

	logger              logr.Logger
	doneCh              chan struct{}
	keysExceedsMaxTimes []string
}

type Count struct {
	Current int
	Max     int
}

func NewMaxTimes(logger logr.Logger) MaxTimes {
	return &maxTimes{
		m:      make(map[string]*Count),
		doneCh: make(chan struct{}),
		logger: logger,
	}
}

func (m *maxTimes) Register(times int, keys ...string) {
	m.lock.Lock()
	defer m.lock.Unlock()

	m.logger.Info("register max times", "keys", keys)
	for _, key := range keys {
		m.m[key] = &Count{
			Max: times,
		}
	}
}

func (m *maxTimes) Do(key string) bool {
	m.lock.Lock()
	defer m.lock.Unlock()

	count, ok := m.m[key]
	if !ok {
		return false
	}
	count.Current += 1
	if count.Current > count.Max {
		m.done(key, true)
		return false
	}

	m.logger.Info("this item is not done", "item", key, "count", count.Current, "max", count.Max)
	return true
}

func (m *maxTimes) Done(key string) {
	m.lock.Lock()
	defer m.lock.Unlock()

	m.done(key, false)
}

func (m *maxTimes) done(key string, exceed bool) {
	delete(m.m, key)
	if exceed {
		m.keysExceedsMaxTimes = append(m.keysExceedsMaxTimes, key)
	}

	if len(m.m) == 0 {
		close(m.doneCh)
		m.logger.Info("all items are done")
		return
	}

	m.logger.Info("not all items are done", "items", len(m.m))
}

func (m *maxTimes) KeysExceedMaxTimes() []string {
	return m.keysExceedsMaxTimes
}

func (m *maxTimes) Run(ctx context.Context, f func()) {
	select {
	case <-ctx.Done():
	case <-m.doneCh:
	}

	m.logger.Info("shutting down")

	f()
}