package agent

import (
	"context"
	stderrors "errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	service_config "github.com/tidbcloud/dataflow-service/internal/config"
	"github.com/tidbcloud/dataflow-service/internal/crypto"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/encoding/protojson"
)

func setTestSharedHTTPClient(t *testing.T, client *http.Client) {
	t.Helper()
	prevClient := sharedHttpClient
	sharedHttpClient = client
	t.Cleanup(func() {
		sharedHttpClient = prevClient
	})

	prevToken := service_config.Global().GrpcConfig.MgmtSvcProxyToken
	service_config.Global().GrpcConfig.MgmtSvcProxyToken = "test-token"
	t.Cleanup(func() {
		service_config.Global().GrpcConfig.MgmtSvcProxyToken = prevToken
	})
}

func TestNGRegionalClientPrecheckDMJobRetriesTransient502(t *testing.T) {
	var attempts atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/v1/dm/jobprecheck", r.URL.Path)

		if attempts.Add(1) == 1 {
			http.Error(w, `{"message":"An invalid response was received from the upstream server"}`, http.StatusBadGateway)
			return
		}

		payload, err := protojson.Marshal(&agentpb.PrecheckDMJobResp{CheckResults: []byte(`{"results":[{"state":"success"}]}`)})
		require.NoError(t, err)
		w.Header().Set("Content-Type", "application/json")
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.PrecheckDMJob(context.Background(), &agentpb.PrecheckDMJobReq{ReqId: "pck-1", ClusterId: "cluster-1"})
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.Equal(t, int32(2), attempts.Load())
}

func TestNGRegionalClientPrecheckDMJobStopsAfterTransientRetryLimit(t *testing.T) {
	var attempts atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		attempts.Add(1)
		http.Error(w, `{"message":"An invalid response was received from the upstream server"}`, http.StatusBadGateway)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	_, err := client.PrecheckDMJob(context.Background(), &agentpb.PrecheckDMJobReq{ReqId: "pck-1", ClusterId: "cluster-1"})
	require.Error(t, err)
	require.Equal(t, int32(3), attempts.Load())
}

func TestNGRegionalClientPrecheckDMJobDoesNotRetryBadRequest(t *testing.T) {
	var attempts atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		attempts.Add(1)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		_, err := w.Write([]byte(`{"code":3,"message":"bad request"}`))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	_, err := client.PrecheckDMJob(context.Background(), &agentpb.PrecheckDMJobReq{ReqId: "pck-1", ClusterId: "cluster-1"})
	require.Error(t, err)
	require.Equal(t, "bad request", err.Error())
	require.Equal(t, int32(1), attempts.Load())
}

func TestIsRetryableUpstreamErrorRecognizesTransportFailure(t *testing.T) {
	err := &urlErrorStub{}
	require.True(t, isRetryableUpstreamError(err))
}

type urlErrorStub struct{}

func (u *urlErrorStub) Error() string   { return "EOF" }
func (u *urlErrorStub) Timeout() bool   { return false }
func (u *urlErrorStub) Temporary() bool { return false }

func TestNGRegionalClientPrecheckDMJobHonorsContextCancellation(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(100 * time.Millisecond)
		http.Error(w, `{"message":"An invalid response was received from the upstream server"}`, http.StatusBadGateway)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
	defer cancel()

	client := &NGRegionalClient{AgentUrl: server.URL}
	_, err := client.PrecheckDMJob(ctx, &agentpb.PrecheckDMJobReq{ReqId: "pck-1", ClusterId: "cluster-1"})
	require.Error(t, err)
}

func TestGetResponseRetriesTransient502ForSharedPath(t *testing.T) {
	var attempts atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		if attempts.Add(1) == 1 {
			http.Error(w, `{"message":"An invalid response was received from the upstream server"}`, http.StatusBadGateway)
			return
		}
		_, err := w.Write([]byte(`{"ok":true}`))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	resp, err := getResponseWithRetry(context.Background(), http.MethodGet, server.URL, nil)
	require.NoError(t, err)
	require.Equal(t, []byte(`{"ok":true}`), resp)
	require.Equal(t, int32(2), attempts.Load())
}

func TestGetResponseDoesNotRetryTransient502ByDefault(t *testing.T) {
	var attempts atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		attempts.Add(1)
		http.Error(w, `{"message":"An invalid response was received from the upstream server"}`, http.StatusBadGateway)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	_, err := getResponse(context.Background(), http.MethodPost, server.URL, nil)
	require.Error(t, err)
	require.Equal(t, int32(1), attempts.Load())
}

