package metrics

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

const (
	serverMessageBucketGranularity    = 5 * time.Minute
	serverMessageDefaultFlushInterval = 30 * time.Second

	// DirectionReceived indicates an inbound message.
	DirectionReceived = "received"
	// DirectionSent indicates an outbound message.
	DirectionSent = "sent"
)

// ServerMessageTimeSeriesPoint represents a single data point in the server message time series.
type ServerMessageTimeSeriesPoint struct {
	Timestamp int64             `json:"timestamp"`
	Counts    map[string]uint64 `json:"counts"`
}

// ServerMessageCollector tracks server message statistics with optional TSDB persistence.
// When a TSDBCollector is provided, it periodically flushes in-memory data to TSDB
// and loads historical data on startup, surviving daemon restarts.
type ServerMessageCollector struct {
	mu       sync.RWMutex
	counters map[string]*serverMessageCounter // keyed by server ID
	buckets  []serverMessageBucket
	// tokenEvents preserves exact event times for point-in-time diagnostics.
	// Coarse dashboard buckets cannot exclude writes later in the same bucket.
	tokenEvents []agentTokenEvent

	// TSDB persistence (nil for pure in-memory mode, e.g. in tests).
	tsdb               *TSDBCollector
	pending            []ServerMessageItem
	pendingAgent       []AgentMessageItem
	pendingAgentTokens []AgentTokenItem

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

type serverMessageCounter struct {
	// Aggregate counters (received + sent).
	total    uint64
	today    uint64
	thisHour uint64

	// Per-direction counters.
	receivedTotal    uint64
	receivedToday    uint64
	receivedThisHour uint64
	sentTotal        uint64
	sentToday        uint64
	sentThisHour     uint64

	lastDay  int
	lastHour int
	lastSeen time.Time
}

type serverMessageBucket struct {
	timestamp     time.Time
	counts        map[string]uint64 // serverID -> count (combined)
	received      map[string]uint64 // serverID -> received count
	sent          map[string]uint64 // serverID -> sent count
	agentReceived map[string]uint64 // agentID -> received count
	agentSent     map[string]uint64 // agentID -> sent count
	agentTokens   map[string]map[string]TokenUsageStats
}

type agentTokenEvent struct {
	timestamp time.Time
	agentID   string
	model     string
	stats     TokenUsageStats
}

// NewServerMessageCollector creates a new server message stats collector.
// If tsdb is non-nil, persistence is enabled; otherwise it operates in pure in-memory mode.
func NewServerMessageCollector(tsdb *TSDBCollector) *ServerMessageCollector {
	return &ServerMessageCollector{
		counters:      make(map[string]*serverMessageCounter),
		tsdb:          tsdb,
		flushInterval: serverMessageDefaultFlushInterval,
		stopCh:        make(chan struct{}),
	}
}

// Start launches background goroutines for TSDB persistence.
// If no TSDB is configured, this is a no-op.
func (c *ServerMessageCollector) Start(ctx context.Context) {
	if c.tsdb == nil {
		c.loaded.Store(true)
		return
	}
	c.ctx, c.cancel = context.WithCancel(ctx)
	go c.loadFromTSDB()
	go c.flushLoop()
}

// Stop signals background goroutines to stop and performs a final flush.
func (c *ServerMessageCollector) Stop() {
	if c.tsdb == nil {
		return
	}
	if c.cancel != nil {
		c.cancel()
	}
	close(c.stopCh)
}

// RecordMessage records a single received message event for a server.
// Deprecated: Use RecordReceived or RecordSent instead. This method delegates
// to RecordReceived for compatibility with older call sites.
func (c *ServerMessageCollector) RecordMessage(serverID string) {
	c.RecordReceived(serverID)
}

// RecordReceived records a single inbound message event for a server.
func (c *ServerMessageCollector) RecordReceived(serverID string) {
	c.recordMessage(serverID, DirectionReceived)
}

// RecordSent records a single outbound message event for a server.
func (c *ServerMessageCollector) RecordSent(serverID string) {
	c.recordMessage(serverID, DirectionSent)
}

// RecordAgentReceived records one inbound message against an agent without
// changing aggregate server message counters.
func (c *ServerMessageCollector) RecordAgentReceived(agentID string) {
	c.recordAgentMessage(agentID, DirectionReceived)
}

// RecordAgentSent records one outbound message against an agent without
// changing aggregate server message counters.
func (c *ServerMessageCollector) RecordAgentSent(agentID string) {
	c.recordAgentMessage(agentID, DirectionSent)
}

// RecordAgentTokenUsage records one runtime token usage event against an agent.
func (c *ServerMessageCollector) RecordAgentTokenUsage(agentID, model string, inputToken, cachedInputToken, outputToken uint64) {
	agentID = strings.TrimSpace(agentID)
	model = strings.TrimSpace(model)
	if agentID == "" || model == "" {
		return
	}
	if inputToken == 0 && cachedInputToken == 0 && outputToken == 0 {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	bucketKey := now.Truncate(serverMessageBucketGranularity)
	bi := len(c.buckets) - 1
	if bi < 0 || !c.buckets[bi].timestamp.Equal(bucketKey) {
		c.buckets = append(c.buckets, serverMessageBucket{
			timestamp:     bucketKey,
			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),
		})
		bi = len(c.buckets) - 1
	}
	c.incrementAgentTokenBucketLocked(bi, agentID, model, inputToken, cachedInputToken, outputToken)
	c.tokenEvents = append(c.tokenEvents, agentTokenEvent{
		timestamp: now,
		agentID:   agentID,
		model:     model,
		stats: TokenUsageStats{
			Count:       1,
			Input:       inputToken,
			CachedInput: cachedInputToken,
			Output:      outputToken,
		},
	})
	c.garbageCollectBuckets(now)

	if c.tsdb != nil {
		c.pendingAgentTokens = append(c.pendingAgentTokens, AgentTokenItem{
			AgentID:          agentID,
			Model:            model,
			Timestamp:        now.UnixMilli(),
			InputToken:       inputToken,
			CachedInputToken: cachedInputToken,
			OutputToken:      outputToken,
		})
	}
}

// RecordReceivedForAgent records an inbound server message and attributes it
// to an agent for per-agent message charts.
func (c *ServerMessageCollector) RecordReceivedForAgent(serverID, agentID string) {
	c.recordMessage(serverID, DirectionReceived)
	c.recordAgentMessage(agentID, DirectionReceived)
}

// RecordSentForAgent records an outbound server message and attributes it
// to an agent for per-agent message charts.
func (c *ServerMessageCollector) RecordSentForAgent(serverID, agentID string) {
	c.recordMessage(serverID, DirectionSent)
	c.recordAgentMessage(agentID, DirectionSent)
}

// recordMessage records a single message event with the specified direction.
func (c *ServerMessageCollector) recordMessage(serverID, direction string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	cc, ok := c.counters[serverID]
	if !ok {
		cc = &serverMessageCounter{}
		c.counters[serverID] = cc
	}
	cc.lastSeen = now

	if cc.lastDay != now.YearDay() {
		cc.today = 0
		cc.receivedToday = 0
		cc.sentToday = 0
		cc.lastDay = now.YearDay()
	}
	if cc.lastHour != now.Hour() {
		cc.thisHour = 0
		cc.receivedThisHour = 0
		cc.sentThisHour = 0
		cc.lastHour = now.Hour()
	}

	cc.total++
	cc.today++
	cc.thisHour++

	switch direction {
	case DirectionSent:
		cc.sentTotal++
		cc.sentToday++
		cc.sentThisHour++
	default: // DirectionReceived
		cc.receivedTotal++
		cc.receivedToday++
		cc.receivedThisHour++
	}

	bucketKey := now.Truncate(serverMessageBucketGranularity)
	bi := len(c.buckets) - 1
	if bi >= 0 && c.buckets[bi].timestamp.Equal(bucketKey) {
		c.buckets[bi].counts[serverID]++
		switch direction {
		case DirectionSent:
			c.buckets[bi].sent[serverID]++
		default:
			c.buckets[bi].received[serverID]++
		}
	} else {
		counts := map[string]uint64{serverID: 1}
		received := make(map[string]uint64)
		sent := make(map[string]uint64)
		switch direction {
		case DirectionSent:
			sent[serverID] = 1
		default:
			received[serverID] = 1
		}
		c.buckets = append(c.buckets, serverMessageBucket{timestamp: bucketKey, counts: counts, received: received, sent: sent})
	}
	c.garbageCollectBuckets(now)

	// Buffer for TSDB persistence.
	if c.tsdb != nil {
		c.pending = append(c.pending, ServerMessageItem{
			ServerID:  serverID,
			Direction: direction,
			Timestamp: now.UnixMilli(),
		})
	}
}

// recordAgentMessage records a single message event in the per-agent bucket
// series used by the dashboard agent detail page.
func (c *ServerMessageCollector) recordAgentMessage(agentID, direction string) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	now := time.Now()
	bucketKey := now.Truncate(serverMessageBucketGranularity)
	bi := len(c.buckets) - 1
	if bi >= 0 && c.buckets[bi].timestamp.Equal(bucketKey) {
		c.incrementAgentBucketLocked(bi, agentID, direction)
	} else {
		bucket := serverMessageBucket{
			timestamp:     bucketKey,
			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),
		}
		c.buckets = append(c.buckets, bucket)
		c.incrementAgentBucketLocked(len(c.buckets)-1, agentID, direction)
	}
	c.garbageCollectBuckets(now)

	if c.tsdb != nil {
		c.pendingAgent = append(c.pendingAgent, AgentMessageItem{
			AgentID:   agentID,
			Direction: direction,
			Timestamp: now.UnixMilli(),
		})
	}
}

