package central_thanos

import (
	"sync"
	"sync/atomic"
	"testing"
)

func TestExactClusterMatcher(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		want        string
	}{
		{name: "dedicated uses cluster_id", servicePlan: "Dedicated", want: `cluster_id="tidb123"`},
		{name: "premium uses tidb_cluster_id", servicePlan: "Premium", want: `tidb_cluster_id="123"`},
		{name: "essential v2 uses premium tidb_cluster_id", servicePlan: "Essential_V2", want: `tidb_cluster_id="123"`},
		{name: "essential uses cluster_id", servicePlan: "Essential", want: `cluster_id="123"`},
		{name: "starter uses cluster_id", servicePlan: "Starter", want: `cluster_id="123"`},
		{name: "empty plan falls back to dedicated", servicePlan: "", want: `cluster_id="tidb123"`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := exactClusterMatcherForServicePlan(tt.servicePlan, 123); got != tt.want {
				t.Fatalf("exactClusterMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestRegexClusterMatcher(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		want        string
	}{
		{name: "dedicated uses cluster_id regex", servicePlan: "Dedicated", want: `cluster_id=~".*123"`},
		{name: "premium uses tidb_cluster_id regex", servicePlan: "Premium", want: `tidb_cluster_id=~".*123"`},
		{name: "essential v2 uses premium tidb_cluster_id regex", servicePlan: "Essential_V2", want: `tidb_cluster_id=~".*123"`},
		{name: "essential uses cluster_id regex", servicePlan: "Essential", want: `cluster_id=~".*123"`},
		{name: "starter uses cluster_id regex", servicePlan: "Starter", want: `cluster_id=~".*123"`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := regexClusterMatcherForServicePlan(tt.servicePlan, 123); got != tt.want {
				t.Fatalf("regexClusterMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestSharedSeriesMatcherForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		keyspaceID  string
		want        string
		wantErr     bool
	}{
		{name: "premium uses tidb_cluster_id", servicePlan: "Premium", want: `tidb_cluster_id="123"`},
		{name: "essential v2 uses premium tidb_cluster_id", servicePlan: "Essential_V2", want: `tidb_cluster_id="123"`},
		{name: "essential uses keyspace_id", servicePlan: "Essential", keyspaceID: "555755", want: `keyspace_id="555755"`},
		{name: "starter uses keyspace_id", servicePlan: "Starter", keyspaceID: "555750", want: `keyspace_id="555750"`},
		{name: "essential requires keyspace", servicePlan: "Essential", wantErr: true},
		{name: "dedicated falls back to legacy regex matcher", servicePlan: "Dedicated", want: `cluster_id=~".*123"`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := sharedSeriesMatcherForServicePlan(tt.servicePlan, 123, tt.keyspaceID)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("sharedSeriesMatcherForServicePlan() error = nil, want error")
				}
				return
			}
			if err != nil {
				t.Fatalf("sharedSeriesMatcherForServicePlan() error = %v", err)
			}
			if got != tt.want {
				t.Fatalf("sharedSeriesMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestAggregateDBMatcherForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		want        string
	}{
		{name: "dedicated keeps legacy regex matcher", servicePlan: "Dedicated", want: `cluster_id=~".*123"`},
		{name: "premium uses tidb_cluster_id exact matcher", servicePlan: "Premium", want: `tidb_cluster_id="123"`},
		{name: "essential v2 uses premium tidb_cluster_id exact matcher", servicePlan: "Essential_V2", want: `tidb_cluster_id="123"`},
		{name: "essential uses cluster_id exact matcher", servicePlan: "Essential", want: `cluster_id="123"`},
		{name: "starter uses cluster_id exact matcher", servicePlan: "Starter", want: `cluster_id="123"`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := aggregateDBMatcherForServicePlan(tt.servicePlan, 123); got != tt.want {
				t.Fatalf("aggregateDBMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestTidbClusterRegexMatcherForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		want        string
	}{
		{name: "dedicated uses tidb_cluster_id regex", servicePlan: "Dedicated", want: `tidb_cluster_id=~".*123"`},
		{name: "premium uses tidb_cluster_id regex", servicePlan: "Premium", want: `tidb_cluster_id=~".*123"`},
		{name: "essential v2 uses premium tidb_cluster_id regex", servicePlan: "Essential_V2", want: `tidb_cluster_id=~".*123"`},
		{name: "essential uses cluster_id regex", servicePlan: "Essential", want: `cluster_id=~".*123"`},
		{name: "starter uses cluster_id regex", servicePlan: "Starter", want: `cluster_id=~".*123"`},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := tidbClusterRegexMatcherForServicePlan(tt.servicePlan, 123); got != tt.want {
				t.Fatalf("tidbClusterRegexMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestRUMatcherForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		clusterID   uint64
		keyspaceID  string
		want        string
		wantErr     bool
	}{
		{name: "dedicated uses legacy tidb cluster regex", servicePlan: "Dedicated", clusterID: 123, want: `tidb_cluster_id=~".*123"`},
		{name: "premium uses tidb cluster regex", servicePlan: "Premium", clusterID: 123, want: `tidb_cluster_id=~".*123"`},
		{name: "essential v2 uses premium tidb cluster regex", servicePlan: "Essential_V2", clusterID: 123, want: `tidb_cluster_id=~".*123"`},
		{name: "essential uses resource_group keyspace", servicePlan: "Essential", clusterID: 123, keyspaceID: "556846", want: `resource_group="556846"`},
		{name: "starter uses resource_group keyspace", servicePlan: "Starter", clusterID: 123, keyspaceID: "556849", want: `resource_group="556849"`},
		{name: "essential requires keyspace", servicePlan: "Essential", clusterID: 123, wantErr: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ruMatcherForServicePlan(tt.servicePlan, tt.clusterID, tt.keyspaceID)
			if tt.wantErr {
				if err == nil {
					t.Fatalf("ruMatcherForServicePlan() error = nil, want error")
				}
				return
			}
			if err != nil {
				t.Fatalf("ruMatcherForServicePlan() error = %v", err)
			}
			if got != tt.want {
				t.Fatalf("ruMatcherForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestBuildServerlessFrontendRUReadPromQL(t *testing.T) {
	want := `(sum(rate(resource_manager_resource_unit_read_request_unit_sum{resource_group="abc123"}[1m])) or vector(0))`
	if got := buildServerlessFrontendRUReadPromQL(`resource_group="abc123"`); got != want {
		t.Fatalf("buildServerlessFrontendRUReadPromQL() = %q, want %q", got, want)
	}
}

func TestBuildServerlessFrontendRUWritePromQL(t *testing.T) {
	want := `(sum(rate(resource_manager_resource_unit_write_request_unit_sum{resource_group="abc123"}[1m])) or vector(0))`
	if got := buildServerlessFrontendRUWritePromQL(`resource_group="abc123"`); got != want {
		t.Fatalf("buildServerlessFrontendRUWritePromQL() = %q, want %q", got, want)
	}
}

func TestBuildServerlessFrontendRUTotalPromQL(t *testing.T) {
	want := `(sum(rate(resource_manager_resource_unit_read_request_unit_sum{resource_group="abc123"}[1m])) or vector(0)) + ` +
		`(sum(rate(resource_manager_resource_unit_write_request_unit_sum{resource_group="abc123"}[1m])) or vector(0)) + ` +
		`(sum(rate(resource_manager_resource_unit_sql_layer_request_unit{resource_group="abc123"}[1m])) or vector(0)) + ` +
		`(sum(rate(resource_manager_resource_unit_sql_layer_request_unit_sum{resource_group="abc123"}[1m])) or vector(0)) + ` +
		`(sum(rate(gateway_connection_flow_bytes_sum{namespace="tidb-admin",direction="out",resource_group="abc123"}[1m])/(1024)) or vector(0)) + ` +
		`(sum(rate(sum(tikv_worker_load_data_wru_cost_counter{resource_group="abc123"})[1m])) or vector(0)) + ` +
		`(sum(increase(sum(export_ru_cost_counter{resource_group="abc123"})[120s])/120) or vector(0))`
	if got := buildServerlessFrontendRUTotalPromQL(`resource_group="abc123"`); got != want {
		t.Fatalf("buildServerlessFrontendRUTotalPromQL() = %q, want %q", got, want)
	}
}

func TestBuildRUTotalPromQL(t *testing.T) {
	legacyMatcher := `tidb_cluster_id=~".*123"`
	legacyWant := `sum by (resource_group) (rate(resource_manager_resource_unit_read_request_unit_sum{tidb_cluster_id=~".*123"}[1m])) + sum by (resource_group) (rate(resource_manager_resource_unit_write_request_unit_sum{tidb_cluster_id=~".*123"}[1m]))`
	if got := buildRUTotalPromQL("Premium", legacyMatcher); got != legacyWant {
		t.Fatalf("buildRUTotalPromQL(Premium) = %q, want %q", got, legacyWant)
	}

	if got := buildRUTotalPromQL("Essential", `resource_group="557850"`); got != buildServerlessFrontendRUTotalPromQL(`resource_group="557850"`) {
		t.Fatalf("buildRUTotalPromQL(Essential) = %q, want frontend serverless formula", got)
	}

	if got := buildRUTotalPromQL("Starter", `resource_group="557850"`); got != buildServerlessFrontendRUTotalPromQL(`resource_group="557850"`) {
		t.Fatalf("buildRUTotalPromQL(Starter) = %q, want frontend serverless formula", got)
	}
}

func TestBuildNodeMemoryUsedBytesPromQLDeduplicatesByJoinKey(t *testing.T) {
	matcher := `cluster_id="tidb123"`
	want := `max by (component, instance) (node_memory_MemTotal_bytes{cluster_id="tidb123"}) - max by (component, instance) (node_memory_MemAvailable_bytes{cluster_id="tidb123"})`
	if got := buildNodeMemoryUsedBytesPromQL(matcher); got != want {
		t.Fatalf("buildNodeMemoryUsedBytesPromQL() = %q, want %q", got, want)
	}
}

func TestBuildNextGenHistoricalKeyspaceLookupPromQL(t *testing.T) {
	got := buildNextGenHistoricalKeyspaceLookupPromQL(`cluster_id="10194534416767893615"`)
	want := `topk(1, last_over_time(tidb_server_connections{cluster_id="10194534416767893615"}[24h])) or topk(1, last_over_time(tidb_server_tokens{cluster_id="10194534416767893615"}[24h])) or topk(1, last_over_time(tidb_server_handle_query_duration_seconds_sum{cluster_id="10194534416767893615"}[24h]))`
	if got != want {
		t.Fatalf("buildNextGenHistoricalKeyspaceLookupPromQL() = %q, want %q", got, want)
	}
}

func TestGetOrResolveNextGenKeyspaceIDCachesResolvedValue(t *testing.T) {
	resetNextGenKeyspaceIDStateForTest()

	var calls atomic.Int32
	resolve := func() (string, error) {
		calls.Add(1)
		return "abc123", nil
	}

	got1, err := getOrResolveNextGenKeyspaceID(123, resolve)
	if err != nil {
		t.Fatalf("getOrResolveNextGenKeyspaceID() first call error = %v", err)
	}
	got2, err := getOrResolveNextGenKeyspaceID(123, resolve)
	if err != nil {
		t.Fatalf("getOrResolveNextGenKeyspaceID() second call error = %v", err)
	}

	if got1 != "abc123" || got2 != "abc123" {
		t.Fatalf("getOrResolveNextGenKeyspaceID() = %q, %q, want both abc123", got1, got2)
	}
	if calls.Load() != 1 {
		t.Fatalf("resolve call count = %d, want 1", calls.Load())
	}
}

func TestGetOrResolveNextGenKeyspaceIDDeduplicatesConcurrentLookups(t *testing.T) {
	resetNextGenKeyspaceIDStateForTest()

	var calls atomic.Int32
	release := make(chan struct{})
	resolve := func() (string, error) {
		calls.Add(1)
		<-release
		return "rg-4", nil
	}

	const workers = 3
	results := make(chan string, workers)
	errs := make(chan error, workers)
	var wg sync.WaitGroup
	wg.Add(workers)
	for range workers {
		go func() {
			defer wg.Done()
			got, err := getOrResolveNextGenKeyspaceID(456, resolve)
			if err != nil {
				errs <- err
				return
			}
			results <- got
		}()
	}

	close(release)
	wg.Wait()
	close(results)
	close(errs)

	for err := range errs {
		t.Fatalf("getOrResolveNextGenKeyspaceID() error = %v", err)
	}
	for got := range results {
		if got != "rg-4" {
			t.Fatalf("getOrResolveNextGenKeyspaceID() = %q, want rg-4", got)
		}
	}
	if calls.Load() != 1 {
		t.Fatalf("resolve call count = %d, want 1", calls.Load())
	}
}

func resetNextGenKeyspaceIDStateForTest() {
	nextGenKeyspaceIDCache.Range(func(key, _ any) bool {
		nextGenKeyspaceIDCache.Delete(key)
		return true
	})
	nextGenKeyspaceIDInflight.Range(func(key, _ any) bool {
		nextGenKeyspaceIDInflight.Delete(key)
		return true
	})
}

func TestUseClusterLevelDBLabelsForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		want        bool
	}{
		{name: "dedicated keeps node-level grouping", servicePlan: "Dedicated", want: false},
		{name: "premium uses instance-level grouping", servicePlan: "Premium", want: true},
		{name: "essential uses instance-level grouping", servicePlan: "Essential", want: true},
		{name: "starter uses instance-level grouping", servicePlan: "Starter", want: true},
		{name: "empty plan falls back to dedicated behavior", servicePlan: "", want: false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := useClusterLevelDBLabelsForServicePlan(tt.servicePlan); got != tt.want {
				t.Fatalf("useClusterLevelDBLabelsForServicePlan() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestDBGroupLabelsForServicePlan(t *testing.T) {
	tests := []struct {
		name        string
		servicePlan string
		extra       []string
		want        string
	}{
		{
			name:        "dedicated keeps component and instance",
			servicePlan: "Dedicated",
			extra:       []string{"sql_type"},
			want:        "component, instance, sql_type",
		},
		{
			name:        "premium removes component and instance",
			servicePlan: "Premium",
			extra:       []string{"sql_type"},
			want:        "sql_type",
		},
		{
			name:        "essential removes component and instance",
			servicePlan: "Essential",
			extra:       []string{"type"},
			want:        "type",
		},
		{
			name:        "starter removes component and instance",
			servicePlan: "Starter",
			extra:       []string{"type", "txn_mode"},
			want:        "type, txn_mode",
		},
		{
			name:        "no extra labels for nextgen collapses to pure sum",
			servicePlan: "Premium",
			extra:       nil,
			want:        "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := dbGroupLabelsForServicePlan(tt.servicePlan, tt.extra...); got != tt.want {
				t.Fatalf("dbGroupLabelsForServicePlan() = %q, want %q", got, tt.want)
			}
		})
	}
}