package notifier

import (
	"sync"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	ctrl "sigs.k8s.io/controller-runtime"
)

const (
	// MaxNotifierChannelSize is the max size of notifier channel.
	MaxNotifierChannelSize = 10
)

var (
	log = ctrl.Log.WithName("notifier")
)

type Notifier struct {
	sync.RWMutex

	keyMap map[string]map[string]chan interface{}
}

func New() *Notifier {
	return &Notifier{keyMap: make(map[string]map[string]chan interface{})}
}

func (n *Notifier) Register(key, id string) chan interface{} {
	defer metrics.NotifierRegisterCounter.WithLabelValues().Inc()

	n.Lock()
	defer n.Unlock()

	if _, ok := n.keyMap[key]; !ok {
		n.keyMap[key] = make(map[string]chan interface{})
	}
	if ch, ok := n.keyMap[key][id]; ok {
		return ch
	}
	ch := make(chan interface{}, MaxNotifierChannelSize)
	n.keyMap[key][id] = ch
	return ch
}

func (n *Notifier) Unregister(key, id string) {
	defer metrics.NotifierUnregisterCounter.WithLabelValues().Inc()

	n.Lock()
	defer n.Unlock()

	n.unregister(key, id)
}

func (n *Notifier) unregister(key, id string) {
	if _, ok := n.keyMap[key]; !ok {
		return
	}
	delete(n.keyMap[key], id)
	if len(n.keyMap[key]) == 0 {
		delete(n.keyMap, key)
	}
}

// Send sends data to all registered channels for the given key with no block.
// If the channel is full, the data will be dropped.
func (n *Notifier) Send(key string, data interface{}) {
	defer metrics.NotifierSendCounter.WithLabelValues().Inc()

	closeChannels := make(map[string]chan interface{})

	n.RLock()
	if m, ok := n.keyMap[key]; ok {
		for id, ch := range m {
			select {
			case ch <- data:
				metrics.NotifierSendChanCounter.WithLabelValues("done").Inc()
			default:
				closeChannels[id] = ch
				metrics.NotifierSendChanCounter.WithLabelValues("block").Inc()
				log.Info("block send data and will close channel", "key", key, "id", id, "data", data)
			}
		}
	}
	n.RUnlock()

	for id, ch := range closeChannels {
		n.Lock()
		n.unregister(key, id)
		close(ch)
		n.Unlock()
	}
}