package cmd

import (
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/tidbcloud/clinic-cli/internal/clinic"
	"github.com/tidbcloud/clinic-cli/internal/config"
)

func TestLokiQueryRangeRequiresClusterIDStartAndEnd(t *testing.T) {
	t.Parallel()

	ctx := &RuntimeContext{Client: &recordingClinicClient{}}
	testCases := []struct {
		name string
		args []string
		want string
	}{
		{
			name: "missing cluster id",
			args: []string{"query", "--start", "10", "--end", "20", `{app="tidb"}`},
			want: "missing required --cluster-id",
		},
		{
			name: "missing start",
			args: []string{"query", "--cluster-id", "cluster-42", "--end", "20", `{app="tidb"}`},
			want: "missing required --start",
		},
		{
			name: "missing end",
			args: []string{"query", "--cluster-id", "cluster-42", "--start", "10", `{app="tidb"}`},
			want: "missing required --end",
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			err := lokiCommandHandler(ctx, tc.args, &strings.Builder{})
			if err == nil {
				t.Fatal("lokiCommandHandler() error = nil, want validation error")
			}
			if got := err.Error(); got != tc.want {
				t.Fatalf("error = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestLokiInstanceMapRequiresClusterIDStartAndEnd(t *testing.T) {
	t.Parallel()

	ctx := &RuntimeContext{Client: &recordingClinicClient{}}
	testCases := []struct {
		name string
		args []string
		want string
	}{
		{
			name: "missing cluster id",
			args: []string{"instance-map", "--start", "10", "--end", "20"},
			want: "missing required --cluster-id",
		},
		{
			name: "missing start",
			args: []string{"instance-map", "--cluster-id", "cluster-42", "--end", "20"},
			want: "missing required --start",
		},
		{
			name: "missing end",
			args: []string{"instance-map", "--cluster-id", "cluster-42", "--start", "10"},
			want: "missing required --end",
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			err := lokiCommandHandler(ctx, tc.args, &strings.Builder{})
			if err == nil {
				t.Fatal("lokiCommandHandler() error = nil, want validation error")
			}
			if got := err.Error(); got != tc.want {
				t.Fatalf("error = %q, want %q", got, tc.want)
			}
		})
	}
}

func TestLokiQueryRangeUsesLatestFiveMinutesWhenStartEndCannotBeResolved(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	before := time.Now().Unix()
	err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--output", "json",
		`{app="tidb", cluster_env="staging", namespace="tidbcluster-42"}`,
	}, &strings.Builder{})
	if err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}
	after := time.Now().Unix()

	startRaw := client.request.Query.Get("start")
	endRaw := client.request.Query.Get("end")
	startUnix, parseErr := strconv.ParseInt(startRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("start query value = %q, want unix nanoseconds: %v", startRaw, parseErr)
	}
	endUnix, parseErr := strconv.ParseInt(endRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("end query value = %q, want unix nanoseconds: %v", endRaw, parseErr)
	}
	if got, want := endUnix-startUnix, int64((5 * time.Minute).Nanoseconds()); got != want {
		t.Fatalf("range length = %d, want %d", got, want)
	}
	endUnixSeconds := endUnix / int64(time.Second)
	if endUnixSeconds < before || endUnixSeconds > after {
		t.Fatalf("end query value = %d, want unix seconds between %d and %d", endUnix, before, after)
	}
}

func TestLokiQueryOptionalLogQLWithCobra(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name      string
		args      []string
		wantQuery string
	}{
		{
			name: "omitted logql builds default selector",
			args: []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				"-o", "json",
			},
			wantQuery: `{cluster_env="staging", namespace="tidbcluster-42"}`,
		},
		{
			name: "provided logql receives default selector labels",
			args: []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				"-o", "json",
				`{component="tidb"}`,
			},
			wantQuery: `{component="tidb", cluster_env="staging", namespace="tidbcluster-42"}`,
		},
		{
			name: "content-only selector receives default selector labels",
			args: []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				"-o", "json",
				`{} |= "Request range exceeds bound"`,
			},
			wantQuery: `{cluster_env="staging", namespace="tidbcluster-42"} |= "Request range exceeds bound"`,
		},
		{
			name: "can skip namespace injection for selector queries",
			args: []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				"--skip-namespace-injection",
				"-o", "json",
				`{component="tidb"}`,
			},
			wantQuery: `{component="tidb", cluster_env="staging"}`,
		},
		{
			name: "can skip namespace injection for inferred default selector",
			args: []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				"--skip-namespace-injection",
				"-o", "json",
			},
			wantQuery: `{cluster_env="staging"}`,
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			client := &recordingClinicClient{
				responses: []clinicResponse{
					{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
				},
			}
			ctx := &RuntimeContext{
				Client:      client,
				SelectedEnv: "staging",
				SelectedProfile: &config.ProfileRecord{
					ClusterID:           "cluster-42",
					OrgID:               "org-1",
					ClusterDeployTypeV2: "dedicated",
				},
			}

			if err := lokiCommandHandler(ctx, tc.args, &strings.Builder{}); err != nil {
				t.Fatalf("lokiCommandHandler() error = %v", err)
			}

			assertQueryValue(t, client.request.Query, "query", tc.wantQuery)
			assertQueryValue(t, client.request.Query, "limit", "10")
		})
	}
}

