package metrics

import (
	"context"
	"math"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"
)

type countingClusterIDProvider struct {
	mu        sync.Mutex
	clusterID map[string]string
	err       error
	calls     int
}

func (p *countingClusterIDProvider) ClusterID(ctx context.Context, keyspaceName string) (string, error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.calls++
	if p.err != nil {
		return "", p.err
	}
	return p.clusterID[keyspaceName], nil
}

func (p *countingClusterIDProvider) CallCount() int {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.calls
}

func TestRequestTrackingEnabled(t *testing.T) {
	SetClusterIDProviderForTest(nil)
	Init(Settings{Enabled: false}, nil)
	if RequestTrackingEnabled() {
		t.Fatal("expected tracking disabled when metrics and slow threshold are unset")
	}

	Init(Settings{Enabled: false, SlowRequestThreshold: 10 * time.Millisecond}, nil)
	if !RequestTrackingEnabled() {
		t.Fatal("expected tracking enabled when slow threshold is set")
	}
	if SlowThreshold() != 10*time.Millisecond {
		t.Fatalf("expected slow threshold 10ms, got %s", SlowThreshold())
	}

	Init(Settings{Enabled: true}, nil)
	if !Enabled() {
		t.Fatal("expected metrics enabled")
	}
}

func TestRequestMetricsContext(t *testing.T) {
	SetClusterIDProviderForTest(nil)
	start := time.Unix(0, 0)
	rm := NewRequestMetrics(5, 12, start)
	ctx := WithRequestMetrics(context.Background(), rm)

	ObserveStageFromContext(ctx, StageHandler, 5*time.Millisecond)
	if got := rm.StageDuration(StageHandler); got != 5*time.Millisecond {
		t.Fatalf("expected handler duration 5ms, got %s", got)
	}

	SetRequestTarget(ctx, "ks", "tbl")
	ks, tbl := rm.Target()
	if ks != "ks" || tbl != "tbl" {
		t.Fatalf("expected target ks/tbl, got %q/%q", ks, tbl)
	}

	rm.SetFrameReadBreakdown(1234, 12*time.Millisecond, 34*time.Millisecond)
	requestBytes, firstByteWait, bodyRead := rm.FrameReadBreakdown()
	if requestBytes != 1234 {
		t.Fatalf("expected request bytes 1234, got %d", requestBytes)
	}
	if firstByteWait != 12*time.Millisecond {
		t.Fatalf("expected first byte wait 12ms, got %s", firstByteWait)
	}
	if bodyRead != 34*time.Millisecond {
		t.Fatalf("expected body read 34ms, got %s", bodyRead)
	}

	RecordRequestSQL(ctx, "SELECT  \n  *  FROM foo WHERE id = ?", "abc")
	sqlCount, sqlBytes, dataBytes, omitted, sqls := rm.SQLSummary()
	if sqlCount != 1 {
		t.Fatalf("expected sql count 1, got %d", sqlCount)
	}
	if sqlBytes != len("SELECT  \n  *  FROM foo WHERE id = ?") {
		t.Fatalf("expected sql bytes %d, got %d", len("SELECT  \n  *  FROM foo WHERE id = ?"), sqlBytes)
	}
	if dataBytes != 3 {
		t.Fatalf("expected data bytes 3, got %d", dataBytes)
	}
	if omitted != 0 {
		t.Fatalf("expected omitted 0, got %d", omitted)
	}
	if len(sqls) != 1 || sqls[0] != "SELECT * FROM foo WHERE id = ?" {
		t.Fatalf("unexpected sql previews: %#v", sqls)
	}

	SetRequestBatchInfo(ctx, "UNLOGGED", 1024, 4, 4, 1024)
	batchType, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount := rm.BatchInfo()
	if batchType != "UNLOGGED" || statementCount != 1024 || executeItemCount != 4 || mergedGroupCount != 4 || mergedStmtCount != 1024 {
		t.Fatalf("unexpected batch info: %q %d %d %d %d", batchType, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount)
	}
}

