Luminary/internal/store/cache/cache.go

package cache

import (
	"encoding/json"
	"sync"
	"sync/atomic"
	"time"

	"github.com/rose_cat707/luminary/internal/auth"
	"github.com/rose_cat707/luminary/internal/model"
)

type IngressKeySnapshot struct {
	Key      model.IngressKey
	KeyPlain string // only set on create, not persisted in cache long-term
}

type UsageDelta struct {
	IngressKeyID  uint
	ProviderKeyID uint
	ProviderID    uint
	ClientIP      string
	Endpoint      string
	Model         string
	TokensIn      int
	TokensOut     int
	Cost          float64
	LatencyMs     int64
	Status        string
	Stream        bool
	ErrorMsg      string
	RequestBody   string
	ResponseBody  string
	BudgetUsed    float64
	Requests      int64
	Tokens        int64
}

type Store struct {
	mu sync.RWMutex

	providers     map[uint]*model.Provider
	providerKeys  map[uint]*model.ProviderKey
	providerByName map[string]*model.Provider
	models        map[uint][]model.ModelEntry

	ingressKeys   map[uint]*model.IngressKey
	ingressByHash map[string]*model.IngressKey

	ipRules []model.IPRule
	routingRules map[uint][]model.RoutingRule // ingressKeyID -> rules

	healthStatus map[uint]model.HealthCheckRecord

	usageDeltas []UsageDelta
	usageMu     sync.Mutex

	// rate limit counters: ingressKeyID -> window start + counts
	rateMu      sync.Mutex
	rateWindows map[uint]*rateWindow
}

type rateWindow struct {
	start    time.Time
	requests int
	tokens   int
}

func New() *Store {
	return &Store{
		providers:      make(map[uint]*model.Provider),
		providerKeys:   make(map[uint]*model.ProviderKey),
		providerByName: make(map[string]*model.Provider),
		models:         make(map[uint][]model.ModelEntry),
		ingressKeys:    make(map[uint]*model.IngressKey),
		ingressByHash:  make(map[string]*model.IngressKey),
		routingRules:   make(map[uint][]model.RoutingRule),
		healthStatus:   make(map[uint]model.HealthCheckRecord),
		rateWindows:    make(map[uint]*rateWindow),
	}
}

func (c *Store) LoadProviders(providers []model.Provider) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.providers = make(map[uint]*model.Provider)
	c.providerByName = make(map[string]*model.Provider)
	for i := range providers {
		p := providers[i]
		c.providers[p.ID] = &p
		c.providerByName[p.Name] = &p
	}
}

func (c *Store) LoadProviderKeys(keys []model.ProviderKey) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.providerKeys = make(map[uint]*model.ProviderKey)
	for i := range keys {
		k := keys[i]
		c.providerKeys[k.ID] = &k
	}
}

func (c *Store) LoadModels(models []model.ModelEntry) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.models = make(map[uint][]model.ModelEntry)
	for _, m := range models {
		c.models[m.ProviderID] = append(c.models[m.ProviderID], m)
	}
}

func (c *Store) LoadIngressKeys(keys []model.IngressKey) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ingressKeys = make(map[uint]*model.IngressKey)
	c.ingressByHash = make(map[string]*model.IngressKey)
	for i := range keys {
		k := keys[i]
		c.ingressKeys[k.ID] = &k
		c.ingressByHash[k.KeyHash] = &k
	}
}

func (c *Store) LoadIPRules(rules []model.IPRule) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ipRules = rules
}

func (c *Store) GetProvider(id uint) (model.Provider, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	p, ok := c.providers[id]
	if !ok {
		return model.Provider{}, false
	}
	return *p, true
}

func (c *Store) GetProviderByName(name string) (model.Provider, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	p, ok := c.providerByName[name]
	if !ok {
		return model.Provider{}, false
	}
	return *p, true
}

func (c *Store) ListProviders() []model.Provider {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]model.Provider, 0, len(c.providers))
	for _, p := range c.providers {
		out = append(out, *p)
	}
	return out
}

func (c *Store) SetProvider(p model.Provider) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.providers[p.ID] = &p
	c.providerByName[p.Name] = &p
}

func (c *Store) DeleteProvider(id uint) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if p, ok := c.providers[id]; ok {
		delete(c.providerByName, p.Name)
	}
	delete(c.providers, id)
	delete(c.models, id)
}

func (c *Store) GetProviderKey(id uint) (model.ProviderKey, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	k, ok := c.providerKeys[id]
	if !ok {
		return model.ProviderKey{}, false
	}
	return *k, true
}

func (c *Store) ListProviderKeys(providerID uint) []model.ProviderKey {
	c.mu.RLock()
	defer c.mu.RUnlock()
	var out []model.ProviderKey
	for _, k := range c.providerKeys {
		if k.ProviderID == providerID {
			out = append(out, *k)
		}
	}
	return out
}

func (c *Store) ListAllProviderKeys() []model.ProviderKey {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]model.ProviderKey, 0, len(c.providerKeys))
	for _, k := range c.providerKeys {
		out = append(out, *k)
	}
	return out
}

func (c *Store) SetProviderKey(k model.ProviderKey) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.providerKeys[k.ID] = &k
}

func (c *Store) DeleteProviderKey(id uint) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.providerKeys, id)
}

func (c *Store) GetModels(providerID uint) []model.ModelEntry {
	c.mu.RLock()
	defer c.mu.RUnlock()
	models := c.models[providerID]
	if models == nil {
		return []model.ModelEntry{}
	}
	return models
}