func TestLokiInstanceMapUsesLatestFiveMinutesWhenStartEndCannotBeResolved(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	before := time.Now().Unix()
	err := lokiCommandHandler(ctx, []string{
		"instance-map",
		"--cluster-id", "cluster-42",
		`--output`, "table",
	}, &strings.Builder{})
	if err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}
	after := time.Now().Unix()

	startRaw := client.request.Query.Get("start")
	endRaw := client.request.Query.Get("end")
	startUnix, parseErr := strconv.ParseInt(startRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("start query value = %q, want unix nanoseconds: %v", startRaw, parseErr)
	}
	endUnix, parseErr := strconv.ParseInt(endRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("end query value = %q, want unix nanoseconds: %v", endRaw, parseErr)
	}
	if got, want := endUnix-startUnix, int64((5 * time.Minute).Nanoseconds()); got != want {
		t.Fatalf("range length = %d, want %d", got, want)
	}
	endUnixSeconds := endUnix / int64(time.Second)
	if endUnixSeconds < before || endUnixSeconds > after {
		t.Fatalf("end query value = %d, want unix seconds between %d and %d", endUnix, before, after)
	}
}

func TestLokiInstanceMapBuildsDefaultQueryAndDefaultLimit(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	if err := lokiCommandHandler(ctx, []string{
		"instance-map",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600",
	}, &strings.Builder{}); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := client.request.Method, http.MethodGet; got != want {
		t.Fatalf("method = %q, want %q", got, want)
	}
	if got, want := client.request.Path, "/data-proxy/loki/cluster-42/api/v1/query_range"; got != want {
		t.Fatalf("path = %q, want %q", got, want)
	}
	assertQueryValue(t, client.request.Query, "query", `{cluster_env="staging", namespace="tidbcluster-42"}`)
	assertQueryValue(t, client.request.Query, "start", "1712123000000000000")
	assertQueryValue(t, client.request.Query, "end", "1712123600000000000")
	assertQueryValue(t, client.request.Query, "limit", "1000")
}

func TestLokiInstanceMapRetriesGatewayTimeout(t *testing.T) {
	t.Parallel()

	t.Run("retries up to success for default output", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					body: []byte(`{"status":"success","data":{"resultType":"streams","result":[{"stream":{"node":"cn-hk.10.0.7.198","instance":"worker-tidb-0"},"values":[["1712123600000000000","line"]]}]}}`),
				},
			},
		}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		var stdout strings.Builder
		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
		}, &stdout)
		if err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}
		if got, want := len(client.requests), 3; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}

		rows := strings.Split(strings.TrimSpace(stdout.String()), "\n")
		if got, want := len(rows), 2; got != want {
			t.Fatalf("row count = %d, want %d", got, want)
		}
		header := strings.Fields(rows[0])
		wantHeader := []string{"NODE", "DATA_SERVER_INSTANCE", "LOKI_INSTANCE", "IP"}
		if got, want := strings.Join(header, "|"), strings.Join(wantHeader, "|"); got != want {
			t.Fatalf("table header = %q, want %q", got, want)
		}
	})

	t.Run("retries up to max then returns last 504 error", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout #1",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout #2",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout #3",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout #4",
					},
				},
			},
		}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want gateway timeout error")
		}

		var httpErr *clinic.HTTPError
		if ok := errors.As(err, &httpErr); !ok {
			t.Fatalf("error type = %T, want *clinic.HTTPError", err)
		}
		if got, want := httpErr.StatusCode, http.StatusGatewayTimeout; got != want {
			t.Fatalf("status code = %d, want %d", got, want)
		}
		if got, want := httpErr.Status, "504 Gateway Timeout #4"; got != want {
			t.Fatalf("status = %q, want %q", got, want)
		}
		if got, want := len(client.requests), 4; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
	})

	t.Run("does not retry for non-504 errors", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusBadGateway,
						Status:     "502 Bad Gateway",
					},
				},
			},
		}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want non-504 error")
		}
		var httpErr *clinic.HTTPError
		if ok := errors.As(err, &httpErr); !ok {
			t.Fatalf("error type = %T, want *clinic.HTTPError", err)
		}
		if got, want := httpErr.StatusCode, http.StatusBadGateway; got != want {
			t.Fatalf("status code = %d, want %d", got, want)
		}
		if got, want := httpErr.Status, "502 Bad Gateway"; got != want {
			t.Fatalf("status = %q, want %q", got, want)
		}
		if got, want := len(client.requests), 1; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
	})

	t.Run("retries for json output mode as well", func(t *testing.T) {
		t.Parallel()

		rawBody := `{"status":"success","data":{"resultType":"streams","result":[{"stream":{"node":"cn-hk.10.0.7.198","instance":"worker-tidb-0"},"values":[["1712123600000000000","line"]]}]}}`
		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					body: []byte(rawBody),
				},
			},
		}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		var stdout strings.Builder
		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			"--output", "json",
		}, &stdout)
		if err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}
		if got, want := len(client.requests), 2; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
		assertJSONEqual(t, stdout.String(), `[
			{"node":"cn-hk.10.0.7.198","data_server_instance":"cn-hk.10.0.7.198","loki_instance":"worker-tidb-0","ip":"10.0.7.198"}
		]`)
	})
}

