package statements

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/jarcoal/httpmock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	apierr "github.com/pingcap-inc/clinic/pkg/ngm/error"
	statementapi "github.com/pingcap-inc/clinic/pkg/ngm/statements/api"
	"github.com/pingcap-inc/clinic/pkg/o11y"
)

type recordingStatementQuery struct {
	listParam   *statementapi.ListParam
	plansParam  *statementapi.PlansParam
	detailParam *statementapi.DetailParam
}

func (q *recordingStatementQuery) List(_ context.Context, param *statementapi.ListParam) ([]*statementapi.Model, error) {
	q.listParam = param
	return []*statementapi.Model{{Digest: "v1-list"}}, nil
}

func (q *recordingStatementQuery) Plans(_ context.Context, param *statementapi.PlansParam) ([]*statementapi.Model, error) {
	q.plansParam = param
	return []*statementapi.Model{{Digest: "v1-plans"}}, nil
}

func (q *recordingStatementQuery) Detail(_ context.Context, param *statementapi.DetailParam) (*statementapi.Model, error) {
	q.detailParam = param
	return &statementapi.Model{Digest: "v1-detail"}, nil
}

func newStatementsTestRouter(svc *Service) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(apierr.ErrorHandlerFn)
	InitStatementsApi(router.Group("/ngm/api/v1"), svc)
	return router
}

func newStatementsTestO11yClient(t *testing.T) *o11y.Client {
	t.Helper()
	cli, err := o11y.New(
		"www.gs.us-west-2.aws.observability-staging.pingcap.cloud",
		"www.ds.%s.%s.observability-staging.pingcap.cloud",
		"client-id",
		"client-secret",
		"dev",
	)
	require.NoError(t, err)
	return cli
}

func activateStatementsO11yResolverMock(t *testing.T) {
	t.Helper()
	httpmock.Activate()
	t.Cleanup(httpmock.DeactivateAndReset)

	httpmock.RegisterResponder(http.MethodPost, "https://tidb-soc2.us.auth0.com/oauth/token", func(_ *http.Request) (*http.Response, error) {
		return httpmock.NewStringResponse(http.StatusOK, `{"access_token":"token","expires_in":86400,"token_type":"Bearer"}`), nil
	})
	httpmock.RegisterResponder(http.MethodGet, "http://www.gs.us-west-2.aws.observability-staging.pingcap.cloud/api/v1/clusters/", func(r *http.Request) (*http.Response, error) {
		query := r.URL.Query()
		assert.Equal(t, "dev-devtier", query.Get("bizType"))
		assert.Equal(t, "aws", query.Get("vendor"))
		assert.Equal(t, "us-east-1", query.Get("region"))
		assert.Equal(t, "100", query.Get("tenantID"))
		assert.Equal(t, "300", query.Get("applicationID"))
		assert.Equal(t, "true", query.Get("showDeleted"))

		body := `{"code":0,"data":{"clusters":[{"applicationID":"300","bizType":"dev-devtier","id":"statement-o11y-id","projectID":"200","region":"us-east-1","tenantID":"100","vendor":"aws","version":"1","internalReadUri":"https://o11y.example.com/internal/metrics/statement-o11y-id/api/v1/","storageInfo":[{"bucketName":"statement-bucket","category":"statements","objectPath":"statement/path"}]}],"continueToken":"","remainingItemCount":0},"message":""}`
		return httpmock.NewStringResponse(http.StatusOK, body), nil
	})
}

func activateDedicatedStatementsStorageMock(t *testing.T) {
	t.Helper()
	httpmock.Activate()
	t.Cleanup(httpmock.DeactivateAndReset)

	httpmock.RegisterResponder(http.MethodPost, "https://tidb-soc2.us.auth0.com/oauth/token", func(_ *http.Request) (*http.Response, error) {
		return httpmock.NewStringResponse(http.StatusOK, `{"access_token":"token","expires_in":86400,"token_type":"Bearer"}`), nil
	})
	httpmock.RegisterResponder(http.MethodGet, "http://www.gs.us-west-2.aws.observability-staging.pingcap.cloud/api/v1/clusters/", func(r *http.Request) (*http.Response, error) {
		query := r.URL.Query()
		assert.Equal(t, "dev-tidbcloud", query.Get("bizType"))
		assert.Equal(t, "aws", query.Get("vendor"))
		assert.Equal(t, "us-east-1", query.Get("region"))
		assert.Equal(t, "100", query.Get("tenantID"))
		assert.Empty(t, query.Get("projectID"))
		assert.Equal(t, "300", query.Get("applicationID"))
		assert.Equal(t, "true", query.Get("showDeleted"))

		body := `{"code":0,"data":{"clusters":[{"applicationID":"300","bizType":"dev-tidbcloud","id":"dedicated-o11y-id","projectID":"historical-o11y-project","region":"us-east-1","tenantID":"100","vendor":"aws","version":"1","storageInfo":[{"bucketName":"dedicated-statement-bucket","category":"statements","objectPath":"canonical/statement/path"}]}],"continueToken":"","remainingItemCount":0},"message":""}`
		return httpmock.NewStringResponse(http.StatusOK, body), nil
	})
}