func TestRequestMetricsWithoutDiagnosticsSkipsSlowLogDetails(t *testing.T) {
	SetClusterIDProviderForTest(nil)
	start := time.Unix(0, 0)
	rm := NewRequestMetricsWithDiagnostics(5, 12, start, false)
	ctx := WithRequestMetrics(context.Background(), rm)

	ObserveStageFromContext(ctx, StageHandler, 5*time.Millisecond)
	if got := rm.StageDuration(StageHandler); got != 0 {
		t.Fatalf("expected handler duration 0 when diagnostics disabled, got %s", got)
	}

	SetRequestTarget(ctx, "ks", "tbl")
	if ks, tbl := rm.Target(); ks != "" || tbl != "" {
		t.Fatalf("expected empty target when diagnostics disabled, got %q/%q", ks, tbl)
	}

	RecordRequestSQL(ctx, "SELECT * FROM foo WHERE id = ?", "abc")
	sqlCount, sqlBytes, dataBytes, omitted, sqls := rm.SQLSummary()
	if sqlCount != 0 || sqlBytes != 0 || dataBytes != 0 || omitted != 0 || len(sqls) != 0 {
		t.Fatalf("expected empty SQL summary when diagnostics disabled, got count=%d sqlBytes=%d dataBytes=%d omitted=%d sqls=%v", sqlCount, sqlBytes, dataBytes, omitted, sqls)
	}

	SetRequestBatchInfo(ctx, "UNLOGGED", 1024, 4, 4, 1024)
	batchType, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount := rm.BatchInfo()
	if batchType != "" || statementCount != 0 || executeItemCount != 0 || mergedGroupCount != 0 || mergedStmtCount != 0 {
		t.Fatalf("expected empty batch info when diagnostics disabled: %q %d %d %d %d", batchType, statementCount, executeItemCount, mergedGroupCount, mergedStmtCount)
	}

	rm.SetFrameReadBreakdown(1234, 12*time.Millisecond, 34*time.Millisecond)
	requestBytes, firstByteWait, bodyRead := rm.FrameReadBreakdown()
	if requestBytes != 1234 || firstByteWait != 12*time.Millisecond || bodyRead != 34*time.Millisecond {
		t.Fatalf("expected frame read breakdown to remain available, got %d %s %s", requestBytes, firstByteWait, bodyRead)
	}
}

func TestMetricsHandlerExportsPoolStats(t *testing.T) {
	SetClusterIDProviderForTest(nil)
	Init(Settings{Enabled: true}, func() []PoolStat {
		return []PoolStat{
			{TenantID: "t1", User: "u1", KeyspaceName: "ks1", ClusterID: "c1", Open: 2, InUse: 1, WaitCount: 3, WaitDuration: 5 * time.Second},
		}
	})

	h := Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()
	if !strings.Contains(body, "proxy_mysql_pool_open") {
		t.Fatalf("expected pool stats in metrics output, got: %s", body)
	}
	for _, label := range []string{`KeyspaceName="ks1"`, `ClusterId="c1"`} {
		if !strings.Contains(body, label) {
			t.Fatalf("expected %s in metrics output, got: %s", label, body)
		}
	}
}

func TestMetricsHandlerExportsPreparedStmtCacheStats(t *testing.T) {
	provider := &countingClusterIDProvider{clusterID: map[string]string{"ks": "123"}}
	SetClusterIDProviderForTest(provider)
	Init(Settings{Enabled: true}, nil)
	SetPreparedStmtCacheStatsProvider(func() []PreparedStmtCacheStat {
		return []PreparedStmtCacheStat{
			{User: "ks.user", Size: 2, Hits: 5, Misses: 3, Evictions: 1, Errors: 4},
		}
	})

	h := Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()

	for _, metricName := range []string{
		"proxy_mysql_prepared_stmt_cache_size",
		"proxy_mysql_prepared_stmt_cache_hits_total",
		"proxy_mysql_prepared_stmt_cache_misses_total",
		"proxy_mysql_prepared_stmt_cache_evictions_total",
		"proxy_mysql_prepared_stmt_cache_errors_total",
	} {
		if !strings.Contains(body, metricName) {
			t.Fatalf("expected %s in metrics output, got: %s", metricName, body)
		}
	}
	for _, label := range []string{
		`user="ks.user"`,
		`KeyspaceName="ks"`,
		`ClusterId="123"`,
	} {
		if !strings.Contains(body, label) {
			t.Fatalf("expected %s in metrics output, got: %s", label, body)
		}
	}
	if provider.CallCount() != 1 {
		t.Fatalf("expected cluster id provider to be called once, got %d", provider.CallCount())
	}
}