func TestLokiInstanceMapSupportsOutputModesAndRejectsInvalid(t *testing.T) {
	t.Parallel()

	rawBody := `{"status":"success","data":{"resultType":"streams","result":[]}}`
	validOutputTests := []struct {
		name         string
		outputArg    string
		outputAlias  string
		expectedPath string
	}{
		{
			name:         "supports -o table",
			outputArg:    "-o",
			outputAlias:  "table",
			expectedPath: "/data-proxy/loki/cluster-42/api/v1/query_range",
		},
		{
			name:         "supports --output json",
			outputArg:    "--output",
			outputAlias:  "json",
			expectedPath: "/data-proxy/loki/cluster-42/api/v1/query_range",
		},
	}

	for _, tc := range validOutputTests {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			client := &recordingClinicClient{
				responses: []clinicResponse{
					{body: []byte(rawBody)},
				},
			}
			ctx := &RuntimeContext{
				Client:      client,
				SelectedEnv: "staging",
				SelectedProfile: &config.ProfileRecord{
					ClusterID:           "cluster-42",
					OrgID:               "org-1",
					ClusterDeployTypeV2: "dedicated",
				},
			}
			var stdout strings.Builder
			err := lokiCommandHandler(ctx, []string{
				"instance-map",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600",
				tc.outputArg, tc.outputAlias,
			}, &stdout)
			if err != nil {
				t.Fatalf("lokiCommandHandler() error = %v", err)
			}
			if got, want := client.request.Path, tc.expectedPath; got != want {
				t.Fatalf("path = %q, want %q", got, want)
			}
			if tc.outputAlias == "json" {
				assertJSONEqual(t, stdout.String(), `[]`)
				return
			}
			if got := stdout.String(); got == "" {
				t.Fatalf("expected non-empty output for -o table, got empty output")
			}
		})
	}

	t.Run("rejects invalid output format", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		ctx.SelectedEnv = "staging"
		ctx.SelectedProfile = &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		}
		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--output", "yaml",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "output must be one of: table, json"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})
}

func TestLokiOutputShortAliasWithCobra(t *testing.T) {
	t.Parallel()

	t.Run("query", func(t *testing.T) {
		t.Parallel()

		rawBody := `{"status":"success","data":{"resultType":"streams","result":[]}}`
		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(rawBody)}}}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		if err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			"-o", "json",
			`{cluster_env="staging", namespace="tidbcluster-42"}`,
		}, &stdout); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		assertJSONEqual(t, stdout.String(), rawBody)
	})

	t.Run("instance-map", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)}}}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		var stdout strings.Builder
		if err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			"-o", "json",
		}, &stdout); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		assertJSONEqual(t, stdout.String(), `[]`)
	})
}

func TestLokiInstanceMapJSONOutputMapsAndSortsRows(t *testing.T) {
	t.Parallel()

	body := `{"status":"success","data":{"resultType":"streams","result":[{"stream":{"node":"cn-hk.10.0.8.9","instance":"db-c"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"db-b"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"db-a"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"db-a"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"db-c"},"values":[["1712123000000000000","line"]]}]}}`
	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(body)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	var stdout strings.Builder
	err := lokiCommandHandler(ctx, []string{
		"instance-map",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600",
		"--output", "json",
	}, &stdout)
	if err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}
	assertJSONEqual(t, stdout.String(), `[
		{"node":"cn-hk.10.0.7.198","data_server_instance":"cn-hk.10.0.7.198","loki_instance":"db-a","ip":"10.0.7.198"},
		{"node":"cn-hk.10.0.7.198","data_server_instance":"cn-hk.10.0.7.198","loki_instance":"db-b","ip":"10.0.7.198"},
		{"node":"cn-hk.10.0.7.198","data_server_instance":"cn-hk.10.0.7.198","loki_instance":"db-c","ip":"10.0.7.198"},
		{"node":"cn-hk.10.0.8.9","data_server_instance":"cn-hk.10.0.8.9","loki_instance":"db-c","ip":"10.0.8.9"}
	]`)
}

