package metrics

import (
	"context"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/prometheus/prometheus/tsdb"
)

// newTestTSDB creates a TSDBCollector backed by a temporary directory.
// The caller should defer cleanup of the returned dir and call tc.Close().
func newTestTSDB(t *testing.T) (*TSDBCollector, string) {
	t.Helper()

	dir := t.TempDir()
	dbPath := filepath.Join(dir, "test.tsdb")

	opts := tsdb.DefaultOptions()
	opts.RetentionDuration = int64(7 * 24 * time.Hour / time.Millisecond)
	opts.MinBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.MaxBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.NoLockfile = true

	db, err := tsdb.Open(dbPath, nil, nil, opts, nil)
	if err != nil {
		t.Fatalf("failed to open test TSDB: %v", err)
	}

	maskMapPath := filepath.Join(dir, "key_mask_map.json")
	tc := &TSDBCollector{
		db:          db,
		promptDB:    db,
		stopCh:      make(chan struct{}),
		keyMaskMap:  make(map[string]string),
		maskMapPath: maskMapPath,
	}

	return tc, dir
}

func waitForMetricsCollectorLoad(t *testing.T, name string, loaded func() bool) {
	t.Helper()

	waitForMetricsCondition(t, name+" collector load from TSDB", 5*time.Second, loaded)
}

func waitForMetricsCondition(t *testing.T, name string, timeout time.Duration, condition func() bool) {
	t.Helper()

	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	timer := time.NewTimer(timeout)
	defer timer.Stop()

	for {
		if condition() {
			return
		}
		select {
		case <-timer.C:
			if condition() {
				return
			}
			t.Fatalf("%s did not complete within %s", name, timeout)
		case <-ticker.C:
		}
	}
}

func waitNextMetricsMillisecond(t *testing.T) {
	t.Helper()

	start := time.Now().UnixMilli()
	ticker := time.NewTicker(time.Millisecond)
	defer ticker.Stop()
	timer := time.NewTimer(time.Second)
	defer timer.Stop()

	for {
		if time.Now().UnixMilli() != start {
			return
		}
		select {
		case <-timer.C:
			if time.Now().UnixMilli() != start {
				return
			}
			t.Fatal("clock did not advance to next millisecond within 1s")
		case <-ticker.C:
		}
	}
}

// --------------------------------------------------------------------------
// Test 1: ServerMessageCollector round-trip through TSDB
// --------------------------------------------------------------------------