func TestMetricsHandlerExportsConnectionStatsAndDisconnections(t *testing.T) {
	provider := &countingClusterIDProvider{clusterID: map[string]string{"ks": "123"}}
	SetClusterIDProviderForTest(provider)
	Init(Settings{Enabled: true}, nil)
	SetConnectionStatsProvider(func() []ConnectionStat {
		return []ConnectionStat{
			{Scope: "all", KeyspaceName: "ks", ClusterID: "123", Count: 2},
			{Scope: "active", KeyspaceName: "ks", ClusterID: "123", Count: 1},
		}
	})
	RecordCQLDisconnection("ks.user")

	h := Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()

	for _, metricName := range []string{
		"proxy_cql_connection_count",
		"proxy_cql_disconnection_total",
	} {
		if !strings.Contains(body, metricName) {
			t.Fatalf("expected %s in metrics output, got: %s", metricName, body)
		}
	}
	for _, label := range []string{
		`scope="all"`,
		`scope="active"`,
		`KeyspaceName="ks"`,
		`ClusterId="123"`,
	} {
		if !strings.Contains(body, label) {
			t.Fatalf("expected %s in metrics output, got: %s", label, body)
		}
	}
	if provider.CallCount() != 1 {
		t.Fatalf("expected cluster id provider to be called once, got %d", provider.CallCount())
	}
}

func TestMetricsHandlerExportsRequestAndMySQLMetrics(t *testing.T) {
	provider := &countingClusterIDProvider{clusterID: map[string]string{"ks": "123"}}
	SetClusterIDProviderForTest(provider)
	Init(Settings{Enabled: true}, nil)
	rm := NewRequestMetrics(0x07, 1, time.Now())
	ctx := WithRequestMetrics(WithUsername(context.Background(), "ks.user"), rm)
	IncRequests(0x07, rm)
	IncInFlightRequests(rm)
	ObserveStageFromContext(ctx, StageBatchParse, time.Millisecond)
	ObserveStageFromContext(ctx, StageBatchDecode, 2*time.Millisecond)
	ObserveStageFromContext(ctx, StageBatchMergePlan, 3*time.Millisecond)
	ObserveStageFromContext(ctx, StageBatchExecuteWait, 4*time.Millisecond)
	ObserveCQLQuery(ctx, "select", true, 2*time.Millisecond)
	ObserveRequestCompleted(7, true, 2*time.Millisecond, rm)
	ObserveRequestCompleted(7, false, 4*time.Millisecond, rm)
	ObserveMySQLOperation(ctx, MySQLOpQueryRead, true, 1*time.Millisecond)
	ObserveMySQLOperation(ctx, MySQLOpExecWrite, false, 3*time.Millisecond)
	RecordMySQLError(ctx, "read", MySQLErrorTimeout)
	ObserveCQLQuery(ctx, "insert", false, 6*time.Millisecond)
	ObserveCQLWriteRows(ctx, "insert", 3)
	ObserveCQLWriteRows(ctx, "update", 2)
	ObserveCQLWriteRows(ctx, "batch", 5)
	ObserveCQLWriteRows(ctx, "select", 7)
	DecInFlightRequests(rm)

	unknownRM := NewRequestMetrics(0x42, 2, time.Now())
	IncRequests(0x42, unknownRM)
	ObserveRequestCompleted(0x42, true, 5*time.Millisecond, unknownRM)
	ObserveRequestCompleted(0x05, false, 3*time.Millisecond, NewRequestMetrics(0x05, 3, time.Now()))

	h := Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()

	for _, metricName := range []string{
		"proxy_request_in_flight",
		"proxy_request_completed_total",
		"proxy_request_duration_seconds",
		"proxy_request_stage_seconds",
		"proxy_cql_query_total",
		"proxy_cql_query_duration_seconds",
		"proxy_cql_write_rows_total",
		"proxy_mysql_operation_total",
		"proxy_mysql_operation_seconds",
		"proxy_mysql_error_total",
	} {
		if !strings.Contains(body, metricName) {
			t.Fatalf("expected %s in metrics output, got: %s", metricName, body)
		}
	}
	for _, label := range []string{
		`opcode="QUERY"`,
		`op_type="select"`,
		`op_type="OPTIONS"`,
		`operation_type="insert"`,
		`operation_type="update"`,
		`operation_type="batch"`,
		`stage="batch_parse"`,
		`stage="batch_decode"`,
		`stage="batch_merge_plan"`,
		`stage="batch_execute_wait"`,
		`KeyspaceName="ks"`,
		`ClusterId="123"`,
	} {
		if !strings.Contains(body, label) {
			t.Fatalf("expected %s in metrics output, got: %s", label, body)
		}
	}
	if strings.Contains(body, `opcode="7"`) {
		t.Fatalf("expected numeric opcode labels to be replaced, got: %s", body)
	}
	if strings.Contains(body, `UNKNOWN_OP_`) {
		t.Fatalf("expected unknown opcode labels to be omitted, got: %s", body)
	}
	if strings.Contains(body, `operation_type="select"`) {
		t.Fatalf("expected select write rows to be ignored, got: %s", body)
	}
	if provider.CallCount() != 1 {
		t.Fatalf("expected cluster id provider to be called once, got %d", provider.CallCount())
	}
}