func TestLokiInstanceMapTableOutputIsSortedAndStable(t *testing.T) {
	t.Parallel()

	body := `{"status":"success","data":{"resultType":"streams","result":[{"stream":{"node":"cn-hk.10.0.8.9","instance":"pod-z"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"pod-gamma"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"pod-alpha"},"values":[["1712123000000000000","line"]]},{"stream":{"node":"cn-hk.10.0.7.198","instance":"pod-alpha"},"values":[["1712123000000000000","line"]]}]}}`
	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(body)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	var stdout strings.Builder
	if err := lokiCommandHandler(ctx, []string{
		"instance-map",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600",
	}, &stdout); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	rows := strings.Split(strings.TrimSpace(stdout.String()), "\n")
	if got, want := len(rows), 4; got != want {
		t.Fatalf("row count = %d, want %d", got, want)
	}
	header := strings.Fields(rows[0])
	wantHeader := []string{"NODE", "DATA_SERVER_INSTANCE", "LOKI_INSTANCE", "IP"}
	if got, want := strings.Join(header, "|"), strings.Join(wantHeader, "|"); got != want {
		t.Fatalf("table header = %q, want %q", got, want)
	}

	cells := make([][]string, 0, len(rows)-1)
	for _, row := range rows[1:] {
		fields := strings.Fields(row)
		if got := len(fields); got != 4 {
			t.Fatalf("row field count = %d, want %d", got, 4)
		}
		cells = append(cells, fields)
	}

	gotDataRows := []string{
		strings.Join(cells[0], "|"),
		strings.Join(cells[1], "|"),
		strings.Join(cells[2], "|"),
	}
	wantDataRows := []string{
		"cn-hk.10.0.7.198|cn-hk.10.0.7.198|pod-alpha|10.0.7.198",
		"cn-hk.10.0.7.198|cn-hk.10.0.7.198|pod-gamma|10.0.7.198",
		"cn-hk.10.0.8.9|cn-hk.10.0.8.9|pod-z|10.0.8.9",
	}
	for i := range wantDataRows {
		if got, want := gotDataRows[i], wantDataRows[i]; got != want {
			t.Fatalf("row %d = %q, want %q", i, got, want)
		}
	}
}

func TestLokiInstanceMapRejectsUnexpectedPositionalArgs(t *testing.T) {
	t.Parallel()

	ctx := &RuntimeContext{
		Client:      &recordingClinicClient{},
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}
	err := lokiCommandHandler(ctx, []string{
		"instance-map",
		"--cluster-id", "cluster-42",
		"--start", "1712123000000000000",
		"--end", "1712123600000000000",
		"unexpected",
	}, &strings.Builder{})
	if err == nil {
		t.Fatal("lokiCommandHandler() error = nil, want validation error")
	}
	if got, want := err.Error(), "unexpected arguments: [unexpected]"; got != want {
		t.Fatalf("error = %q, want %q", got, want)
	}
}

func TestLokiInstanceMapRejectsInvalidLimit(t *testing.T) {
	t.Parallel()

	t.Run("rejects non-positive limit", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--limit", "0",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "limit must be greater than 0"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("rejects limit above max", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--limit", "1001",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "limit must be less than or equal to 1000"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("supports custom valid limit", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
			},
		}
		ctx := &RuntimeContext{
			Client:      client,
			SelectedEnv: "staging",
			SelectedProfile: &config.ProfileRecord{
				ClusterID:           "cluster-42",
				OrgID:               "org-1",
				ClusterDeployTypeV2: "dedicated",
			},
		}

		if err := lokiCommandHandler(ctx, []string{
			"instance-map",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			"--limit", "500",
		}, &strings.Builder{}); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}
		assertQueryValue(t, client.request.Query, "limit", "500")
	})
}

func TestLokiQueryRangeSupportsNowExpressions(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := legacyLokiTestContext(client)

	before := time.Now().Unix()
	err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "now-1d",
		"--end", "now",
		"--output", "json",
		`{app="tidb", cluster_env="staging", namespace="tidbcluster-42"}`,
	}, &strings.Builder{})
	if err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}
	after := time.Now().Unix()

	startRaw := client.request.Query.Get("start")
	endRaw := client.request.Query.Get("end")
	startUnix, parseErr := strconv.ParseInt(startRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("start query value = %q, want unix nanoseconds: %v", startRaw, parseErr)
	}
	endUnix, parseErr := strconv.ParseInt(endRaw, 10, 64)
	if parseErr != nil {
		t.Fatalf("end query value = %q, want unix nanoseconds: %v", endRaw, parseErr)
	}
	wantMin := int64((24 * time.Hour).Nanoseconds())
	wantMax := wantMin + int64(time.Second)
	if got := endUnix - startUnix; got < wantMin || got > wantMax {
		t.Fatalf("range length = %d, want between %d and %d", got, wantMin, wantMax)
	}
	endUnixSeconds := endUnix / int64(time.Second)
	if endUnixSeconds < before || endUnixSeconds > after {
		t.Fatalf("end query value = %d, want unix seconds between %d and %d", endUnix, before, after)
	}
}

