package metrics

import (
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestServerMessageCollector_RecordAndSummary(t *testing.T) {
	c := NewServerMessageCollector(nil)

	c.RecordReceived("ch1")
	c.RecordReceived("ch1")
	c.RecordSent("ch2")

	summary := c.Summary()
	if len(summary) != 2 {
		t.Fatalf("Summary length = %d, want 2", len(summary))
	}

	var ch1, ch2 *ServerMessageSummary
	for i := range summary {
		if summary[i].ServerID == "ch1" {
			ch1 = &summary[i]
		}
		if summary[i].ServerID == "ch2" {
			ch2 = &summary[i]
		}
	}

	if ch1 == nil {
		t.Fatal("ch1 not found in summary")
	}
	if ch1.Total != 2 {
		t.Errorf("ch1 Total = %d, want 2", ch1.Total)
	}
	if ch1.ReceivedTotal != 2 {
		t.Errorf("ch1 ReceivedTotal = %d, want 2", ch1.ReceivedTotal)
	}
	if ch1.SentTotal != 0 {
		t.Errorf("ch1 SentTotal = %d, want 0", ch1.SentTotal)
	}
	if ch2 == nil {
		t.Fatal("ch2 not found in summary")
	}
	if ch2.Total != 1 {
		t.Errorf("ch2 Total = %d, want 1", ch2.Total)
	}
	if ch2.ReceivedTotal != 0 {
		t.Errorf("ch2 ReceivedTotal = %d, want 0", ch2.ReceivedTotal)
	}
	if ch2.SentTotal != 1 {
		t.Errorf("ch2 SentTotal = %d, want 1", ch2.SentTotal)
	}
}

func TestServerMessageCollector_Empty(t *testing.T) {
	c := NewServerMessageCollector(nil)
	summary := c.Summary()
	if len(summary) != 0 {
		t.Errorf("Summary length = %d, want 0", len(summary))
	}
}

func TestServerMessageCollector_TodayAndThisHour(t *testing.T) {
	c := NewServerMessageCollector(nil)
	c.RecordReceived("ch1")

	summary := c.Summary()
	if len(summary) != 1 {
		t.Fatalf("Summary length = %d, want 1", len(summary))
	}
	if summary[0].Today != 1 {
		t.Errorf("Today = %d, want 1", summary[0].Today)
	}
	if summary[0].ThisHour != 1 {
		t.Errorf("ThisHour = %d, want 1", summary[0].ThisHour)
	}
	if summary[0].ReceivedToday != 1 {
		t.Errorf("ReceivedToday = %d, want 1", summary[0].ReceivedToday)
	}
	if summary[0].ReceivedThisHour != 1 {
		t.Errorf("ReceivedThisHour = %d, want 1", summary[0].ReceivedThisHour)
	}
	if summary[0].SentToday != 0 {
		t.Errorf("SentToday = %d, want 0", summary[0].SentToday)
	}
	if summary[0].SentThisHour != 0 {
		t.Errorf("SentThisHour = %d, want 0", summary[0].SentThisHour)
	}
}

func TestServerMessageCollector_MultipleServers(t *testing.T) {
	c := NewServerMessageCollector(nil)
	for i := 0; i < 100; i++ {
		c.RecordReceived("ch1")
		c.RecordSent("ch2")
		c.RecordReceived("ch3")
	}

	summary := c.Summary()
	if len(summary) != 3 {
		t.Fatalf("Summary length = %d, want 3", len(summary))
	}
	for _, s := range summary {
		if s.Total != 100 {
			t.Errorf("server %s Total = %d, want 100", s.ServerID, s.Total)
		}
	}
}

func TestServerMessageCollector_ReceivedAndSentBreakdown(t *testing.T) {
	c := NewServerMessageCollector(nil)

	c.RecordReceived("ch1")
	c.RecordReceived("ch1")
	c.RecordReceived("ch1")
	c.RecordSent("ch1")
	c.RecordSent("ch1")

	summary := c.Summary()
	if len(summary) != 1 {
		t.Fatalf("Summary length = %d, want 1", len(summary))
	}

	s := summary[0]
	if s.Total != 5 {
		t.Errorf("Total = %d, want 5", s.Total)
	}
	if s.ReceivedTotal != 3 {
		t.Errorf("ReceivedTotal = %d, want 3", s.ReceivedTotal)
	}
	if s.SentTotal != 2 {
		t.Errorf("SentTotal = %d, want 2", s.SentTotal)
	}
	if s.Today != 5 {
		t.Errorf("Today = %d, want 5", s.Today)
	}
	if s.ReceivedToday != 3 {
		t.Errorf("ReceivedToday = %d, want 3", s.ReceivedToday)
	}
	if s.SentToday != 2 {
		t.Errorf("SentToday = %d, want 2", s.SentToday)
	}
}

func TestServerMessageCollector_DirectionalTimeSeriesAggregatesByWindowInterval(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)
	firstBucket := now.Add(-20 * time.Minute)
	secondBucket := now.Add(-10 * time.Minute)

	c.buckets = []serverMessageBucket{
		{
			timestamp: firstBucket,
			received:  map[string]uint64{"ch1": 1},
			sent:      map[string]uint64{"ch1": 2},
		},
		{
			timestamp: firstBucket.Add(5 * time.Minute),
			received:  map[string]uint64{"ch1": 3},
			sent:      map[string]uint64{"ch1": 4},
		},
		{
			timestamp: secondBucket,
			received:  map[string]uint64{"ch1": 5},
			sent:      map[string]uint64{"ch1": 6},
		},
	}

	result := c.QueryDirectionalTimeSeries(TimeWindow12Hours)
	require.Equal(t, int((10 * time.Minute).Seconds()), result.Interval)

	byTimestamp := make(map[int64]ServerMessageDirectionalTimeSeriesPoint)
	for _, point := range result.Data {
		byTimestamp[point.Timestamp] = point
	}

	first := byTimestamp[firstBucket.Truncate(10*time.Minute).Unix()]
	require.Equal(t, uint64(4), first.Received)
	require.Equal(t, uint64(6), first.Sent)

	second := byTimestamp[secondBucket.Truncate(10*time.Minute).Unix()]
	require.Equal(t, uint64(5), second.Received)
	require.Equal(t, uint64(6), second.Sent)
}

