// Package runtimecache provides small process-local caches for runtime-derived
// facts. Callers own the source of truth and must wire invalidation from that
// owner; this package only stores derived snapshots.
package runtimecache

import (
	"context"
	"sync"
)

// Loader builds a value for key from the caller-owned source of truth.
type Loader[K comparable, V any] func(context.Context, K) (V, error)

// CloneFunc returns a defensive copy of a cached value. It is optional for
// immutable values and recommended for slices, maps, pointers, and structs that
// contain mutable children.
type CloneFunc[V any] func(V) V

// Map is a process-local derived cache keyed by a comparable value. It collapses
// concurrent misses for the same key and epoch, and Invalidate starts a new
// epoch without blocking readers during the caller's expensive reload.
type Map[K comparable, V any] struct {
	mu      sync.RWMutex
	values  map[K]V
	epoch   uint64
	loading map[loadKey[K]]*loadCall[V]
	clone   CloneFunc[V]
}

type loadKey[K comparable] struct {
	key   K
	epoch uint64
}

type loadCall[V any] struct {
	wg    sync.WaitGroup
	value V
	err   error
	stale bool
}

// Option configures a Map.
type Option[K comparable, V any] func(*Map[K, V])

// WithClone configures defensive copies on cache store and cache read.
func WithClone[K comparable, V any](clone CloneFunc[V]) Option[K, V] {
	return func(cache *Map[K, V]) {
		cache.clone = clone
	}
}

// NewMap creates an empty derived cache.
func NewMap[K comparable, V any](opts ...Option[K, V]) *Map[K, V] {
	cache := &Map[K, V]{}
	for _, opt := range opts {
		if opt != nil {
			opt(cache)
		}
	}
	return cache
}

// GetOrLoad returns the cached value for key, or calls loader once for all
// concurrent misses in the current epoch. If an invalidation races with a
// successful load, callers retry in the new epoch rather than observing a stale
// value.
func (cache *Map[K, V]) GetOrLoad(ctx context.Context, key K, loader Loader[K, V]) (V, error) {
	var zero V
	if cache == nil {
		if loader == nil {
			return zero, nil
		}
		return loader(ctx, key)
	}
	if loader == nil {
		return zero, nil
	}

	for {
		cache.mu.RLock()
		if cache.values != nil {
			if value, ok := cache.values[key]; ok {
				cache.mu.RUnlock()
				return cache.cloneValue(value), nil
			}
		}
		loadEpoch := cache.epoch
		cache.mu.RUnlock()

		lk := loadKey[K]{key: key, epoch: loadEpoch}
		if call, owner := cache.loadCall(key, loadEpoch, lk); !owner {
			call.wg.Wait()
			if call.stale && call.err == nil {
				continue
			}
			if call.err != nil {
				return zero, call.err
			}
			return cache.cloneValue(call.value), nil
		} else {
			value, err := loader(ctx, key)
			var stored V
			if err == nil {
				stored = cache.cloneValue(value)
			}

			cache.mu.Lock()
			stale := cache.epoch != loadEpoch
			if !stale && err == nil {
				if cache.values == nil {
					cache.values = make(map[K]V)
				}
				cache.values[key] = stored
			}
			delete(cache.loading, lk)
			cache.mu.Unlock()

			call.value = stored
			call.err = err
			call.stale = stale
			call.wg.Done()

			if stale && err == nil {
				continue
			}
			if err != nil {
				return zero, err
			}
			return cache.cloneValue(stored), nil
		}
	}
}

func (cache *Map[K, V]) loadCall(key K, epoch uint64, lk loadKey[K]) (*loadCall[V], bool) {
	cache.mu.Lock()
	defer cache.mu.Unlock()

	if cache.values != nil {
		if value, ok := cache.values[key]; ok {
			call := &loadCall[V]{value: cache.cloneValue(value)}
			call.wg.Add(1)
			call.wg.Done()
			return call, false
		}
	}
	if cache.epoch != epoch {
		call := &loadCall[V]{stale: true}
		call.wg.Add(1)
		call.wg.Done()
		return call, false
	}
	if cache.loading == nil {
		cache.loading = make(map[loadKey[K]]*loadCall[V])
	}
	if call, ok := cache.loading[lk]; ok {
		return call, false
	}
	call := &loadCall[V]{}
	call.wg.Add(1)
	cache.loading[lk] = call
	return call, true
}

// Invalidate clears cached values and advances the epoch. Loads started before
// the invalidation are not published into the new epoch.
func (cache *Map[K, V]) Invalidate() {
	if cache == nil {
		return
	}
	cache.mu.Lock()
	cache.values = nil
	cache.epoch++
	cache.mu.Unlock()
}

func (cache *Map[K, V]) cloneValue(value V) V {
	if cache != nil && cache.clone != nil {
		return cache.clone(value)
	}
	return value
}