func TestLokiQueryRangeBuildsDefaultQueryFromClusterMetadata(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name              string
		clusterID         string
		clusterSearchItem string
		wantNamespace     string
	}{
		{
			name:              "premium uses parent namespace",
			clusterID:         "tenant-42",
			clusterSearchItem: `{"clusterID":"tenant-42","orgID":"org-1","clusterDeployTypeV2":"premium","parentID":"pool-9"}`,
			wantNamespace:     "tidbpool-9",
		},
		{
			name:              "dedicated uses cluster namespace",
			clusterID:         "cluster-42",
			clusterSearchItem: `{"clusterID":"cluster-42","orgID":"org-1","clusterDeployTypeV2":"dedicated"}`,
			wantNamespace:     "tidbcluster-42",
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			client := &recordingClinicClient{
				responses: []clinicResponse{
					{body: []byte(`{"items":[` + tc.clusterSearchItem + `]}`)},
					{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
				},
			}
			ctx := &RuntimeContext{Client: client, SelectedEnv: "staging"}

			var stdout strings.Builder
			if err := lokiCommandHandler(ctx, []string{
				"query",
				"--cluster-id", tc.clusterID,
				"--start", "1712123000",
				"--end", "1712123600000",
				"--output", "json",
			}, &stdout); err != nil {
				t.Fatalf("lokiCommandHandler() error = %v", err)
			}

			if got, want := len(client.requests), 2; got != want {
				t.Fatalf("request count = %d, want %d", got, want)
			}
			if got, want := client.requests[0].Path, dashboardClustersPath; got != want {
				t.Fatalf("cluster search path = %q, want %q", got, want)
			}
			assertQueryValue(t, client.requests[0].Query, "cluster_id", tc.clusterID)
			assertQueryValue(t, client.requests[0].Query, "show_deleted", "true")
			assertQueryValue(t, client.requests[1].Query, "query", fmt.Sprintf(`{cluster_env="staging", namespace="%s"}`, tc.wantNamespace))
			assertQueryValue(t, client.requests[1].Query, "start", "1712123000000000000")
			assertQueryValue(t, client.requests[1].Query, "end", "1712123600000000000")
		})
	}
}

func TestLokiQueryRangeUsesSelectedProfileMetadataWithoutRemoteLookup(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	var stdout strings.Builder
	if err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600000",
		"--output", "json",
	}, &stdout); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := len(client.requests), 1; got != want {
		t.Fatalf("request count = %d, want %d", got, want)
	}
	assertQueryValue(t, client.request.Query, "query", `{cluster_env="staging", namespace="tidbcluster-42"}`)
	assertQueryValue(t, client.request.Query, "start", "1712123000000000000")
	assertQueryValue(t, client.request.Query, "end", "1712123600000000000")
}

func TestLokiQueryRangeFallsBackToRemoteWhenSelectedProfileMetadataIncomplete(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"items":[{"clusterID":"tenant-42","orgID":"org-1","clusterDeployTypeV2":"premium","parentID":"pool-9"}]}`)},
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "tenant-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "premium",
		},
	}

	var stdout strings.Builder
	if err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "tenant-42",
		"--start", "1712123000",
		"--end", "1712123600000",
		"--output", "json",
	}, &stdout); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := len(client.requests), 2; got != want {
		t.Fatalf("request count = %d, want %d", got, want)
	}
	assertQueryValue(t, client.requests[0].Query, "cluster_id", "tenant-42")
	assertQueryValue(t, client.requests[1].Query, "query", `{cluster_env="staging", namespace="tidbpool-9"}`)
	assertQueryValue(t, client.requests[1].Query, "start", "1712123000000000000")
	assertQueryValue(t, client.requests[1].Query, "end", "1712123600000000000")
}

func TestLokiQueryRangeDefaultQueryRequiresParentIDForPremium(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"items":[{"clusterID":"tenant-42","orgID":"org-1","clusterDeployTypeV2":"premium"}]}`)},
		},
	}
	ctx := &RuntimeContext{Client: client, SelectedEnv: "staging"}

	err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "tenant-42",
		"--start", "1712123000",
		"--end", "1712123600000",
	}, &strings.Builder{})
	if err == nil {
		t.Fatal("lokiCommandHandler() error = nil, want validation error")
	}
	if got, want := err.Error(), "premium cluster has no parentID: tenant-42"; got != want {
		t.Fatalf("error = %q, want %q", got, want)
	}
}