func TestServerMessageCollector_PersistenceRoundTrip(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Phase 1: Create first collector, record messages, flush to TSDB.
	c1 := NewServerMessageCollector(tsdbCol)
	c1.flushInterval = 100 * time.Millisecond // speed up for test
	c1.Start(ctx)
	waitForMetricsCollectorLoad(t, "server message collector", c1.loaded.Load)

	c1.RecordReceived("ch-agentteam-1")
	waitNextMetricsMillisecond(t)
	c1.RecordReceived("ch-agentteam-1")
	waitNextMetricsMillisecond(t)
	c1.RecordSent("ch-agentteam-1")
	waitNextMetricsMillisecond(t)
	c1.RecordReceived("ch-agentteam-2")
	waitNextMetricsMillisecond(t)
	c1.RecordSent("ch-agentteam-2")
	c1.flush()

	// Verify in-memory state of first collector.
	summary1 := c1.Summary()
	if len(summary1) != 2 {
		t.Fatalf("Phase 1: expected 2 servers in summary, got %d", len(summary1))
	}
	var agentteamOneTotal, agentteamTwoTotal uint64
	for _, s := range summary1 {
		switch s.ServerID {
		case "ch-agentteam-1":
			agentteamOneTotal = s.Total
		case "ch-agentteam-2":
			agentteamTwoTotal = s.Total
		}
	}
	if agentteamOneTotal != 3 {
		t.Errorf("Phase 1: ch-agentteam-1 total = %d, want 3", agentteamOneTotal)
	}
	if agentteamTwoTotal != 2 {
		t.Errorf("Phase 1: ch-agentteam-2 total = %d, want 2", agentteamTwoTotal)
	}

	// Stop first collector (triggers final flush).
	c1.Stop()

	// Verify data is in TSDB.
	records := tsdbCol.QueryServerMessages(TimeWindow6Hours)
	if len(records) != 5 {
		t.Fatalf("Phase 1 TSDB: expected 5 records, got %d", len(records))
	}

	// Phase 2: Create a NEW collector from the same TSDB (simulates restart).
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()

	c2 := NewServerMessageCollector(tsdbCol)
	c2.flushInterval = 100 * time.Millisecond
	c2.Start(ctx2)

	// Wait for loadFromTSDB to complete.
	waitForMetricsCollectorLoad(t, "second server message collector", c2.loaded.Load)

	// Verify restored summary.
	summary2 := c2.Summary()
	if len(summary2) != 2 {
		t.Fatalf("Phase 2: expected 2 servers in summary, got %d", len(summary2))
	}

	agentteamOneTotal, agentteamTwoTotal = 0, 0
	for _, s := range summary2 {
		switch s.ServerID {
		case "ch-agentteam-1":
			agentteamOneTotal = s.Total
		case "ch-agentteam-2":
			agentteamTwoTotal = s.Total
		}
	}
	if agentteamOneTotal != 3 {
		t.Errorf("Phase 2: ch-agentteam-1 total = %d, want 3", agentteamOneTotal)
	}
	if agentteamTwoTotal != 2 {
		t.Errorf("Phase 2: ch-agentteam-2 total = %d, want 2", agentteamTwoTotal)
	}

	// Verify QueryTimeSeries also has data after restore.
	ts := c2.QueryTimeSeries(TimeWindow6Hours)
	if len(ts.Data) == 0 {
		t.Error("Phase 2: QueryTimeSeries returned no data after restore")
	}

	// Verify the time series has counts for our servers.
	totalFromTS := uint64(0)
	for _, dp := range ts.Data {
		for _, count := range dp.Counts {
			totalFromTS += count
		}
	}
	if totalFromTS != 5 {
		t.Errorf("Phase 2: total from TimeSeries = %d, want 5", totalFromTS)
	}

	c2.Stop()
}

// --------------------------------------------------------------------------
// Test 2: ServerMessageCollector with nil TSDB (pure in-memory, no persistence)
// --------------------------------------------------------------------------

func TestServerMessageCollector_NilTSDB_NoPanic(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c := NewServerMessageCollector(nil)
	c.Start(ctx) // should be a no-op for nil TSDB

	c.RecordMessage("ch1")
	c.RecordMessage("ch2")

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

	ts := c.QueryTimeSeries(TimeWindow6Hours)
	if ts.Data == nil {
		t.Error("expected non-nil time series data")
	}

	c.Stop() // should not panic
}

// --------------------------------------------------------------------------
// Test 3: Shared TSDB between ServerMessageCollector and MemoryCollector (LLM)
// --------------------------------------------------------------------------

