package metrics

import (
	"encoding/json"
	"log/slog"
	"os"
	"path/filepath"
	"time"
)

func (mc *MemoryCollector) loadFromTSDB() {
	slog.Info("Loading historical stats from TSDB into memory...")
	start := time.Now()

	summary := mc.tsdb.StatsSummary()

	tempBuckets := make(map[string]map[string]map[int64]*requestBucket)
	tempKeyMasks := make(map[string]string)

	for keyID, aks := range summary {
		tempKeyMasks[keyID] = aks.MaskedKey
		for model, ms := range aks.ModelStats {
			if tempBuckets[keyID] == nil {
				tempBuckets[keyID] = make(map[string]map[int64]*requestBucket)
			}
			if tempBuckets[keyID][model] == nil {
				tempBuckets[keyID][model] = make(map[int64]*requestBucket)
			}
			ms.mu.RLock()
			for _, tc := range ms.Counters {
				bucketTs := tc.Timestamp.Unix()
				b, exists := tempBuckets[keyID][model][bucketTs]
				if !exists {
					b = &requestBucket{}
					tempBuckets[keyID][model][bucketTs] = b
				}
				b.Count += tc.Count
				b.InputToken += tc.InputToken
				b.CachedInputToken += tc.CachedInputToken
				b.OutputToken += tc.OutputToken
				b.Duration += tc.Duration
			}
			ms.mu.RUnlock()
		}
	}

	errorTS := mc.tsdb.QueryErrorTimeSeries(TimeWindow1Week)
	tempErrorBuckets := make(map[int64]map[ErrorType]uint64)
	for _, dp := range errorTS {
		if tempErrorBuckets[dp.Timestamp] == nil {
			tempErrorBuckets[dp.Timestamp] = make(map[ErrorType]uint64)
		}
		for et, count := range dp.Errors {
			tempErrorBuckets[dp.Timestamp][et] += count
		}
	}

	errorSum := mc.tsdb.QueryErrorSummary()
	tempLastErrorTime := errorSum.LastErrorTime

	mc.mu.Lock()
	for keyID, maskedKey := range tempKeyMasks {
		if _, exists := mc.keyMaskMap[keyID]; !exists {
			mc.keyMaskMap[keyID] = maskedKey
		}
	}
	for keyID, models := range tempBuckets {
		if mc.requestBuckets[keyID] == nil {
			mc.requestBuckets[keyID] = make(map[string]map[int64]*requestBucket)
		}
		for model, buckets := range models {
			if mc.requestBuckets[keyID][model] == nil {
				mc.requestBuckets[keyID][model] = make(map[int64]*requestBucket)
			}
			for ts, b := range buckets {
				existing, exists := mc.requestBuckets[keyID][model][ts]
				if !exists {
					mc.requestBuckets[keyID][model][ts] = b
				} else {
					existing.Count += b.Count
					existing.InputToken += b.InputToken
					existing.CachedInputToken += b.CachedInputToken
					existing.OutputToken += b.OutputToken
					existing.Duration += b.Duration
				}
			}
		}
	}
	for ts, errors := range tempErrorBuckets {
		if mc.errorBuckets[ts] == nil {
			mc.errorBuckets[ts] = make(map[ErrorType]uint64)
		}
		for et, count := range errors {
			mc.errorBuckets[ts][et] += count
		}
	}
	if tempLastErrorTime > mc.lastErrorTime {
		mc.lastErrorTime = tempLastErrorTime
	}
	mc.invalidateAllCachesLocked()
	mc.mu.Unlock()

	mc.loaded.Store(true)
	slog.Info("Historical stats loaded from TSDB",
		"duration", time.Since(start),
		"keys", len(mc.requestBuckets),
	)
}

func (mc *MemoryCollector) flushLoop() {
	ticker := time.NewTicker(mc.flushInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			mc.flush()
			mc.pruneOldBuckets()
		case <-mc.stopCh:
			mc.flush()
			return
		case <-mc.ctx.Done():
			mc.flush()
			return
		}
	}
}

func (mc *MemoryCollector) flush() {
	mc.mu.Lock()
	requests := mc.pendingRequests
	errors := mc.pendingErrors
	dirty := mc.keyMaskDirty
	mc.pendingRequests = nil
	mc.pendingErrors = nil
	mc.keyMaskDirty = false
	mc.mu.Unlock()

	if len(requests) > 0 {
		mc.tsdb.BatchRecordRequests(requests)
	}

	if len(errors) > 0 {
		mc.tsdb.BatchRecordErrors(errors)
	}

	if dirty {
		if err := mc.saveKeyMaskMap(); err != nil {
			slog.Warn("Failed to save key mask map", "error", err)
		}
	}
}

func (mc *MemoryCollector) saveKeyMaskMap() error {
	maskMapPath := filepath.Join(statsDir, "key_mask_map.json")

	mc.mu.RLock()
	data, err := json.MarshalIndent(mc.keyMaskMap, "", "  ")
	mc.mu.RUnlock()
	if err != nil {
		return err
	}

	return os.WriteFile(maskMapPath, data, 0600)
}

func (mc *MemoryCollector) pruneOldBuckets() {
	mc.mu.Lock()
	defer mc.mu.Unlock()

	cutoff := time.Now().Add(-TimeWindow1Week.Duration()).Unix()

	for keyID, models := range mc.requestBuckets {
		for model, buckets := range models {
			for ts := range buckets {
				if ts < cutoff {
					delete(buckets, ts)
				}
			}
			if len(buckets) == 0 {
				delete(models, model)
			}
		}
		if len(models) == 0 {
			delete(mc.requestBuckets, keyID)
		}
	}

	for ts := range mc.errorBuckets {
		if ts < cutoff {
			delete(mc.errorBuckets, ts)
		}
	}
	mc.invalidateAllCachesLocked()
}