func TestGetResponseDoesNotRetryTransportFailureByDefault(t *testing.T) {
	var attempts atomic.Int32
	setTestSharedHTTPClient(t, &http.Client{
		Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) {
			attempts.Add(1)
			return nil, &urlErrorStub{}
		}),
	})

	_, err := getResponse(context.Background(), http.MethodPost, "http://example.com", nil)
	require.Error(t, err)
	var retryErr *retryableUpstreamError
	require.True(t, stderrors.As(err, &retryErr))
	require.Equal(t, int32(1), attempts.Load())
}

func TestGetResponsePropagatesTraceparent(t *testing.T) {
	var traceparent string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		traceparent = r.Header.Get("Traceparent")
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	prevClient := sharedHttpClient
	sharedHttpClient = &http.Client{
		Transport: otelhttp.NewTransport(server.Client().Transport),
	}
	t.Cleanup(func() {
		sharedHttpClient = prevClient
	})

	prevToken := service_config.Global().GrpcConfig.MgmtSvcProxyToken
	service_config.Global().GrpcConfig.MgmtSvcProxyToken = "test-token"
	t.Cleanup(func() {
		service_config.Global().GrpcConfig.MgmtSvcProxyToken = prevToken
	})

	prevProvider := otel.GetTracerProvider()
	prevPropagator := otel.GetTextMapPropagator()
	provider := sdktrace.NewTracerProvider()
	otel.SetTracerProvider(provider)
	otel.SetTextMapPropagator(propagation.TraceContext{})
	t.Cleanup(func() {
		otel.SetTracerProvider(prevProvider)
		otel.SetTextMapPropagator(prevPropagator)
		_ = provider.Shutdown(context.Background())
	})

	ctx, span := provider.Tracer("regional-ng-test").Start(context.Background(), "propagate-trace")
	defer span.End()

	_, err := getResponse(ctx, http.MethodPost, server.URL, nil)
	require.NoError(t, err)
	require.NotEmpty(t, traceparent)
}

func TestCreateSchemasErrorFromResp(t *testing.T) {
	t.Run("returns pem changed error", func(t *testing.T) {
		err := createSchemasErrorFromResp(&agentpb.CreateSchemasResp{PemChanged: true})
		require.ErrorIs(t, err, PEMChangedError)
	})

	t.Run("returns error response from top level error msg", func(t *testing.T) {
		msg := "failed to create schemas: err=boom"
		err := createSchemasErrorFromResp(&agentpb.CreateSchemasResp{ErrorMsg: &msg})
		require.EqualError(t, err, msg)
	})

	t.Run("returns nil for success response", func(t *testing.T) {
		require.NoError(t, createSchemasErrorFromResp(&agentpb.CreateSchemasResp{}))
		require.NoError(t, createSchemasErrorFromResp(nil))
	})
}

func TestNGRegionalClientCreateTiKVCapacityReservation(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/premium/v1/sharedPools/101/tikvCapacityReservations", r.URL.Path)
		require.Contains(t, readBody(t, r), `"reservationId":"import-1"`)
		w.Header().Set("Content-Type", "application/json")
		_, err := w.Write([]byte(`{"sharedPoolId":"101","bizType":"TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT","reservationId":"import-1","needScale":true,"status":"TIKV_CAPACITY_RESERVATION_STATUS_PROVISIONING","from":{"nodeCount":3,"totalCapacityBytes":"1000","totalAvailableBytes":"200","usagePercentage":80},"to":{"nodeCount":4,"totalCapacityBytes":"1500","totalAvailableBytes":"700","usagePercentage":53}}`))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{BaseUrl: server.URL}
	resp, err := client.CreateTiKVCapacityReservation(context.Background(), "101", "import-1", 2048)
	require.NoError(t, err)
	require.Equal(t, "import-1", resp.ReservationID)
	require.Equal(t, TiKVCapacityReservationStatusProvisioning, resp.Status)
}