func TestSharedTSDB_ServerMessagesAndLLM(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Create both collectors on the same TSDB.
	serverMessageCol := NewServerMessageCollector(tsdbCol)
	serverMessageCol.flushInterval = 100 * time.Millisecond
	serverMessageCol.Start(ctx)

	llmCol := NewMemoryCollectorWithTSDB(tsdbCol)
	llmCol.flushInterval = 100 * time.Millisecond
	llmCol.Start(ctx)

	waitForMetricsCollectorLoad(t, "server message collector", serverMessageCol.loaded.Load)
	waitForMetricsCollectorLoad(t, "LLM", llmCol.loaded.Load)

	// Record server messages.
	serverMessageCol.RecordMessage("ch-agentteam-1")
	waitNextMetricsMillisecond(t)
	serverMessageCol.RecordMessage("ch-agentteam-1")
	waitNextMetricsMillisecond(t)
	serverMessageCol.RecordMessage("ch-agentteam-2")

	// Record LLM requests.
	llmCol.RecordRequest("sk-test-key-1", "openai/gpt-4o", 100, 50, 200, 500*time.Millisecond)
	waitNextMetricsMillisecond(t)
	llmCol.RecordRequest("sk-test-key-1", "openai/gpt-4o", 150, 60, 250, 600*time.Millisecond)

	serverMessageCol.flush()
	llmCol.flush()

	// Verify server message data is in TSDB.
	serverMessageRecords := tsdbCol.QueryServerMessages(TimeWindow6Hours)
	if len(serverMessageRecords) != 3 {
		t.Errorf("expected 3 server message records in TSDB, got %d", len(serverMessageRecords))
	}

	// Verify LLM data is in TSDB — use StatsSummary which reads all request metrics.
	llmSummary := tsdbCol.StatsSummary()
	if len(llmSummary) == 0 {
		t.Error("expected LLM stats in TSDB, got none")
	}

	// Verify server message collector doesn't see LLM data and vice versa.
	serverMessageSummary := serverMessageCol.Summary()
	serverMessageTotalMessages := uint64(0)
	for _, s := range serverMessageSummary {
		serverMessageTotalMessages += s.Total
	}
	if serverMessageTotalMessages != 3 {
		t.Errorf("server message total messages = %d, want 3", serverMessageTotalMessages)
	}

	llmStatsSummary := llmCol.StatsSummary()
	llmTotalRequests := uint64(0)
	for _, aks := range llmStatsSummary {
		for _, ms := range aks.ModelStats {
			stats := ms.Stats()
			llmTotalRequests += stats[TimeWindow6Hours].Count
		}
	}
	if llmTotalRequests != 2 {
		t.Errorf("LLM total requests = %d, want 2", llmTotalRequests)
	}

	// Stop both collectors.
	serverMessageCol.Stop()
	llmCol.Stop()

	// Phase 2: Simulate restart — create new collectors from same TSDB.
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()

	serverMessageCol2 := NewServerMessageCollector(tsdbCol)
	serverMessageCol2.flushInterval = 100 * time.Millisecond
	serverMessageCol2.Start(ctx2)

	llmCol2 := NewMemoryCollectorWithTSDB(tsdbCol)
	llmCol2.flushInterval = 100 * time.Millisecond
	llmCol2.Start(ctx2)

	waitForMetricsCollectorLoad(t, "server message collector", serverMessageCol2.loaded.Load)
	waitForMetricsCollectorLoad(t, "LLM", llmCol2.loaded.Load)

	// Verify restored server message data.
	restoredServerMessageSummary := serverMessageCol2.Summary()
	restoredTotal := uint64(0)
	for _, s := range restoredServerMessageSummary {
		restoredTotal += s.Total
	}
	if restoredTotal != 3 {
		t.Errorf("restored server message total = %d, want 3", restoredTotal)
	}

	// Verify restored LLM data.
	restoredLLMSummary := llmCol2.StatsSummary()
	restoredLLMTotal := uint64(0)
	for _, aks := range restoredLLMSummary {
		for _, ms := range aks.ModelStats {
			stats := ms.Stats()
			restoredLLMTotal += stats[TimeWindow6Hours].Count
		}
	}
	if restoredLLMTotal != 2 {
		t.Errorf("restored LLM total = %d, want 2", restoredLLMTotal)
	}

	serverMessageCol2.Stop()
	llmCol2.Stop()
}

// --------------------------------------------------------------------------
// Test 4: ownsDB flag — Close only closes TSDB when ownsDB is true
// --------------------------------------------------------------------------

func TestMemoryCollector_OwnsDB(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)

	// Shared mode: ownsDB = false.
	mc := NewMemoryCollectorWithTSDB(tsdbCol)
	if mc.ownsDB {
		t.Error("NewMemoryCollectorWithTSDB should set ownsDB to false")
	}
	if err := mc.Close(); err != nil {
		t.Errorf("Close() on shared collector should be no-op, got error: %v", err)
	}

	// TSDB should still be usable after shared collector Close().
	records := tsdbCol.QueryServerMessages(TimeWindow6Hours)
	if records == nil {
		// nil is fine (no data), but shouldn't panic
		_ = records
	}

	// Now close the actual TSDB.
	if err := tsdbCol.Close(); err != nil {
		t.Errorf("TSDBCollector.Close() failed: %v", err)
	}
}

// --------------------------------------------------------------------------
// Test 5: Messages recorded during loadFromTSDB are merged correctly
// --------------------------------------------------------------------------