func TestServerMessageCollector_QueryAgentMessageBarSeries(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)
	firstBucket := now.Add(-2 * time.Hour)
	secondBucket := now.Add(-30 * time.Minute)

	c.buckets = []serverMessageBucket{
		{
			timestamp:     firstBucket,
			agentReceived: map[string]uint64{"agent-a": 2, "agent-b": 9},
			agentSent:     map[string]uint64{"agent-a": 1},
		},
		{
			timestamp:     secondBucket,
			agentReceived: map[string]uint64{"agent-a": 3},
			agentSent:     map[string]uint64{"agent-a": 4, "agent-b": 7},
		},
	}

	result := c.QueryAgentMessageBarSeries("agent-a", TimeWindow6Hours)
	require.Equal(t, int((30 * time.Minute).Seconds()), result.Interval)
	require.Equal(t, []BarSeriesEntry{
		{Key: "inbound", Label: "Inbound"},
		{Key: "outbound", Label: "Outbound"},
	}, result.Series)

	byTimestamp := make(map[int64]BarSeriesPoint)
	for _, point := range result.Data {
		byTimestamp[point.Timestamp] = point
	}

	first := byTimestamp[firstBucket.Truncate(30*time.Minute).Unix()]
	require.Equal(t, uint64(2), first.Values["inbound"])
	require.Equal(t, uint64(1), first.Values["outbound"])

	second := byTimestamp[secondBucket.Truncate(30*time.Minute).Unix()]
	require.Equal(t, uint64(3), second.Values["inbound"])
	require.Equal(t, uint64(4), second.Values["outbound"])
}