func TestLokiQueryRangeInjectsDefaultSelectorLabelsWhenMissing(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	if err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600000",
		"--output", "json",
		`{component="worker-tidb"}`,
	}, &strings.Builder{}); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := len(client.requests), 1; got != want {
		t.Fatalf("request count = %d, want %d", got, want)
	}
	assertQueryValue(t, client.request.Query, "query", `{component="worker-tidb", cluster_env="staging", namespace="tidbcluster-42"}`)
}

func TestLokiQueryRangeDoesNotOverrideProvidedSelectorLabels(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	ctx := &RuntimeContext{
		Client:      client,
		SelectedEnv: "staging",
		SelectedProfile: &config.ProfileRecord{
			ClusterID:           "cluster-42",
			OrgID:               "org-1",
			ClusterDeployTypeV2: "dedicated",
		},
	}

	query := `{component="worker-tidb", cluster_env="prod", namespace="custom-ns"}`
	if err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600000",
		"--output", "json",
		query,
	}, &strings.Builder{}); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := len(client.requests), 1; got != want {
		t.Fatalf("request count = %d, want %d", got, want)
	}
	assertQueryValue(t, client.request.Query, "query", query)
}

func TestLokiQueryWarnsForExplicitScopeLabels(t *testing.T) {
	t.Parallel()

	testCases := []struct {
		name  string
		query string
		want  []string
	}{
		{
			name:  "explicit cluster env and namespace",
			query: `{component="worker-tidb", cluster_env="staging", namespace="tidbcluster-42"} |= "error"`,
			want: []string{
				`[clinic warning]`,
				`explicit Loki scope label(s): cluster_env, namespace`,
				`clinic automatically scopes Loki queries`,
				`--allow-explicit-scope`,
			},
		},
		{
			name:  "namespace wildcard",
			query: `{cluster_env="staging", namespace=~".*"} |= "error"`,
			want: []string{
				`[clinic warning]`,
				`namespace=~".*"`,
				`cross-namespace`,
				`--allow-explicit-scope`,
			},
		},
	}

	for _, tc := range testCases {
		tc := tc
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			client := &recordingClinicClient{
				responses: []clinicResponse{
					{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
				},
			}
			var stderr strings.Builder
			ctx := legacyLokiTestContext(client)
			ctx.DiagnosticsOut = &stderr

			if err := lokiCommandHandler(ctx, []string{
				"query",
				"--cluster-id", "cluster-42",
				"--start", "1712123000",
				"--end", "1712123600000",
				"--output", "json",
				tc.query,
			}, &strings.Builder{}); err != nil {
				t.Fatalf("lokiCommandHandler() error = %v", err)
			}

			assertContainsAll(t, stderr.String(), tc.want)
			assertQueryValue(t, client.request.Query, "query", tc.query)
		})
	}
}

func TestLokiQueryAllowExplicitScopeSuppressesScopeWarning(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{
		responses: []clinicResponse{
			{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)},
		},
	}
	var stderr strings.Builder
	ctx := legacyLokiTestContext(client)
	ctx.DiagnosticsOut = &stderr
	query := `{cluster_env="staging", namespace=~".*"} |= "error"`

	if err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "1712123000",
		"--end", "1712123600000",
		"--output", "json",
		"--allow-explicit-scope",
		query,
	}, &strings.Builder{}); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got := stderr.String(); got != "" {
		t.Fatalf("stderr = %q, want no explicit-scope warning", got)
	}
	assertQueryValue(t, client.request.Query, "query", query)
}