func registerStatementsProxyResponder(
	t *testing.T,
	upstreamPath string,
	status int,
	body string,
	assertRequest func(*http.Request),
) *int {
	t.Helper()
	called := 0
	pattern := regexp.MustCompile(`^https://o11y\.example\.com` + regexp.QuoteMeta(upstreamPath) + `(?:\?.*)?$`)
	httpmock.RegisterRegexpResponder(http.MethodGet, pattern, func(r *http.Request) (*http.Response, error) {
		called++
		assert.Equal(t, upstreamPath, r.URL.Path)
		assert.Equal(t, []string{"dev"}, statementHeaderValues(r.Header, "x-env"))
		assert.Equal(t, []string{"dev-tidbcloud"}, statementHeaderValues(r.Header, "x-biz-type"))
		assert.Equal(t, "aws", r.Header.Get("x-provider"))
		assert.Equal(t, "us-east-1", r.Header.Get("x-region"))
		assert.Equal(t, "100", r.Header.Get("x-org-id"))
		if projectID := r.Header.Get("x-project-id"); projectID != "" {
			assert.Equal(t, "200", projectID)
		}
		assert.Equal(t, "300", r.Header.Get("x-cluster-id"))
		assert.Equal(t, "shared", r.Header.Get("x-deploy-type"))
		assert.Equal(t, "client-marker", r.Header.Get("x-client-marker"))
		if assertRequest != nil {
			assertRequest(r)
		}

		response := httpmock.NewStringResponse(status, body)
		response.Header.Set("Content-Type", "application/json")
		response.Header.Set("X-Upstream-Marker", "preserved")
		return response, nil
	})
	return &called
}

func statementHeaderValues(header http.Header, name string) []string {
	for key, values := range header {
		if strings.EqualFold(key, name) {
			return values
		}
	}
	return nil
}

func newStatementsRequest(t *testing.T, target, deployType string) *http.Request {
	t.Helper()
	req := httptest.NewRequest(http.MethodGet, target, nil)
	req.Header.Set("x-provider", "aws")
	req.Header.Set("x-region", "us-east-1")
	req.Header.Set("x-org-id", "100")
	req.Header.Set("x-project-id", "200")
	req.Header.Set("x-cluster-id", "300")
	req.Header.Set("x-deploy-type", deployType)
	req.Header.Set("x-client-marker", "client-marker")
	return req
}

func TestIsStatementV2DataSource(t *testing.T) {
	tests := []struct {
		name       string
		dataSource string
		want       bool
	}{
		{name: "missing"},
		{name: "empty", dataSource: " "},
		{name: "v1", dataSource: "v1"},
		{name: "unknown", dataSource: "deltalake"},
		{name: "v2", dataSource: "v2", want: true},
		{name: "different case", dataSource: "V2"},
		{name: "whitespace", dataSource: "  v2\t"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, isStatementV2DataSource(tt.dataSource))
		})
	}
}

func TestStatementV2FixedProxyRoutes(t *testing.T) {
	tests := []struct {
		name         string
		route        string
		upstreamPath string
		status       int
	}{
		{name: "datasource", route: "/statements/datasource", upstreamPath: statementsDatasourceO11yPath, status: http.StatusOK},
		{name: "field values", route: "/statements/field_values", upstreamPath: statementsFieldValuesO11yPath, status: http.StatusOK},
		{name: "summary", route: "/statements/summary", upstreamPath: statementsSummaryO11yPath, status: http.StatusAccepted},
		{name: "topstats upstream error", route: "/statements/topstats", upstreamPath: statementsTopStatsO11yPath, status: http.StatusServiceUnavailable},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			activateStatementsO11yResolverMock(t)
			called := registerStatementsProxyResponder(t, tt.upstreamPath, tt.status, `{"proxied":true}`, func(r *http.Request) {
				assert.Equal(t, "v2", r.URL.Query().Get("data_source"))
				assert.Equal(t, []string{"first", "second"}, r.URL.Query()["custom"])
				assert.Equal(t, `[{"field":"schema_name"}]`, r.URL.Query().Get("filters"))
			})
			router := newStatementsTestRouter(NewStatementsService("dev", nil, newStatementsTestO11yClient(t)))
			response := httptest.NewRecorder()
			target := "/ngm/api/v1" + tt.route + `?data_source=v2&custom=first&custom=second&filters=%5B%7B%22field%22%3A%22schema_name%22%7D%5D`

			router.ServeHTTP(response, newStatementsRequest(t, target, TierShared))

			require.Equal(t, 1, *called)
			require.Equal(t, tt.status, response.Code)
			require.JSONEq(t, `{"proxied":true}`, response.Body.String())
			require.Equal(t, "application/json", response.Header().Get("Content-Type"))
			require.Equal(t, "preserved", response.Header().Get("X-Upstream-Marker"))
		})
	}
}