func TestServerMessageCollector_QueryAgentMessageSummary(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)

	c.buckets = []serverMessageBucket{
		{
			timestamp:     now.Add(-2 * time.Hour),
			agentReceived: map[string]uint64{"agent-a": 2, "agent-b": 9},
			agentSent:     map[string]uint64{"agent-a": 1},
		},
		{
			timestamp:     now.Add(-30 * time.Minute),
			agentReceived: map[string]uint64{"agent-a": 3},
			agentSent:     map[string]uint64{"agent-a": 4, "agent-b": 7},
		},
		{
			timestamp:     now.Add(-7 * time.Hour),
			agentReceived: map[string]uint64{"agent-a": 100},
			agentSent:     map[string]uint64{"agent-a": 100},
		},
	}

	summary := c.QueryAgentMessageSummary("agent-a", TimeWindow6Hours)
	require.Equal(t, "agent-a", summary.AgentID)
	require.Equal(t, uint64(5), summary.Received)
	require.Equal(t, uint64(5), summary.Sent)
	require.Equal(t, uint64(10), summary.Total)
}

func TestServerMessageCollector_QueryAgentMessageDirectionalTimeSeries(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)
	firstBucket := now.Add(-20 * time.Minute)
	secondBucket := now.Add(-10 * time.Minute)

	c.buckets = []serverMessageBucket{
		{
			timestamp:     firstBucket,
			agentReceived: map[string]uint64{"agent-a": 1, "agent-b": 9},
			agentSent:     map[string]uint64{"agent-a": 2},
		},
		{
			timestamp:     firstBucket.Add(5 * time.Minute),
			agentReceived: map[string]uint64{"agent-a": 3},
			agentSent:     map[string]uint64{"agent-a": 4},
		},
		{
			timestamp:     secondBucket,
			agentReceived: map[string]uint64{"agent-a": 5},
			agentSent:     map[string]uint64{"agent-a": 6, "agent-b": 7},
		},
	}

	result := c.QueryAgentMessageDirectionalTimeSeries("agent-a", TimeWindow12Hours)
	require.Equal(t, int((10 * time.Minute).Seconds()), result.Interval)

	byTimestamp := make(map[int64]ServerMessageDirectionalTimeSeriesPoint)
	for _, point := range result.Data {
		byTimestamp[point.Timestamp] = point
	}

	first := byTimestamp[firstBucket.Truncate(10*time.Minute).Unix()]
	require.Equal(t, uint64(4), first.Received)
	require.Equal(t, uint64(6), first.Sent)

	second := byTimestamp[secondBucket.Truncate(10*time.Minute).Unix()]
	require.Equal(t, uint64(5), second.Received)
	require.Equal(t, uint64(6), second.Sent)
}

func TestServerMessageCollector_QueryAgentTokenStats(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)
	firstBucket := now.Add(-2 * time.Hour)
	secondBucket := now.Add(-30 * time.Minute)

	c.buckets = []serverMessageBucket{
		{
			timestamp: firstBucket,
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {
					"openai/gpt-4o":    {Count: 1, Input: 100, CachedInput: 30, Output: 20},
					"anthropic/sonnet": {Count: 1, Input: 40, CachedInput: 10, Output: 15},
				},
				"agent-b": {
					"openai/gpt-4o": {Count: 1, Input: 999, CachedInput: 999, Output: 999},
				},
			},
		},
		{
			timestamp: secondBucket,
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {
					"openai/gpt-4o": {Count: 2, Input: 150, CachedInput: 20, Output: 70},
				},
			},
		},
		{
			timestamp: now.Add(-7 * time.Hour),
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {
					"openai/gpt-4o": {Count: 1, Input: 1000, CachedInput: 1000, Output: 1000},
				},
			},
		},
	}

	summary := c.QueryAgentTokenSummary("agent-a", TimeWindow6Hours)
	require.Equal(t, "agent-a", summary.AgentID)
	require.Equal(t, uint64(4), summary.Count)
	require.Equal(t, uint64(290), summary.Input)
	require.Equal(t, uint64(60), summary.Cached)
	require.Equal(t, uint64(105), summary.Output)
	require.Equal(t, uint64(395), summary.Total)

	trend := c.QueryAgentTokenTimeSeries("agent-a", TimeWindow6Hours)
	require.Equal(t, int((5 * time.Minute).Seconds()), trend.Interval)
	trendByTimestamp := make(map[int64]TimeSeriesDataPoint)
	for _, point := range trend.Data {
		trendByTimestamp[point.Timestamp] = point
	}
	firstTrend := trendByTimestamp[firstBucket.Truncate(5*time.Minute).Unix()]
	require.Equal(t, uint64(2), firstTrend.Count)
	require.Equal(t, uint64(140), firstTrend.Input)
	require.Equal(t, uint64(40), firstTrend.Cached)
	require.Equal(t, uint64(35), firstTrend.Output)

	bar := c.QueryAgentModelTokenBarSeries("agent-a", TimeWindow6Hours, "input")
	require.Equal(t, int((30 * time.Minute).Seconds()), bar.Interval)
	require.Equal(t, []BarSeriesEntry{
		{Key: "anthropic/sonnet", Label: "sonnet"},
		{Key: "openai/gpt-4o", Label: "gpt-4o"},
	}, bar.Series)
	barByTimestamp := make(map[int64]BarSeriesPoint)
	for _, point := range bar.Data {
		barByTimestamp[point.Timestamp] = point
	}
	firstBar := barByTimestamp[firstBucket.Truncate(30*time.Minute).Unix()]
	require.Equal(t, uint64(40), firstBar.Values["anthropic/sonnet"])
	require.Equal(t, uint64(100), firstBar.Values["openai/gpt-4o"])
	secondBar := barByTimestamp[secondBucket.Truncate(30*time.Minute).Unix()]
	require.Equal(t, uint64(150), secondBar.Values["openai/gpt-4o"])
}