func (c *Store) SetModels(providerID uint, models []model.ModelEntry) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.models[providerID] = models
}

func (c *Store) ListIngressKeys() []model.IngressKey {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]model.IngressKey, 0, len(c.ingressKeys))
	for _, k := range c.ingressKeys {
		out = append(out, *k)
	}
	return out
}

func (c *Store) GetIngressKey(id uint) (model.IngressKey, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	k, ok := c.ingressKeys[id]
	if !ok {
		return model.IngressKey{}, false
	}
	return *k, true
}

func (c *Store) SetIngressKey(k model.IngressKey) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ingressKeys[k.ID] = &k
	c.ingressByHash[k.KeyHash] = &k
}

func (c *Store) DeleteIngressKey(id uint) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if k, ok := c.ingressKeys[id]; ok {
		delete(c.ingressByHash, k.KeyHash)
	}
	delete(c.ingressKeys, id)
}

func (c *Store) LookupIngressKey(plain string) (*model.IngressKey, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	hash := auth.HashIngressKey(plain)
	if k, ok := c.ingressByHash[hash]; ok && k.Enabled {
		cp := *k
		return &cp, true
	}
	for _, k := range c.ingressKeys {
		if !k.Enabled || !auth.IsLegacyIngressHash(k.KeyHash) {
			continue
		}
		if auth.CheckIngressKey(k.KeyHash, plain) {
			cp := *k
			return &cp, true
		}
	}
	return nil, false
}

func (c *Store) ResetProviderKeyDailyCounts() {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, k := range c.providerKeys {
		k.RequestsToday = 0
		k.TokensToday = 0
	}
}

func (c *Store) ListIPRules() []model.IPRule {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make([]model.IPRule, len(c.ipRules))
	copy(out, c.ipRules)
	return out
}

func (c *Store) SetIPRules(rules []model.IPRule) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.ipRules = rules
}

func (c *Store) SetHealthStatus(providerID uint, rec model.HealthCheckRecord) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.healthStatus[providerID] = rec
	if p, ok := c.providers[providerID]; ok {
		p.Status = rec.Status
	}
}

func (c *Store) GetHealthStatus() map[uint]model.HealthCheckRecord {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := make(map[uint]model.HealthCheckRecord, len(c.healthStatus))
	for k, v := range c.healthStatus {
		out[k] = v
	}
	return out
}

func (c *Store) RecordUsage(delta UsageDelta) {
	c.mu.Lock()
	if k, ok := c.ingressKeys[delta.IngressKeyID]; ok {
		k.BudgetUsed += delta.BudgetUsed
		k.RequestCount += delta.Requests
		k.TokenCount += delta.Tokens
	}
	if pk, ok := c.providerKeys[delta.ProviderKeyID]; ok {
		pk.RequestsToday += int(delta.Requests)
		pk.TokensToday += delta.Tokens
	}
	c.mu.Unlock()

	c.usageMu.Lock()
	c.usageDeltas = append(c.usageDeltas, delta)
	c.usageMu.Unlock()
}

func (c *Store) DrainUsageDeltas() []UsageDelta {
	c.usageMu.Lock()
	defer c.usageMu.Unlock()
	if len(c.usageDeltas) == 0 {
		return nil
	}
	out := c.usageDeltas
	c.usageDeltas = nil
	return out
}

func (c *Store) CheckRateLimit(ingressKeyID uint, rpm, tpm int, tokens int) bool {
	if rpm <= 0 && tpm <= 0 {
		return true
	}
	c.rateMu.Lock()
	defer c.rateMu.Unlock()
	now := time.Now()
	w, ok := c.rateWindows[ingressKeyID]
	if !ok || now.Sub(w.start) >= time.Minute {
		c.rateWindows[ingressKeyID] = &rateWindow{start: now, requests: 1, tokens: tokens}
		return true
	}
	if rpm > 0 && w.requests >= rpm {
		return false
	}
	if tpm > 0 && w.tokens+tokens > tpm {
		return false
	}
	w.requests++
	w.tokens += tokens
	return true
}

func ParseJSONList(s string) []string {
	var out []string
	if s == "" {
		return []string{"*"}
	}
	_ = json.Unmarshal([]byte(s), &out)
	if len(out) == 0 {
		return []string{"*"}
	}
	return out
}

func MatchesList(list []string, value string) bool {
	for _, item := range list {
		if item == "*" || item == value {
			return true
		}
	}
	return false
}

func ModelsToJSON(models []string) string {
	b, _ := json.Marshal(models)
	return string(b)
}

func (c *Store) LoadRoutingRules(rules []model.RoutingRule) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.routingRules = make(map[uint][]model.RoutingRule)
	for _, r := range rules {
		c.routingRules[r.IngressKeyID] = append(c.routingRules[r.IngressKeyID], r)
	}
}

func (c *Store) ListRoutingRules(ingressKeyID uint) []model.RoutingRule {
	c.mu.RLock()
	defer c.mu.RUnlock()
	out := c.routingRules[ingressKeyID]
	cp := make([]model.RoutingRule, len(out))
	copy(cp, out)
	return cp
}

func (c *Store) SetRoutingRules(ingressKeyID uint, rules []model.RoutingRule) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.routingRules[ingressKeyID] = rules
}

func (c *Store) DeleteRoutingRulesForIngress(ingressKeyID uint) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.routingRules, ingressKeyID)
}

type Counter struct {
	v atomic.Int64
}

func (c *Counter) Add(n int64) {
	c.v.Add(n)
}

func (c *Counter) Load() int64 {
	return c.v.Load()
}