func TestNGRegionalClientGetTiKVCapacityReservation(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		require.Equal(t, "/api/premium/v1/sharedPools/101/tikvCapacityReservations/TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT/import-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		_, err := w.Write([]byte(`{"sharedPoolId":"101","bizType":"TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT","reservationId":"import-1","needScale":false,"status":"TIKV_CAPACITY_RESERVATION_STATUS_READY","from":{"nodeCount":3,"totalCapacityBytes":"1000","totalAvailableBytes":"200","usagePercentage":80},"to":{"nodeCount":3,"totalCapacityBytes":"1000","totalAvailableBytes":"200","usagePercentage":80}}`))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{BaseUrl: server.URL}
	resp, err := client.GetTiKVCapacityReservation(context.Background(), "101", "import-1")
	require.NoError(t, err)
	require.Equal(t, TiKVCapacityReservationStatusReady, resp.Status)
	require.False(t, resp.NeedScale)
}

func TestNGRegionalClientDeleteTiKVCapacityReservation(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodDelete, r.Method)
		require.Equal(t, "/api/premium/v1/sharedPools/101/tikvCapacityReservations/TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT/import-1", r.URL.Path)
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{BaseUrl: server.URL}
	require.NoError(t, client.DeleteTiKVCapacityReservation(context.Background(), "101", "import-1"))
}

func TestNGRegionalClientCreateImportDataSizeEstimateTaskIncludesTableMeta(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/v1/import/estimate-data-size-tasks", r.URL.Path)
		body := readBody(t, r)
		require.Contains(t, body, `"taskId":"estimate-1"`)
		require.Contains(t, body, `"database":"db1"`)
		require.Contains(t, body, `"table":"t1"`)
		require.Contains(t, body, `"wildcardPath":"s3://bucket/db1/t1*.csv"`)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.CreateImportDataSizeEstimateTaskResp{TaskId: "estimate-1"})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.CreateImportDataSizeEstimateTask(
		context.Background(),
		"estimate-1",
		101,
		"s3://bucket/path",
		-1,
		nil,
		&agentpb.CreateImportDataSizeEstimateTaskReq_TableMeta{
			Database:     "db1",
			Table:        "t1",
			WildcardPath: "s3://bucket/db1/t1*.csv",
		},
	)
	require.NoError(t, err)
	require.Equal(t, "estimate-1", resp.GetTaskId())
}

func TestNGRegionalClientCreateImportFilesTask(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/v1/import/files-tasks", r.URL.Path)
		body := readBody(t, r)
		require.Contains(t, body, `"taskId":"files-1"`)
		require.Contains(t, body, `"database":"db1"`)
		require.Contains(t, body, `"table":"t1"`)
		require.Contains(t, body, `"wildcardPath":"s3://bucket/db1/t1*.csv"`)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.CreateImportFilesTaskResp{TaskId: "files-1"})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.CreateImportFilesTask(
		context.Background(),
		"files-1",
		101,
		"s3://bucket/path",
		-1,
		&agentpb.GetImportFilesReq_TableMeta{
			Database:     "db1",
			Table:        "t1",
			WildcardPath: "s3://bucket/db1/t1*.csv",
		},
	)
	require.NoError(t, err)
	require.Equal(t, "files-1", resp.GetTaskId())
}

func TestNGRegionalClientGetImportFilesTask(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		require.Equal(t, "/api/v1/import/files-tasks/files-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.GetImportFilesTaskResp{
			TaskId: "files-1",
			Status: agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED,
			Result: &agentpb.GetImportFilesResp{
				TableMeta: []*agentpb.GetImportFilesResp_TableMeta{
					{
						Database:  "db1",
						Table:     "t1",
						TotalSize: 1024,
					},
				},
			},
		})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.GetImportFilesTask(context.Background(), "files-1")
	require.NoError(t, err)
	require.Equal(t, agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_SUCCEEDED, resp.GetStatus())
	require.NotNil(t, resp.GetResult())
	require.Len(t, resp.GetResult().GetTableMeta(), 1)
	require.Equal(t, uint64(1024), resp.GetResult().GetTableMeta()[0].GetTotalSize())
}

func TestNGRegionalClientCreateSchemasTask(t *testing.T) {
	pubPEM := mustTestPublicKeyPEM(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/crypto/pem":
			require.Equal(t, http.MethodPost, r.Method)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.GetPublicKeyResponse{Pem: pubPEM})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		case "/api/v1/import/create-schemas-tasks":
			require.Equal(t, http.MethodPost, r.Method)
			body := readBody(t, r)
			require.Contains(t, body, `"taskId":"schema-1"`)
			require.Contains(t, body, `"clusterId":"101"`)
			require.NotContains(t, body, `"password":"secret"`)
			require.Contains(t, body, `"passwordEncrypted":true`)
			require.Contains(t, body, `"user":"root"`)
			require.Contains(t, body, `"database":"db1"`)
			require.Contains(t, body, `"table":"t1"`)
			require.Contains(t, body, `"database":"db2"`)
			require.Contains(t, body, `"table":"t2"`)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.CreateSchemasTaskResp{TaskId: "schema-1"})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.CreateSchemasTask(
		context.Background(),
		"schema-1",
		101,
		"s3://bucket/path",
		"root",
		"secret",
		&agentpb.CreateSchemasReq_Schema{
			Database: "db1",
			Table:    "t1",
		},
		&agentpb.CreateSchemasReq_Schema{
			Database: "db2",
			Table:    "t2",
		},
	)
	require.NoError(t, err)
	require.Equal(t, "schema-1", resp.GetTaskId())
}