func TestInFlightRequestsUseIncrementLabelsWhenUsernameChanges(t *testing.T) {
	provider := &countingClusterIDProvider{clusterID: map[string]string{"ks": "123"}}
	SetClusterIDProviderForTest(provider)
	Init(Settings{Enabled: true}, nil)

	rm := NewRequestMetrics(0x0F, 1, time.Now())
	ctx := WithRequestMetrics(context.Background(), rm)
	IncInFlightRequests(rm)
	ctx = WithUsername(ctx, "ks.user")
	DecInFlightRequests(rm)

	h := Handler()
	if h == nil {
		t.Fatal("expected metrics handler when enabled")
	}

	rec := httptest.NewRecorder()
	req := httptest.NewRequest("GET", "/metrics", nil)
	h.ServeHTTP(rec, req)
	body := rec.Body.String()

	if !strings.Contains(body, `proxy_request_in_flight{ClusterId="-1",KeyspaceName="unknown"} 0`) {
		t.Fatalf("expected original in-flight label series to return to zero, got: %s", body)
	}
	if strings.Contains(body, `proxy_request_in_flight{ClusterId="123",KeyspaceName="ks"} -1`) {
		t.Fatalf("expected username update not to decrement a different in-flight series, got: %s", body)
	}
}

func TestResolveMetricLabelsCachesFailuresAndSuccesses(t *testing.T) {
	provider := &countingClusterIDProvider{
		clusterID: map[string]string{
			"ks_ok": "999",
		},
	}
	SetClusterIDProviderForTest(provider)

	labels := ResolveMetricLabels(context.Background(), "ks_ok.branch")
	if labels.KeyspaceName != "ks_ok" || labels.ClusterID != "999" {
		t.Fatalf("unexpected labels for success: %+v", labels)
	}
	labels = ResolveMetricLabels(context.Background(), "ks_ok.other")
	if provider.CallCount() != 1 {
		t.Fatalf("expected cached success lookup, got %d calls", provider.CallCount())
	}

	provider.err = context.DeadlineExceeded
	SetClusterIDProviderForTest(provider)
	labels = ResolveMetricLabels(context.Background(), "ks_fail.env")
	if labels.KeyspaceName != "ks_fail" || labels.ClusterID != unknownClusterID {
		t.Fatalf("unexpected labels for failure: %+v", labels)
	}
	labels = ResolveMetricLabels(context.Background(), "ks_fail.env2")
	if provider.CallCount() != 2 {
		t.Fatalf("expected cached failure lookup, got %d calls", provider.CallCount())
	}
}