func TestLokiQueryRangeUsesOptionalParamsAndRejectsInvalidDirection(t *testing.T) {
	t.Parallel()
	explicitQuery := `{app="tidb", cluster_env="staging", namespace="tidbcluster-42"}`

	t.Run("forwards optional params and prints default log lines", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[{"stream":{"component":"tidb","instance":"worker-tidb-0"},"values":[["1712123600000000000","log line 1"],["1712123599000000000","log line 2"]]}]}}`)}}}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		if err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600000",
			"--limit", "500",
			"--direction", "backward",
			"--step", "30s",
			explicitQuery,
		}, &stdout); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		if got, want := client.request.Method, http.MethodGet; got != want {
			t.Fatalf("method = %q, want %q", got, want)
		}
		if got, want := len(client.requests), 1; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
		if got, want := client.request.Path, "/data-proxy/loki/cluster-42/api/v1/query_range"; got != want {
			t.Fatalf("path = %q, want %q", got, want)
		}
		assertQueryValue(t, client.request.Query, "query", explicitQuery)
		assertQueryValue(t, client.request.Query, "start", "1712123000000000000")
		assertQueryValue(t, client.request.Query, "end", "1712123600000000000")
		assertQueryValue(t, client.request.Query, "limit", "500")
		assertQueryValue(t, client.request.Query, "direction", "backward")
		assertQueryValue(t, client.request.Query, "step", "30s")
		if got, want := stdout.String(), "[instance=worker-tidb-0] log line 1\n[instance=worker-tidb-0] log line 2\n"; got != want {
			t.Fatalf("stdout = %q, want %q", got, want)
		}
	})

	t.Run("uses default limit when not provided", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(`{"status":"success","data":{"resultType":"streams","result":[]}}`)}}}
		ctx := legacyLokiTestContext(client)

		if err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600000",
			explicitQuery,
		}, &strings.Builder{}); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		assertQueryValue(t, client.request.Query, "limit", "10")
	})

	t.Run("supports json output", func(t *testing.T) {
		t.Parallel()

		rawBody := `{"status":"success","data":{"resultType":"streams","result":[],"stats":{"summary":{"totalEntriesReturned":0}}}}`
		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(rawBody)}}}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		if err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"-o", "json",
			explicitQuery,
		}, &stdout); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		assertJSONEqual(t, stdout.String(), rawBody)
	})

	t.Run("supports --output json alias", func(t *testing.T) {
		t.Parallel()

		rawBody := `{"status":"success","data":{"resultType":"streams","result":[]}}`
		client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(rawBody)}}}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		if err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--output", "json",
			explicitQuery,
		}, &stdout); err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}

		assertJSONEqual(t, stdout.String(), rawBody)
	})

	t.Run("rejects invalid direction", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--direction", "sideways",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "direction must be one of: forward, backward"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("rejects invalid output format", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"-o", "yaml",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "output must be one of: json"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("rejects non-positive limit", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--limit", "0",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "limit must be greater than 0"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("rejects limit above server max", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
			"--limit", "1001",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "limit must be less than or equal to 1000"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})

	t.Run("rejects removed --query flag", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--query", `{app="tidb"}`,
			"--start", "1712123000000000000",
			"--end", "1712123600000000000",
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got := err.Error(); !strings.Contains(got, "unknown flag") {
			t.Fatalf("error = %q, want unknown flag error", got)
		}
	})

	t.Run("rejects multiple positional query arguments", func(t *testing.T) {
		t.Parallel()

		ctx := &RuntimeContext{Client: &recordingClinicClient{}}
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			`{app="tidb"}`,
			`{component="tikv"}`,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want validation error")
		}
		if got, want := err.Error(), "unexpected arguments: [{app=\"tidb\"} {component=\"tikv\"}]"; got != want {
			t.Fatalf("error = %q, want %q", got, want)
		}
	})
}

func TestLokiQueryRangeRetriesGatewayTimeout(t *testing.T) {
	t.Parallel()

	const explicitQuery = `{app="tidb", cluster_env="staging", namespace="tidbcluster-42"}`

	t.Run("retries up to success for default output", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					body: []byte(`{"status":"success","data":{"resultType":"streams","result":[{"stream":{"component":"tidb","instance":"worker-tidb-0"},"values":[["1712123600000000000","log line 1"]]}]}}`),
				},
			},
		}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			explicitQuery,
		}, &stdout)
		if err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}
		if got, want := len(client.requests), 3; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
		if got, want := stdout.String(), "[instance=worker-tidb-0] log line 1\n"; got != want {
			t.Fatalf("stdout = %q, want %q", got, want)
		}
	})

	t.Run("retries up to max then returns last 504 error", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
			},
		}
		ctx := legacyLokiTestContext(client)

		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want gateway timeout error")
		}

		var httpErr *clinic.HTTPError
		if ok := errors.As(err, &httpErr); !ok {
			t.Fatalf("error type = %T, want *clinic.HTTPError", err)
		}
		if got, want := httpErr.StatusCode, http.StatusGatewayTimeout; got != want {
			t.Fatalf("status code = %d, want %d", got, want)
		}
		if got, want := len(client.requests), 4; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
	})

	t.Run("does not retry for non-504 errors", func(t *testing.T) {
		t.Parallel()

		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusBadGateway,
						Status:     "502 Bad Gateway",
					},
				},
			},
		}
		ctx := legacyLokiTestContext(client)

		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			explicitQuery,
		}, &strings.Builder{})
		if err == nil {
			t.Fatal("lokiCommandHandler() error = nil, want non-504 error")
		}
		if got, want := len(client.requests), 1; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
	})

	t.Run("retries for json output mode as well", func(t *testing.T) {
		t.Parallel()

		rawBody := `{"status":"success","data":{"resultType":"streams","result":[]}}`
		client := &recordingClinicClient{
			responses: []clinicResponse{
				{
					err: &clinic.HTTPError{
						StatusCode: http.StatusGatewayTimeout,
						Status:     "504 Gateway Timeout",
					},
				},
				{
					body: []byte(rawBody),
				},
			},
		}
		ctx := legacyLokiTestContext(client)

		var stdout strings.Builder
		err := lokiCommandHandler(ctx, []string{
			"query",
			"--cluster-id", "cluster-42",
			"--start", "1712123000",
			"--end", "1712123600",
			"--output", "json",
			explicitQuery,
		}, &stdout)
		if err != nil {
			t.Fatalf("lokiCommandHandler() error = %v", err)
		}
		if got, want := len(client.requests), 2; got != want {
			t.Fatalf("request count = %d, want %d", got, want)
		}
		assertJSONEqual(t, stdout.String(), rawBody)
	})
}

func TestLokiLabelsUsesOptionalStartAndEnd(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(`{"status":"success","data":["cluster","component"]}`)}}}
	ctx := &RuntimeContext{Client: client}

	var stdout strings.Builder
	expectedStart := strconv.FormatInt(time.Date(2026, time.April, 3, 10, 4, 43, 0, time.Local).UnixNano(), 10)
	expectedEnd := strconv.FormatInt(time.Date(2026, time.April, 3, 10, 5, 43, 0, time.Local).UnixNano(), 10)
	if err := lokiCommandHandler(ctx, []string{
		"labels",
		"--cluster-id", "cluster-42",
		"--start", "2026-04-03 10:04:43",
		"--end", "2026-04-03 10:05:43",
	}, &stdout); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := client.request.Method, http.MethodGet; got != want {
		t.Fatalf("method = %q, want %q", got, want)
	}
	if got, want := client.request.Path, "/data-proxy/loki/cluster-42/api/v1/labels"; got != want {
		t.Fatalf("path = %q, want %q", got, want)
	}
	assertQueryValue(t, client.request.Query, "start", expectedStart)
	assertQueryValue(t, client.request.Query, "end", expectedEnd)
	assertJSONEqual(t, stdout.String(), `{"status":"success","data":["cluster","component"]}`)
}

func TestLokiLabelValuesUsesLabelPathAndOptionalRange(t *testing.T) {
	t.Parallel()

	client := &recordingClinicClient{responses: []clinicResponse{{body: []byte(`{"status":"success","data":["tidb","tikv"]}`)}}}
	ctx := &RuntimeContext{Client: client}

	var stdout strings.Builder
	if err := lokiCommandHandler(ctx, []string{
		"label-values",
		"--cluster-id", "cluster-42",
		"--label", "component",
		"--start", "2026-04-03 02:04:43 +0000",
		"--end", "1775181883000000000",
	}, &stdout); err != nil {
		t.Fatalf("lokiCommandHandler() error = %v", err)
	}

	if got, want := client.request.Method, http.MethodGet; got != want {
		t.Fatalf("method = %q, want %q", got, want)
	}
	if got, want := client.request.Path, "/data-proxy/loki/cluster-42/api/v1/label/component/values"; got != want {
		t.Fatalf("path = %q, want %q", got, want)
	}
	assertQueryValue(t, client.request.Query, "start", "1775181883000000000")
	assertQueryValue(t, client.request.Query, "end", "1775181883000000000")
	assertJSONEqual(t, stdout.String(), `{"status":"success","data":["tidb","tikv"]}`)
}

func TestLokiQueryRangeRejectsInvalidTimeInput(t *testing.T) {
	t.Parallel()

	ctx := &RuntimeContext{Client: &recordingClinicClient{}}
	err := lokiCommandHandler(ctx, []string{
		"query",
		"--cluster-id", "cluster-42",
		"--start", "not-a-time",
		"--end", "1712123600",
		`{app="tidb"}`,
	}, &strings.Builder{})
	if err == nil {
		t.Fatal("lokiCommandHandler() error = nil, want validation error")
	}
	if got := err.Error(); !strings.Contains(got, "invalid --start") {
		t.Fatalf("error = %q, want invalid --start", got)
	}
}

func TestLokiLabelValuesRequiresLabel(t *testing.T) {
	t.Parallel()

	ctx := &RuntimeContext{Client: &recordingClinicClient{}}
	err := lokiCommandHandler(ctx, []string{
		"label-values",
		"--cluster-id", "cluster-42",
	}, &strings.Builder{})
	if err == nil {
		t.Fatal("lokiCommandHandler() error = nil, want validation error")
	}
	if got, want := err.Error(), "missing required --label"; got != want {
		t.Fatalf("error = %q, want %q", got, want)
	}
}