func (c *ServerMessageCollector) incrementAgentBucketLocked(index int, agentID, direction string) {
	switch direction {
	case DirectionSent:
		if c.buckets[index].agentSent == nil {
			c.buckets[index].agentSent = make(map[string]uint64)
		}
		c.buckets[index].agentSent[agentID]++
	default:
		if c.buckets[index].agentReceived == nil {
			c.buckets[index].agentReceived = make(map[string]uint64)
		}
		c.buckets[index].agentReceived[agentID]++
	}
}

func (c *ServerMessageCollector) incrementAgentTokenBucketLocked(index int, agentID, model string, inputToken, cachedInputToken, outputToken uint64) {
	if c.buckets[index].agentTokens == nil {
		c.buckets[index].agentTokens = make(map[string]map[string]TokenUsageStats)
	}
	models := c.buckets[index].agentTokens[agentID]
	if models == nil {
		models = make(map[string]TokenUsageStats)
		c.buckets[index].agentTokens[agentID] = models
	}
	stats := models[model]
	stats.Count++
	stats.Input += inputToken
	stats.CachedInput += cachedInputToken
	stats.Output += outputToken
	models[model] = stats
}

func (c *ServerMessageCollector) garbageCollectBuckets(now time.Time) {
	cutoff := now.Add(-TimeWindow1Week.Duration())
	i := 0
	for i < len(c.buckets) && c.buckets[i].timestamp.Before(cutoff) {
		i++
	}
	if i > 0 {
		c.buckets = c.buckets[i:]
	}
	i = 0
	for i < len(c.tokenEvents) && c.tokenEvents[i].timestamp.Before(cutoff) {
		i++
	}
	if i > 0 {
		c.tokenEvents = c.tokenEvents[i:]
	}
	c.pruneInactiveCountersLocked(cutoff)
}

func (c *ServerMessageCollector) pruneInactiveCountersLocked(cutoff time.Time) {
	for serverID, cc := range c.counters {
		if cc.lastSeen.IsZero() || !cc.lastSeen.Before(cutoff) {
			continue
		}
		delete(c.counters, serverID)
	}
}