func TestResolveMetricLabelsUsesMySQLHostPatternWhenMatched(t *testing.T) {
	provider := &countingClusterIDProvider{
		clusterID: map[string]string{
			"ignored": "999",
		},
	}
	SetClusterIDProviderForTest(provider)
	SetMetricLabelsFromMySQLHost("db-10477550877178126640-tiproxy.tidb2034467628784816128.svc.cluster.local:4000")

	labels := ResolveMetricLabels(context.Background(), "ignored.user")
	if labels.KeyspaceName != "10477550877178126640" || labels.ClusterID != "10477550877178126640" {
		t.Fatalf("unexpected labels from mysql host pattern: %+v", labels)
	}
	if provider.CallCount() != 0 {
		t.Fatalf("expected host-derived labels to skip cluster provider, got %d calls", provider.CallCount())
	}
}

func TestResolveMetricLabelsIgnoresMySQLHostPatternWhenUnmatched(t *testing.T) {
	provider := &countingClusterIDProvider{
		clusterID: map[string]string{
			"ks": "123",
		},
	}
	SetClusterIDProviderForTest(provider)
	SetMetricLabelsFromMySQLHost("mysql.example.com:4000")

	labels := ResolveMetricLabels(context.Background(), "ks.user")
	if labels.KeyspaceName != "ks" || labels.ClusterID != "123" {
		t.Fatalf("unexpected labels when mysql host does not match pattern: %+v", labels)
	}
	if provider.CallCount() != 1 {
		t.Fatalf("expected cluster id provider to be called once, got %d", provider.CallCount())
	}
}

func TestDurationBucketConfigs(t *testing.T) {
	SetClusterIDProviderForTest(nil)
	cases := []struct {
		name    string
		buckets []float64
		min     float64
		max     float64
		under5  int
	}{
		{name: "stage", buckets: stageDurationBuckets, min: 0.00005, max: 320, under5: 28},
		{name: "request", buckets: requestDurationBuckets, min: 0.00025, max: 640, under5: 20},
		{name: "cql", buckets: cqlQueryDurationBuckets, min: 0.00025, max: 640, under5: 20},
		{name: "mysql", buckets: mysqlOperationDurationBuckets, min: 0.0001, max: 960, under5: 21},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			if len(tc.buckets) < 20 {
				t.Fatalf("expected %s buckets to be reasonably granular, got %d", tc.name, len(tc.buckets))
			}
			if math.Abs(tc.buckets[0]-tc.min) > 1e-12 {
				t.Fatalf("unexpected first bucket for %s: got %v want %v", tc.name, tc.buckets[0], tc.min)
			}
			if math.Abs(tc.buckets[len(tc.buckets)-1]-tc.max) > 1e-12 {
				t.Fatalf("unexpected last bucket for %s: got %v want %v", tc.name, tc.buckets[len(tc.buckets)-1], tc.max)
			}
			under5 := 0
			for i := 1; i < len(tc.buckets); i++ {
				if tc.buckets[i] <= tc.buckets[i-1] {
					t.Fatalf("%s buckets must be strictly increasing: %v <= %v at %d", tc.name, tc.buckets[i], tc.buckets[i-1], i)
				}
			}
			for _, bucket := range tc.buckets {
				if bucket <= 0.005 {
					under5++
				}
			}
			if under5 < tc.under5 {
				t.Fatalf("expected at least %d buckets within 5ms for %s, got %d", tc.under5, tc.name, under5)
			}
		})
	}

	if reflect.DeepEqual(stageDurationBuckets, requestDurationBuckets) {
		t.Fatal("stage and request buckets should be tuned separately")
	}
	if reflect.DeepEqual(requestDurationBuckets, mysqlOperationDurationBuckets) {
		t.Fatal("request and mysql buckets should be tuned separately")
	}
}