func TestNGRegionalClientGetSchemasTask(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		require.Equal(t, "/api/v1/import/create-schemas-tasks/schema-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.GetSchemasTaskResp{
			TaskId: "schema-1",
			Status: agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED,
			Result: &agentpb.CreateSchemasResp{},
		})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.GetSchemasTask(context.Background(), "schema-1")
	require.NoError(t, err)
	require.Equal(t, agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_SUCCEEDED, resp.GetStatus())
	require.NotNil(t, resp.GetResult())
	require.False(t, resp.GetResult().GetPemChanged())
}

func TestNGRegionalClientCreateImportIntoJobTask(t *testing.T) {
	pubPEM := mustTestPublicKeyPEM(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/crypto/pem":
			require.Equal(t, http.MethodPost, r.Method)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.GetPublicKeyResponse{Pem: pubPEM})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		case "/api/v1/import/import-into-job-tasks":
			require.Equal(t, http.MethodPost, r.Method)
			body := readBody(t, r)
			require.Contains(t, body, `"taskId":"submit-1"`)
			require.Contains(t, body, `"clusterId":"101"`)
			require.NotContains(t, body, `"password":"secret"`)
			require.Contains(t, body, `"passwordEncrypted":true`)
			require.Contains(t, body, `"user":"root"`)
			require.Contains(t, body, `"sql":"IMPORT INTO`)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.CreateImportIntoJobTaskResp{TaskId: "submit-1"})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.CreateImportIntoJobTask(
		context.Background(),
		"submit-1",
		101,
		"root",
		"secret",
		"IMPORT INTO `db1`.`t1` FROM 's3://bucket/path/*.csv' WITH DETACHED",
		nil,
	)
	require.NoError(t, err)
	require.Equal(t, "submit-1", resp.GetTaskId())
}

func TestNGRegionalClientGetImportIntoJobTask(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		require.Equal(t, "/api/v1/import/import-into-job-tasks/submit-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.GetImportIntoJobTaskResp{
			TaskId: "submit-1",
			Status: agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED,
			Result: &agentpb.GetImportIntoJobTaskResp_Result{
				JobId: 12345,
			},
		})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.GetImportIntoJobTask(context.Background(), "submit-1")
	require.NoError(t, err)
	require.Equal(t, agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED, resp.GetStatus())
	require.NotNil(t, resp.GetResult())
	require.Equal(t, int64(12345), resp.GetResult().GetJobId())
}

func TestNGRegionalClientGetImportIntoJobTaskNotFoundPlainText(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodGet, r.Method)
		require.Equal(t, "/api/v1/import/import-into-job-tasks/submit-1", r.URL.Path)
		http.NotFound(w, r)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	_, err := client.GetImportIntoJobTask(context.Background(), "submit-1")
	require.Error(t, err)
	var errResp *ErrorResponse
	require.ErrorAs(t, err, &errResp)
	require.Equal(t, int64(codes.NotFound), errResp.Code)
	require.Equal(t, "404 page not found", errResp.Message)
}

