package metrics

import (
	"log/slog"
	"sort"
	"strings"
	"time"
)

// loadFromTSDB loads historical server message data from TSDB into memory.
func (c *ServerMessageCollector) loadFromTSDB() {
	slog.Info("Loading historical server message stats from TSDB into memory...")
	start := time.Now()

	records := c.tsdb.QueryServerMessages(TimeWindow1Week)
	agentRecords := c.tsdb.QueryAgentMessages(TimeWindow1Week)
	agentTokenRecords := c.tsdb.QueryAgentTokenUsage(TimeWindow1Week)
	if len(records) == 0 && len(agentRecords) == 0 && len(agentTokenRecords) == 0 {
		c.loaded.Store(true)
		slog.Info("No historical server message stats found in TSDB")
		return
	}

	now := time.Now()
	todayYearDay := now.YearDay()
	currentHour := now.Hour()

	// Rebuild counters and buckets from TSDB records.
	tempCounters := make(map[string]*serverMessageCounter)
	type bucketData struct {
		counts        map[string]uint64
		received      map[string]uint64
		sent          map[string]uint64
		agentReceived map[string]uint64
		agentSent     map[string]uint64
		agentTokens   map[string]map[string]TokenUsageStats
	}
	tempBucketMap := make(map[time.Time]*bucketData) // bucketTimestamp -> data

	for _, rec := range records {
		recTime := time.UnixMilli(rec.Timestamp)

		// Update counter.
		cc, ok := tempCounters[rec.ServerID]
		if !ok {
			cc = &serverMessageCounter{
				lastDay:  recTime.YearDay(),
				lastHour: recTime.Hour(),
				lastSeen: recTime,
			}
			tempCounters[rec.ServerID] = cc
		}
		if recTime.After(cc.lastSeen) {
			cc.lastSeen = recTime
		}
		cc.total++

		switch rec.Direction {
		case DirectionSent:
			cc.sentTotal++
		default:
			cc.receivedTotal++
		}

		if recTime.YearDay() == todayYearDay {
			cc.today++
			cc.lastDay = todayYearDay
			switch rec.Direction {
			case DirectionSent:
				cc.sentToday++
			default:
				cc.receivedToday++
			}
		}
		if recTime.YearDay() == todayYearDay && recTime.Hour() == currentHour {
			cc.thisHour++
			cc.lastHour = currentHour
			switch rec.Direction {
			case DirectionSent:
				cc.sentThisHour++
			default:
				cc.receivedThisHour++
			}
		}

		// Update bucket.
		bucketKey := recTime.Truncate(serverMessageBucketGranularity)
		bd := tempBucketMap[bucketKey]
		if bd == nil {
			bd = &bucketData{
				counts:        make(map[string]uint64),
				received:      make(map[string]uint64),
				sent:          make(map[string]uint64),
				agentReceived: make(map[string]uint64),
				agentSent:     make(map[string]uint64),
				agentTokens:   make(map[string]map[string]TokenUsageStats),
			}
			tempBucketMap[bucketKey] = bd
		}
		bd.counts[rec.ServerID]++
		switch rec.Direction {
		case DirectionSent:
			bd.sent[rec.ServerID]++
		default:
			bd.received[rec.ServerID]++
		}
	}

	for _, rec := range agentRecords {
		agentID := strings.TrimSpace(rec.AgentID)
		if agentID == "" {
			continue
		}
		bucketKey := time.UnixMilli(rec.Timestamp).Truncate(serverMessageBucketGranularity)
		bd := tempBucketMap[bucketKey]
		if bd == nil {
			bd = &bucketData{
				counts:        make(map[string]uint64),
				received:      make(map[string]uint64),
				sent:          make(map[string]uint64),
				agentReceived: make(map[string]uint64),
				agentSent:     make(map[string]uint64),
				agentTokens:   make(map[string]map[string]TokenUsageStats),
			}
			tempBucketMap[bucketKey] = bd
		}
		switch rec.Direction {
		case DirectionSent:
			bd.agentSent[agentID]++
		default:
			bd.agentReceived[agentID]++
		}
	}

	for _, rec := range agentTokenRecords {
		agentID := strings.TrimSpace(rec.AgentID)
		model := strings.TrimSpace(rec.Model)
		if agentID == "" || model == "" {
			continue
		}
		bucketKey := time.UnixMilli(rec.Timestamp).Truncate(serverMessageBucketGranularity)
		bd := tempBucketMap[bucketKey]
		if bd == nil {
			bd = &bucketData{
				counts:        make(map[string]uint64),
				received:      make(map[string]uint64),
				sent:          make(map[string]uint64),
				agentReceived: make(map[string]uint64),
				agentSent:     make(map[string]uint64),
				agentTokens:   make(map[string]map[string]TokenUsageStats),
			}
			tempBucketMap[bucketKey] = bd
		}
		models := bd.agentTokens[agentID]
		if models == nil {
			models = make(map[string]TokenUsageStats)
			bd.agentTokens[agentID] = models
		}
		stats := models[model]
		stats.Count++
		stats.Input += rec.InputToken
		stats.CachedInput += rec.CachedInputToken
		stats.Output += rec.OutputToken
		models[model] = stats
	}
	tokenEvents := make([]agentTokenEvent, 0, len(agentTokenRecords))
	for _, record := range agentTokenRecords {
		agentID := strings.TrimSpace(record.AgentID)
		model := strings.TrimSpace(record.Model)
		if agentID == "" || model == "" {
			continue
		}
		tokenEvents = append(tokenEvents, agentTokenEvent{
			timestamp: time.UnixMilli(record.Timestamp),
			agentID:   agentID,
			model:     model,
			stats: TokenUsageStats{
				Count:       1,
				Input:       record.InputToken,
				CachedInput: record.CachedInputToken,
				Output:      record.OutputToken,
			},
		})
	}
	sort.Slice(tokenEvents, func(i, j int) bool {
		return tokenEvents[i].timestamp.Before(tokenEvents[j].timestamp)
	})

	// Convert bucket map to sorted slice.
	tempBuckets := make([]serverMessageBucket, 0, len(tempBucketMap))
	for ts, bd := range tempBucketMap {
		tempBuckets = append(tempBuckets, serverMessageBucket{
			timestamp:     ts,
			counts:        bd.counts,
			received:      bd.received,
			sent:          bd.sent,
			agentReceived: bd.agentReceived,
			agentSent:     bd.agentSent,
			agentTokens:   bd.agentTokens,
		})
	}
	sort.Slice(tempBuckets, func(i, j int) bool {
		return tempBuckets[i].timestamp.Before(tempBuckets[j].timestamp)
	})

	// Merge into live state (new messages may have arrived while loading).
	c.mu.Lock()
	for chID, tempCC := range tempCounters {
		existing, exists := c.counters[chID]
		if !exists {
			c.counters[chID] = tempCC
		} else {
			existing.total += tempCC.total
			existing.today += tempCC.today
			existing.thisHour += tempCC.thisHour
			existing.receivedTotal += tempCC.receivedTotal
			existing.receivedToday += tempCC.receivedToday
			existing.receivedThisHour += tempCC.receivedThisHour
			existing.sentTotal += tempCC.sentTotal
			existing.sentToday += tempCC.sentToday
			existing.sentThisHour += tempCC.sentThisHour
			if tempCC.lastSeen.After(existing.lastSeen) {
				existing.lastSeen = tempCC.lastSeen
			}
		}
	}
	// Merge buckets: prepend historical buckets before any live buckets.
	if len(c.buckets) > 0 {
		// Merge overlapping buckets.
		liveBucketMap := make(map[time.Time]int) // timestamp -> index in c.buckets
		for i, b := range c.buckets {
			liveBucketMap[b.timestamp] = i
		}
		for _, hb := range tempBuckets {
			if idx, exists := liveBucketMap[hb.timestamp]; exists {
				// Merge counts into existing live bucket.
				for chID, count := range hb.counts {
					c.buckets[idx].counts[chID] += count
				}
				for chID, count := range hb.received {
					if c.buckets[idx].received == nil {
						c.buckets[idx].received = make(map[string]uint64)
					}
					c.buckets[idx].received[chID] += count
				}
				for chID, count := range hb.sent {
					if c.buckets[idx].sent == nil {
						c.buckets[idx].sent = make(map[string]uint64)
					}
					c.buckets[idx].sent[chID] += count
				}
				for agentID, count := range hb.agentReceived {
					if c.buckets[idx].agentReceived == nil {
						c.buckets[idx].agentReceived = make(map[string]uint64)
					}
					c.buckets[idx].agentReceived[agentID] += count
				}
				for agentID, count := range hb.agentSent {
					if c.buckets[idx].agentSent == nil {
						c.buckets[idx].agentSent = make(map[string]uint64)
					}
					c.buckets[idx].agentSent[agentID] += count
				}
				for agentID, models := range hb.agentTokens {
					if c.buckets[idx].agentTokens == nil {
						c.buckets[idx].agentTokens = make(map[string]map[string]TokenUsageStats)
					}
					existingModels := c.buckets[idx].agentTokens[agentID]
					if existingModels == nil {
						existingModels = make(map[string]TokenUsageStats)
						c.buckets[idx].agentTokens[agentID] = existingModels
					}
					for model, stats := range models {
						existing := existingModels[model]
						existing.Count += stats.Count
						existing.Input += stats.Input
						existing.CachedInput += stats.CachedInput
						existing.Output += stats.Output
						existingModels[model] = existing
					}
				}
			} else {
				c.buckets = append(c.buckets, hb)
			}
		}
		// Re-sort after merge.
		sort.Slice(c.buckets, func(i, j int) bool {
			return c.buckets[i].timestamp.Before(c.buckets[j].timestamp)
		})
	} else {
		c.buckets = tempBuckets
	}
	c.tokenEvents = append(tokenEvents, c.tokenEvents...)
	c.mu.Unlock()

	c.loaded.Store(true)
	slog.Info("Historical server message stats loaded from TSDB",
		"duration", time.Since(start),
		"records", len(records),
		"agent_records", len(agentRecords),
		"agent_token_records", len(agentTokenRecords),
		"servers", len(tempCounters),
	)
}

// flushLoop periodically flushes pending server messages to TSDB.
func (c *ServerMessageCollector) flushLoop() {
	ticker := time.NewTicker(c.flushInterval)
	defer ticker.Stop()

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

// flush writes pending server messages to TSDB.
func (c *ServerMessageCollector) flush() {
	c.mu.Lock()
	pending := c.pending
	pendingAgent := c.pendingAgent
	pendingAgentTokens := c.pendingAgentTokens
	c.pending = nil
	c.pendingAgent = nil
	c.pendingAgentTokens = nil
	c.mu.Unlock()

	if len(pending) > 0 {
		c.tsdb.BatchRecordServerMessages(pending)
	}
	if len(pendingAgent) > 0 {
		c.tsdb.BatchRecordAgentMessages(pendingAgent)
	}
	if len(pendingAgentTokens) > 0 {
		c.tsdb.BatchRecordAgentTokenUsage(pendingAgentTokens)
	}
}