func TestServerMessageCollector_QueryAgentsDashboardScopesAndAggregates(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now().Truncate(time.Hour)
	c.buckets = []serverMessageBucket{
		{
			timestamp:     now.Add(-30 * time.Minute),
			agentReceived: map[string]uint64{"agent-a": 2, "agent-b": 3, "other": 99},
			agentSent:     map[string]uint64{"agent-a": 1, "agent-b": 4, "other": 99},
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {"model-a": {Count: 1, Input: 10, CachedInput: 2, Output: 3}},
				"agent-b": {"model-b": {Count: 2, Input: 20, CachedInput: 4, Output: 6}},
				"other":   {"model-x": {Count: 9, Input: 900, CachedInput: 900, Output: 900}},
			},
		},
		{
			timestamp:     now.Add(-8 * time.Hour),
			agentReceived: map[string]uint64{"agent-a": 100},
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {"model-a": {Count: 1, Input: 100, Output: 100}},
			},
		},
	}

	result := c.QueryAgentsDashboard([]string{"agent-b", "agent-a", "agent-a", ""}, TimeWindow6Hours)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.MessageTrend.Interval)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.MessageBar.Interval)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.TokenTrend.Interval)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.InputTokenBar.Interval)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.CachedTokenBar.Interval)
	require.Equal(t, int((20 * time.Minute).Seconds()), result.OutputTokenBar.Interval)
	require.Equal(t, uint64(5), result.MessageSummary.Received)
	require.Equal(t, uint64(5), result.MessageSummary.Sent)
	require.Equal(t, uint64(10), result.MessageSummary.Total)
	require.Equal(t, uint64(30), result.TokenSummary.Input)
	require.Equal(t, uint64(6), result.TokenSummary.Cached)
	require.Equal(t, uint64(9), result.TokenSummary.Output)
	require.Equal(t, uint64(39), result.TokenSummary.Total)
	require.Equal(t, []string{"agent-a", "agent-b"}, []string{result.Agents[0].AgentID, result.Agents[1].AgentID})
	require.Equal(t, uint64(13), result.Agents[0].Total)
	require.Equal(t, uint64(26), result.Agents[1].Total)
	require.NotEmpty(t, result.MessageTrend.Data)
	require.NotEmpty(t, result.MessageBar.Data)
	require.NotEmpty(t, result.TokenTrend.Data)
	require.Equal(t, []BarSeriesEntry{{Key: "model-a", Label: "model-a"}, {Key: "model-b", Label: "model-b"}}, result.InputTokenBar.Series)
	require.NotEmpty(t, result.InputTokenBar.Data)
	require.NotEmpty(t, result.CachedTokenBar.Data)
	require.NotEmpty(t, result.OutputTokenBar.Data)
	dashboardBucket := now.Add(-30 * time.Minute).Truncate(MachineDashboardOverviewAggregationInterval(TimeWindow6Hours)).Unix()
	messagePoints := make(map[int64]BarSeriesPoint, len(result.MessageBar.Data))
	for _, point := range result.MessageBar.Data {
		messagePoints[point.Timestamp] = point
	}
	messagePoint := messagePoints[dashboardBucket]
	require.Equal(t, uint64(5), messagePoint.Values["inbound"])
	require.Equal(t, uint64(5), messagePoint.Values["outbound"])
	inputPoints := make(map[int64]BarSeriesPoint, len(result.InputTokenBar.Data))
	for _, point := range result.InputTokenBar.Data {
		inputPoints[point.Timestamp] = point
	}
	inputPoint := inputPoints[dashboardBucket]
	require.Equal(t, uint64(10), inputPoint.Values["model-a"])
	require.Equal(t, uint64(20), inputPoint.Values["model-b"])
}