func TestNGRegionalClientBatchCreateImportIntoJobTasks(t *testing.T) {
	pubPEM := mustTestPublicKeyPEM(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/crypto/pem":
			require.Equal(t, http.MethodPost, r.Method)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.GetPublicKeyResponse{Pem: pubPEM})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		case "/api/v1/import/import-into-job-tasks:batchCreate":
			require.Equal(t, http.MethodPost, r.Method)
			body := readBody(t, r)
			require.Contains(t, body, `"clusterId":"101"`)
			require.NotContains(t, body, `"password":"secret"`)
			require.Contains(t, body, `"passwordEncrypted":true`)
			require.Contains(t, body, `"user":"root"`)
			require.Contains(t, body, `"taskId":"submit-1"`)
			require.Contains(t, body, `"taskId":"submit-2"`)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.BatchCreateImportIntoJobTasksResp{
				Tasks: []*agentpb.CreateImportIntoJobTaskResp{{TaskId: "submit-1"}, {TaskId: "submit-2"}},
			})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.BatchCreateImportIntoJobTasks(
		context.Background(),
		101,
		"root",
		"secret",
		[]*agentpb.BatchCreateImportIntoJobTasksReq_Task{
			{TaskId: "submit-1", Sql: "IMPORT INTO `db1`.`t1` FROM 's3://bucket/path/*.csv' WITH DETACHED"},
			{TaskId: "submit-2", Sql: "IMPORT INTO `db1`.`t2` FROM 's3://bucket/path/*.csv' WITH DETACHED"},
		},
	)
	require.NoError(t, err)
	require.Len(t, resp.GetTasks(), 2)
	require.Equal(t, "submit-1", resp.GetTasks()[0].GetTaskId())
	require.Equal(t, "submit-2", resp.GetTasks()[1].GetTaskId())
}

func TestNGRegionalClientBatchGetImportIntoJobTasks(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/v1/import/import-into-job-tasks:batchGet", r.URL.Path)
		body := readBody(t, r)
		require.JSONEq(t, `{"taskIds":["submit-1","submit-2"]}`, body)
		w.Header().Set("Content-Type", "application/json")
		payload, err := protojson.Marshal(&agentpb.BatchGetImportIntoJobTasksResp{
			Tasks: []*agentpb.GetImportIntoJobTaskResp{
				{
					TaskId: "submit-1",
					Status: agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_SUCCEEDED,
					Result: &agentpb.GetImportIntoJobTaskResp_Result{JobId: 12345},
				},
			},
		})
		require.NoError(t, err)
		_, err = w.Write(payload)
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.BatchGetImportIntoJobTasks(context.Background(), []string{"submit-1", "submit-2"})
	require.NoError(t, err)
	require.Len(t, resp.GetTasks(), 1)
	require.Equal(t, "submit-1", resp.GetTasks()[0].GetTaskId())
	require.Equal(t, int64(12345), resp.GetTasks()[0].GetResult().GetJobId())
}

func TestNGRegionalClientBatchCancelImportIntoJobs(t *testing.T) {
	pubPEM := mustTestPublicKeyPEM(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/v1/crypto/pem":
			require.Equal(t, http.MethodPost, r.Method)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.GetPublicKeyResponse{Pem: pubPEM})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		case "/api/v1/import/import-into-jobs:batchCancel":
			require.Equal(t, http.MethodPost, r.Method)
			body := strings.ReplaceAll(readBody(t, r), " ", "")
			require.Contains(t, body, `"clusterId":"101"`)
			require.NotContains(t, body, `"password":"secret"`)
			require.Contains(t, body, `"passwordEncrypted":true`)
			require.Contains(t, body, `"jobIds":["123","456"]`)
			w.Header().Set("Content-Type", "application/json")
			payload, err := protojson.Marshal(&agentpb.BatchCancelImportIntoJobsResp{
				Results: []*agentpb.BatchCancelImportIntoJobsResp_Result{
					{JobId: 123},
					{JobId: 456, ErrorMsg: stringPtr("already done")},
				},
			})
			require.NoError(t, err)
			_, err = w.Write(payload)
			require.NoError(t, err)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.BatchCancelImportIntoJobs(context.Background(), 101, "root", "secret", []uint64{123, 456})
	require.NoError(t, err)
	require.Len(t, resp.GetResults(), 2)
	require.Equal(t, uint64(123), resp.GetResults()[0].GetJobId())
	require.Empty(t, resp.GetResults()[0].GetErrorMsg())
	require.Equal(t, uint64(456), resp.GetResults()[1].GetJobId())
	require.Equal(t, "already done", resp.GetResults()[1].GetErrorMsg())
}

func mustTestPublicKeyPEM(t *testing.T) []byte {
	t.Helper()
	enc, err := crypto.NewRSAEncryptor(2048)
	require.NoError(t, err)
	pem, err := enc.ExportPublicKeyAsPEM()
	require.NoError(t, err)
	return pem
}

func readBody(t *testing.T, r *http.Request) string {
	t.Helper()
	body, err := io.ReadAll(r.Body)
	require.NoError(t, err)
	return string(body)
}

func stringPtr(s string) *string {
	return &s
}

type roundTripperFunc func(*http.Request) (*http.Response, error)

func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
	return f(req)
}