func TestStatementV2ExistingRoutesBypassV1QueryValidation(t *testing.T) {
	tests := []struct {
		name         string
		route        string
		upstreamPath string
	}{
		{name: "list", route: "/statements/list", upstreamPath: statementsListO11yPath},
		{name: "plans", route: "/statements/plans", upstreamPath: statementsPlansO11yPath},
		{name: "plan detail", route: "/statements/plan/detail", upstreamPath: statementsPlanDetailO11yPath},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			activateStatementsO11yResolverMock(t)
			called := registerStatementsProxyResponder(t, tt.upstreamPath, http.StatusAccepted, `{"source":"v2"}`, func(r *http.Request) {
				assert.Equal(t, "v2", r.URL.Query().Get("data_source"))
				assert.Equal(t, []string{"one", "two"}, r.URL.Query()["extra"])
				assert.Empty(t, r.URL.Query().Get("begin_time"), "the V2 proxy must run before V1 required-query validation")
			})
			router := newStatementsTestRouter(NewStatementsService("dev", nil, newStatementsTestO11yClient(t)))
			response := httptest.NewRecorder()
			target := "/ngm/api/v1" + tt.route + "?data_source=v2&extra=one&extra=two"

			router.ServeHTTP(response, newStatementsRequest(t, target, TierShared))

			require.Equal(t, 1, *called)
			require.Equal(t, http.StatusAccepted, response.Code)
			require.JSONEq(t, `{"source":"v2"}`, response.Body.String())
		})
	}
}

func TestStatementV2PlanDetailAddsBinaryPlanText(t *testing.T) {
	activateStatementsO11yResolverMock(t)
	called := registerStatementsProxyResponder(
		t,
		statementsPlanDetailO11yPath,
		http.StatusOK,
		`{"binary_plan":"{\"main\":{\"name\":\"TableFullScan\",\"estRows\":1,\"taskType\":\"root\"}}","digest":"sql-digest"}`,
		func(r *http.Request) {
			assert.Equal(t, "v2", r.URL.Query().Get("data_source"))
		},
	)
	router := newStatementsTestRouter(NewStatementsService("dev", nil, newStatementsTestO11yClient(t)))
	response := httptest.NewRecorder()
	req := newStatementsRequest(
		t,
		"/ngm/api/v1/statements/plan/detail?data_source=v2&digest=sql-digest",
		TierShared,
	)

	router.ServeHTTP(response, req)

	require.Equal(t, 1, *called)
	require.Equal(t, http.StatusOK, response.Code)
	var body map[string]interface{}
	require.NoError(t, json.Unmarshal(response.Body.Bytes(), &body))
	require.Equal(t, `{"main":{"name":"TableFullScan","estRows":1,"taskType":"root"}}`, body["binary_plan"])
	require.Contains(t, body, "binary_plan_text")
	require.Contains(t, body["binary_plan_text"], "TableFullScan")
}

func TestAddBinaryPlanTextModifierRecoversMalformedBinaryPlan(t *testing.T) {
	body := []byte(`{"binary_plan":"{}","digest":"sql-digest"}`)

	modified, err := addBinaryPlanTextModifier(body)

	require.NoError(t, err)
	require.Equal(t, body, modified)
}

func TestAddBinaryPlanTextModifierLeavesResponseWithoutBinaryPlanUnchanged(t *testing.T) {
	body := []byte(`{"digest":"sql-digest"}`)

	modified, err := addBinaryPlanTextModifier(body)

	require.NoError(t, err)
	require.Equal(t, body, modified)
}

func TestStatementV2ProxyRejectsInvalidHeaders(t *testing.T) {
	targets := []string{
		"/ngm/api/v1/statements/datasource",
		"/ngm/api/v1/statements/field_values",
		"/ngm/api/v1/statements/summary",
		"/ngm/api/v1/statements/topstats",
		"/ngm/api/v1/statements/list?data_source=v2",
		"/ngm/api/v1/statements/plans?data_source=v2",
		"/ngm/api/v1/statements/plan/detail?data_source=v2",
	}

	for _, target := range targets {
		t.Run(target, func(t *testing.T) {
			router := newStatementsTestRouter(&Service{})
			response := httptest.NewRecorder()
			req := httptest.NewRequest(http.MethodGet, target, nil)

			router.ServeHTTP(response, req)

			require.Equal(t, http.StatusBadRequest, response.Code)
		})
	}
}

