package pushy

import (
	"sync"
	"time"
)

// Registry maps deviceID -> owning node, with TTLs so a crashed node's
// entries expire (you can't rely on a clean Unregister after a crash).
// In production this is a low-latency KV service (Netflix: KeyValue,
// formerly Dynomite).
type Registry interface {
	Register(deviceID, node string, ttl time.Duration)
	Touch(deviceID, node string, ttl time.Duration)
	Unregister(deviceID, node string)
	Lookup(deviceID string) (node string, ok bool)
}

type entry struct {
	node    string
	expires time.Time
}

// MemRegistry is an in-memory Registry with an injectable clock.
type MemRegistry struct {
	mu    sync.RWMutex
	m     map[string]entry
	nowFn func() time.Time
}

func NewMemRegistry() *MemRegistry {
	return &MemRegistry{m: map[string]entry{}, nowFn: time.Now}
}

func (r *MemRegistry) now() time.Time { return r.nowFn() }

func (r *MemRegistry) Register(d, n string, ttl time.Duration) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.m[d] = entry{node: n, expires: r.now().Add(ttl)}
}

func (r *MemRegistry) Touch(d, n string, ttl time.Duration) { r.Register(d, n, ttl) }

func (r *MemRegistry) Unregister(d, n string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	// Only remove if WE still own it (avoid clobbering a device that
	// already reconnected to another node).
	if e, ok := r.m[d]; ok && e.node == n {
		delete(r.m, d)
	}
}

func (r *MemRegistry) Lookup(d string) (string, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	e, ok := r.m[d]
	if !ok || r.now().After(e.expires) {
		return "", false // expired entries are misses (crashed-node cleanup)
	}
	return e.node, true
}

// Len returns the number of (possibly expired) entries (for tests/obs).
func (r *MemRegistry) Len() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.m)
}