package metrics

import (
	"context"
	"sync"
	"sync/atomic"
	"time"
)

const (
	defaultCacheTTL           = 5 * time.Second
	defaultFlushInterval      = 30 * time.Second
	maxTimeSeriesCacheEntries = 512
)

type requestBucket struct {
	Count            uint64
	InputToken       uint64
	CachedInputToken uint64
	OutputToken      uint64
	Duration         uint64
}

type cacheEntry[T any] struct {
	data     T
	cachedAt time.Time
}

func (e *cacheEntry[T]) get(ttl time.Duration) (T, bool) {
	if e != nil && time.Since(e.cachedAt) < ttl {
		return e.data, true
	}
	var zero T
	return zero, false
}

func (mc *MemoryCollector) invalidateRequestCachesLocked() {
	mc.summaryCache = nil
	for key := range mc.tsCache {
		delete(mc.tsCache, key)
	}
}

func (mc *MemoryCollector) invalidateErrorCachesLocked() {
	mc.errorSummaryCache = nil
	for key := range mc.errorTsCache {
		delete(mc.errorTsCache, key)
	}
}

func (mc *MemoryCollector) invalidateAllCachesLocked() {
	mc.invalidateRequestCachesLocked()
	mc.invalidateErrorCachesLocked()
}

// MemoryCollector provides in-memory real-time aggregation and caching for statistics,
// with periodic batch flushing to TSDB for persistence. On startup, historical data
// is asynchronously loaded from TSDB into memory.
type MemoryCollector struct {
	tsdb *TSDBCollector

	mu sync.RWMutex

	requestBuckets map[string]map[string]map[int64]*requestBucket
	errorBuckets   map[int64]map[ErrorType]uint64
	keyMaskMap     map[string]string
	lastErrorTime  int64

	pendingRequests []BatchRequestItem
	pendingErrors   []BatchErrorItem
	keyMaskDirty    bool

	cacheTTL          time.Duration
	summaryCache      *cacheEntry[map[string]*APIKeyStats]
	tsCache           map[string]*cacheEntry[[]TimeSeriesDataPoint]
	errorTsCache      map[TimeWindow]*cacheEntry[[]ErrorTimeSeriesDataPoint]
	errorSummaryCache *cacheEntry[ErrorSummary]

	flushInterval time.Duration

	ctx    context.Context
	cancel context.CancelFunc
	stopCh chan struct{}
	loaded atomic.Bool

	ownsDB bool // true when this collector created its own TSDBCollector
}

// NewMemoryCollector creates a new MemoryCollector with its own TSDB backend.
func NewMemoryCollector() (*MemoryCollector, error) {
	tsdb, err := NewTSDBCollector()
	if err != nil {
		return nil, err
	}
	mc := NewMemoryCollectorWithTSDB(tsdb)
	mc.ownsDB = true
	return mc, nil
}

// NewMemoryCollectorWithTSDB creates a new MemoryCollector using the given
// shared TSDBCollector. The caller is responsible for closing the TSDBCollector.
func NewMemoryCollectorWithTSDB(tsdb *TSDBCollector) *MemoryCollector {
	mc := &MemoryCollector{
		tsdb:           tsdb,
		requestBuckets: make(map[string]map[string]map[int64]*requestBucket),
		errorBuckets:   make(map[int64]map[ErrorType]uint64),
		keyMaskMap:     make(map[string]string),
		cacheTTL:       defaultCacheTTL,
		flushInterval:  defaultFlushInterval,
		tsCache:        make(map[string]*cacheEntry[[]TimeSeriesDataPoint]),
		errorTsCache:   make(map[TimeWindow]*cacheEntry[[]ErrorTimeSeriesDataPoint]),
		stopCh:         make(chan struct{}),
	}

	tsdb.mu.RLock()
	for k, v := range tsdb.keyMaskMap {
		mc.keyMaskMap[k] = v
	}
	tsdb.mu.RUnlock()

	return mc
}

// Start begins loading historical TSDB data and flushing pending samples.
func (mc *MemoryCollector) Start(ctx context.Context) {
	mc.ctx, mc.cancel = context.WithCancel(ctx)
	go mc.loadFromTSDB()
	go mc.flushLoop()
}

// Stop cancels background work and performs a final flush.
func (mc *MemoryCollector) Stop() {
	mc.cancel()
	close(mc.stopCh)
}

// Close closes the owned TSDB backend, if this collector created it.
func (mc *MemoryCollector) Close() error {
	if mc.ownsDB {
		return mc.tsdb.Close()
	}
	return nil
}

// RecordRequest records one LLM request sample in memory and queues it for TSDB.
func (mc *MemoryCollector) RecordRequest(
	apiKey, model string,
	inputToken, cachedInputToken, outputToken uint64,
	duration time.Duration,
) {

	keyID := KeyID(apiKey)
	now := time.Now()
	bucketTs := now.Truncate(bucketGranularity).Unix()
	timestamp := now.UnixMilli()
	durationMs := uint64(duration.Milliseconds())

	mc.mu.Lock()

	if _, exists := mc.keyMaskMap[keyID]; !exists {
		mc.keyMaskMap[keyID] = MaskAPIKey(apiKey)
		mc.keyMaskDirty = true
	}

	if mc.requestBuckets[keyID] == nil {
		mc.requestBuckets[keyID] = make(map[string]map[int64]*requestBucket)
	}
	if mc.requestBuckets[keyID][model] == nil {
		mc.requestBuckets[keyID][model] = make(map[int64]*requestBucket)
	}
	b, exists := mc.requestBuckets[keyID][model][bucketTs]
	if !exists {
		b = &requestBucket{}
		mc.requestBuckets[keyID][model][bucketTs] = b
	}
	b.Count++
	b.InputToken += inputToken
	b.CachedInputToken += cachedInputToken
	b.OutputToken += outputToken
	b.Duration += durationMs

	mc.pendingRequests = append(mc.pendingRequests, BatchRequestItem{
		KeyID:            keyID,
		Model:            model,
		Timestamp:        timestamp,
		InputToken:       inputToken,
		CachedInputToken: cachedInputToken,
		OutputToken:      outputToken,
		DurationMs:       durationMs,
	})

	mc.invalidateRequestCachesLocked()
	mc.mu.Unlock()
}

// RecordError records one model error sample in memory and queues it for TSDB.
func (mc *MemoryCollector) RecordError(model string, errorType ErrorType) {
	now := time.Now()
	bucketTs := now.Truncate(bucketGranularity).Unix()
	timestamp := now.UnixMilli()

	mc.mu.Lock()

	if mc.errorBuckets[bucketTs] == nil {
		mc.errorBuckets[bucketTs] = make(map[ErrorType]uint64)
	}
	mc.errorBuckets[bucketTs][errorType]++

	nowUnix := now.Unix()
	if nowUnix > mc.lastErrorTime {
		mc.lastErrorTime = nowUnix
	}

	mc.pendingErrors = append(mc.pendingErrors, BatchErrorItem{
		Model:     model,
		ErrorType: errorType,
		Timestamp: timestamp,
	})

	mc.invalidateErrorCachesLocked()
	mc.mu.Unlock()
}