func TestServerMessageCollector_QueryAgentsDashboardAtExcludesFutureBuckets(t *testing.T) {
	c := NewServerMessageCollector(nil)
	snapshotAt := time.Now().Truncate(time.Minute)
	c.buckets = []serverMessageBucket{
		{
			timestamp: snapshotAt.Add(-time.Minute),
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {"model": {Input: 10}},
			},
		},
		{
			timestamp: snapshotAt.Add(time.Minute),
			agentTokens: map[string]map[string]TokenUsageStats{
				"agent-a": {"model": {Input: 90}},
			},
		},
	}
	c.tokenEvents = []agentTokenEvent{
		{
			timestamp: snapshotAt.Add(-time.Second),
			agentID:   "agent-a",
			model:     "model",
			stats:     TokenUsageStats{Count: 1, Input: 10},
		},
		{
			timestamp: snapshotAt.Add(time.Second),
			agentID:   "agent-a",
			model:     "model",
			stats:     TokenUsageStats{Count: 1, Input: 90},
		},
	}

	report := c.QueryAgentsDashboardAt([]string{"agent-a"}, TimeWindow6Hours, snapshotAt)
	require.Equal(t, uint64(10), report.TokenSummary.Input)
}

func TestServerMessageCollector_BackwardCompatRecordMessage(t *testing.T) {
	c := NewServerMessageCollector(nil)

	// RecordMessage should behave as RecordReceived.
	c.RecordMessage("ch1")
	c.RecordMessage("ch1")

	summary := c.Summary()
	if len(summary) != 1 {
		t.Fatalf("Summary length = %d, want 1", len(summary))
	}

	s := summary[0]
	if s.Total != 2 {
		t.Errorf("Total = %d, want 2", s.Total)
	}
	if s.ReceivedTotal != 2 {
		t.Errorf("ReceivedTotal = %d, want 2", s.ReceivedTotal)
	}
	if s.SentTotal != 0 {
		t.Errorf("SentTotal = %d, want 0", s.SentTotal)
	}
}

func TestServerMessageCollectorPrunesInactiveCounters(t *testing.T) {
	c := NewServerMessageCollector(nil)
	now := time.Now()
	c.counters["inactive"] = &serverMessageCounter{
		total:    1,
		lastSeen: now.Add(-TimeWindow1Week.Duration() - time.Minute),
	}
	c.counters["active"] = &serverMessageCounter{
		total:    1,
		lastSeen: now,
	}
	c.buckets = []serverMessageBucket{
		{
			timestamp: now.Add(-TimeWindow1Week.Duration() - time.Minute),
			counts:    map[string]uint64{"inactive": 1},
			received:  map[string]uint64{"inactive": 1},
			sent:      map[string]uint64{},
		},
		{
			timestamp: now.Truncate(serverMessageBucketGranularity),
			counts:    map[string]uint64{"active": 1},
			received:  map[string]uint64{"active": 1},
			sent:      map[string]uint64{},
		},
	}

	c.garbageCollectBuckets(now)

	require.NotContains(t, c.counters, "inactive")
	require.Contains(t, c.counters, "active")
	require.Len(t, c.buckets, 1)
	require.Contains(t, c.buckets[0].counts, "active")
}