package metrics

import (
	"sync"

	"github.com/hashicorp/golang-lru/simplelru"
	"github.com/prometheus/client_golang/prometheus"
)

const (
	defaultMaxTokens = 2048
)

type histEntry struct {
	vec      *prometheus.HistogramVec
	tokenKey string
}

func (e histEntry) DeletePartialMatch(token string) {
	if e.vec == nil || e.tokenKey == "" {
		return
	}
	e.vec.DeletePartialMatch(prometheus.Labels{e.tokenKey: token})
}

type ctrEntry struct {
	vec      *prometheus.CounterVec
	tokenKey string
}

func (e ctrEntry) DeletePartialMatch(token string) {
	if e.vec == nil || e.tokenKey == "" {
		return
	}
	e.vec.DeletePartialMatch(prometheus.Labels{e.tokenKey: token})
}

type ReclaimerMetricList struct {
	mu sync.Mutex

	// registered metric vecs
	histogramVecs map[string]histEntry
	counterVecs   map[string]ctrEntry

	// LRU cache for tokens
	tokenCache simplelru.LRUCache
}

func newReclaimerMetricList() *ReclaimerMetricList {
	r := &ReclaimerMetricList{
		histogramVecs: make(map[string]histEntry),
		counterVecs:   make(map[string]ctrEntry),
	}

	cache, err := simplelru.NewLRU(defaultMaxTokens, r.onEvict)
	if err != nil {
		panic(err)
	}
	r.tokenCache = cache

	return r
}

var Reclaimer = newReclaimerMetricList()

func (r *ReclaimerMetricList) AddHistogramVec(name string, tokenKey string, vec *prometheus.HistogramVec) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.histogramVecs[name] = histEntry{vec: vec, tokenKey: tokenKey}
}

func (r *ReclaimerMetricList) AddCounterVec(name string, tokenKey string, vec *prometheus.CounterVec) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.counterVecs[name] = ctrEntry{vec: vec, tokenKey: tokenKey}
}

func (r *ReclaimerMetricList) GetHistogramVec(name, token string) prometheus.ObserverVec {
	r.mu.Lock()
	defer r.mu.Unlock()

	entry, ok := r.histogramVecs[name]
	if !ok {
		return nil
	}

	// update token access in LRU
	r.tokenCache.Add(token, struct{}{})

	return entry.vec.MustCurryWith(prometheus.Labels{entry.tokenKey: token})
}

func (r *ReclaimerMetricList) GetCounterVec(name, token string) *prometheus.CounterVec {
	r.mu.Lock()
	defer r.mu.Unlock()

	entry, ok := r.counterVecs[name]
	if !ok {
		return nil
	}

	// update token access in LRU
	r.tokenCache.Add(token, struct{}{})

	return entry.vec.MustCurryWith(prometheus.Labels{entry.tokenKey: token})
}

func (r *ReclaimerMetricList) DeleteVec(name, token string) {
	r.mu.Lock()
	defer r.mu.Unlock()

	// remove token access in LRU
	r.tokenCache.Remove(token)

	if entry, ok := r.histogramVecs[name]; ok {
		entry.DeletePartialMatch(token)
		return
	}
	if entry, ok := r.counterVecs[name]; ok {
		entry.DeletePartialMatch(token)
	}
}

func (r *ReclaimerMetricList) onEvict(key, value any) {
	token := key.(string)

	// NOTE: no need to lock here since onEvict is called from Get methods which already hold the lock

	// delete metrics for this token from all registered vecs
	for _, entry := range r.histogramVecs {
		entry.DeletePartialMatch(token)
	}

	for _, entry := range r.counterVecs {
		entry.DeletePartialMatch(token)
	}
}