func TestServerMessageCollector_ConcurrentRecordDuringLoad(t *testing.T) {
	tsdbCol, _ := newTestTSDB(t)
	defer tsdbCol.Close()

	// Pre-populate TSDB with some historical data.
	items := []ServerMessageItem{
		{ServerID: "ch1", Timestamp: time.Now().Add(-10 * time.Minute).UnixMilli()},
		{ServerID: "ch1", Timestamp: time.Now().Add(-9 * time.Minute).UnixMilli()},
		{ServerID: "ch2", Timestamp: time.Now().Add(-8 * time.Minute).UnixMilli()},
	}
	tsdbCol.BatchRecordServerMessages(items)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c := NewServerMessageCollector(tsdbCol)
	c.flushInterval = 100 * time.Millisecond
	c.Start(ctx)

	// Immediately record new messages (before load completes).
	c.RecordMessage("ch1")
	c.RecordMessage("ch3")

	// Wait for load + flush.
	waitForMetricsCollectorLoad(t, "server message collector", c.loaded.Load)

	summary := c.Summary()

	// Should have ch1 (2 historical + 1 live), ch2 (1 historical), ch3 (1 live).
	if len(summary) != 3 {
		t.Fatalf("expected 3 servers, got %d", len(summary))
	}

	totals := make(map[string]uint64)
	for _, s := range summary {
		totals[s.ServerID] = s.Total
	}

	if totals["ch1"] != 3 {
		t.Errorf("ch1 total = %d, want 3 (2 historical + 1 live)", totals["ch1"])
	}
	if totals["ch2"] != 1 {
		t.Errorf("ch2 total = %d, want 1", totals["ch2"])
	}
	if totals["ch3"] != 1 {
		t.Errorf("ch3 total = %d, want 1", totals["ch3"])
	}

	c.Stop()
}

// --------------------------------------------------------------------------
// Test 6: TSDB data directory is created and non-empty after flush
// --------------------------------------------------------------------------

func TestServerMessageCollector_TSDBDataOnDisk(t *testing.T) {
	dir := t.TempDir()
	dbPath := filepath.Join(dir, "test.tsdb")

	opts := tsdb.DefaultOptions()
	opts.RetentionDuration = int64(7 * 24 * time.Hour / time.Millisecond)
	opts.MinBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.MaxBlockDuration = int64(2 * time.Hour / time.Millisecond)
	opts.NoLockfile = true

	db, err := tsdb.Open(dbPath, nil, nil, opts, nil)
	if err != nil {
		t.Fatalf("failed to open TSDB: %v", err)
	}

	tc := &TSDBCollector{
		db:          db,
		stopCh:      make(chan struct{}),
		keyMaskMap:  make(map[string]string),
		maskMapPath: filepath.Join(dir, "key_mask_map.json"),
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c := NewServerMessageCollector(tc)
	c.flushInterval = 100 * time.Millisecond
	c.Start(ctx)

	for i := 0; i < 10; i++ {
		c.RecordMessage("ch1")
		if i < 9 {
			waitNextMetricsMillisecond(t)
		}
	}

	waitForMetricsCollectorLoad(t, "server message collector", c.loaded.Load)
	c.flush()
	c.Stop()
	tc.Close()

	// Verify TSDB directory exists and has data files.
	entries, err := os.ReadDir(dbPath)
	if err != nil {
		t.Fatalf("failed to read TSDB directory: %v", err)
	}
	if len(entries) == 0 {
		t.Error("TSDB directory should not be empty after flush")
	}

	// Reopen TSDB and verify data.
	db2, err := tsdb.Open(dbPath, nil, nil, opts, nil)
	if err != nil {
		t.Fatalf("failed to reopen TSDB: %v", err)
	}
	defer db2.Close()

	tc2 := &TSDBCollector{
		db:          db2,
		stopCh:      make(chan struct{}),
		keyMaskMap:  make(map[string]string),
		maskMapPath: filepath.Join(dir, "key_mask_map.json"),
	}

	records := tc2.QueryServerMessages(TimeWindow6Hours)
	if len(records) != 10 {
		t.Errorf("reopened TSDB: expected 10 records, got %d", len(records))
	}
}