func TestStatementV2ProxyAllowsMissingProjectHeader(t *testing.T) {
	activateStatementsO11yResolverMock(t)
	called := registerStatementsProxyResponder(t, statementsDatasourceO11yPath, http.StatusOK, `{"v2_available":true}`, func(r *http.Request) {
		assert.Empty(t, r.Header.Get("x-project-id"))
	})
	router := newStatementsTestRouter(NewStatementsService("dev", nil, newStatementsTestO11yClient(t)))
	response := httptest.NewRecorder()
	req := newStatementsRequest(t, "/ngm/api/v1/statements/datasource", TierShared)
	req.Header.Del("x-project-id")

	router.ServeHTTP(response, req)

	require.Equal(t, 1, *called)
	require.Equal(t, http.StatusOK, response.Code)
	require.JSONEq(t, `{"v2_available":true}`, response.Body.String())
}

func TestStatementV1RoutesStillUseLegacyQueryImplementation(t *testing.T) {
	tests := []struct {
		name       string
		target     string
		assertCall func(*testing.T, *recordingStatementQuery)
	}{
		{
			name:   "list",
			target: "/ngm/api/v1/statements/list?begin_time=1&end_time=2",
			assertCall: func(t *testing.T, query *recordingStatementQuery) {
				require.NotNil(t, query.listParam)
				require.Equal(t, int64(1), query.listParam.StartUnixSec)
				require.Equal(t, int64(2), query.listParam.EndUnixSec)
				require.False(t, query.listParam.SharedTier)
				require.Equal(t, "dedicated-statement-bucket", query.listParam.StorageBucket)
				require.Equal(t, "canonical/statement/path", query.listParam.StoragePrefix)
			},
		},
		{
			name:   "plans",
			target: "/ngm/api/v1/statements/plans?begin_time=1&end_time=2&digest=sql-digest",
			assertCall: func(t *testing.T, query *recordingStatementQuery) {
				require.NotNil(t, query.plansParam)
				require.Equal(t, "sql-digest", query.plansParam.Digest)
				require.False(t, query.plansParam.SharedTier)
				require.Equal(t, "dedicated-statement-bucket", query.plansParam.StorageBucket)
				require.Equal(t, "canonical/statement/path", query.plansParam.StoragePrefix)
			},
		},
		{
			name:   "plan detail",
			target: "/ngm/api/v1/statements/plan/detail?begin_time=1&end_time=2&digest=sql-digest&plans=plan-digest",
			assertCall: func(t *testing.T, query *recordingStatementQuery) {
				require.NotNil(t, query.detailParam)
				require.Equal(t, "sql-digest", query.detailParam.Digest)
				require.Equal(t, []string{"plan-digest"}, query.detailParam.Plans)
				require.False(t, query.detailParam.SharedTier)
				require.Equal(t, "dedicated-statement-bucket", query.detailParam.StorageBucket)
				require.Equal(t, "canonical/statement/path", query.detailParam.StoragePrefix)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			activateDedicatedStatementsStorageMock(t)
			query := &recordingStatementQuery{}
			router := newStatementsTestRouter(&Service{Query: query, Env: "dev", o11yCli: newStatementsTestO11yClient(t)})
			response := httptest.NewRecorder()

			router.ServeHTTP(response, newStatementsRequest(t, tt.target, TierDedicated))

			require.Equal(t, http.StatusOK, response.Code)
			tt.assertCall(t, query)
		})
	}
}

func TestSharedStatementV1ListKeepsStoragePathWithoutV2Selector(t *testing.T) {
	activateStatementsO11yResolverMock(t)
	query := &recordingStatementQuery{}
	router := newStatementsTestRouter(&Service{Query: query, Env: "dev", o11yCli: newStatementsTestO11yClient(t)})
	response := httptest.NewRecorder()
	req := newStatementsRequest(t, "/ngm/api/v1/statements/list?begin_time=1&end_time=2&data_source=v1", TierShared)

	router.ServeHTTP(response, req)

	require.Equal(t, http.StatusOK, response.Code)
	require.NotNil(t, query.listParam)
	require.True(t, query.listParam.SharedTier)
	require.Equal(t, "statement-bucket", query.listParam.StorageBucket)
	require.Equal(t, "statement/path", query.listParam.StoragePrefix)
	for request := range httpmock.GetCallCountInfo() {
		require.False(t, strings.Contains(request, "/internal/diagnosis/api/v1/statements/list"))
	}
}