package agent

import (
	"context"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/pkg/crypto"
	agentpb "github.com/tidbcloud/dataflow-service/dataflow-agent-ng/idl/pbgen/proto/ng_agent"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

func TestNGRegionalClientGetImportFiles(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/get-import-files", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		_, err := w.Write([]byte(`{"tableMeta":[{"database":"db1","table":"t1","wildcardPath":"s3://bucket/path/db1.t1.*.csv"}]}`))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.GetImportFiles(context.Background(), 123, "s3://bucket/path", 2000)
	require.NoError(t, err)
	require.Len(t, resp.GetTableMeta(), 1)
	require.Equal(t, "db1", resp.GetTableMeta()[0].GetDatabase())
}

func TestNGRegionalClientCreateAndGetImportFilesTask(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/files-tasks":
			w.Header().Set("Content-Type", "application/json")
			_, err := w.Write([]byte(`{"taskId":"987"}`))
			require.NoError(t, err)
		case r.Method == http.MethodGet && r.URL.Path == "/api/v1/import/files-tasks/987":
			w.Header().Set("Content-Type", "application/json")
			_, err := w.Write([]byte(`{"status":"IMPORT_FILES_TASK_STATUS_RUNNING"}`))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected route: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	createResp, err := client.CreateImportFilesTask(context.Background(), "987", 123, "s3://bucket/path", -1)
	require.NoError(t, err)
	require.Equal(t, "987", createResp.GetTaskId())
	getResp, err := client.GetImportFilesTask(context.Background(), "987")
	require.NoError(t, err)
	require.Equal(t, agentpb.ImportFilesTaskStatus_IMPORT_FILES_TASK_STATUS_RUNNING, getResp.GetStatus())
}

func TestNGRegionalClientCheckSourceAccessEndpoints(t *testing.T) {
	tests := map[string]struct {
		path         string
		responseBody string
		call         func(t *testing.T, c *NGRegionalClient)
	}{
		"s3": {
			path:         "/api/v1/s3/check_access",
			responseBody: `{"errorMsg":"","region":"us-east-1","pemChanged":true}`,
			call: func(t *testing.T, c *NGRegionalClient) {
				resp, err := c.CheckS3Access(context.Background(), &agentpb.CheckS3AccessReq{})
				require.NoError(t, err)
				require.Equal(t, "", resp.GetErrorMsg())
				require.Equal(t, "us-east-1", resp.GetRegion())
				require.True(t, resp.GetPemChanged())
			},
		},
		"oss": {
			path:         "/api/v1/alioss/check_access",
			responseBody: `{"errorMsg":"","region":"oss-cn-hangzhou","pemChanged":false}`,
			call: func(t *testing.T, c *NGRegionalClient) {
				resp, err := c.CheckAliOSSAccess(context.Background(), &agentpb.CheckAliOSSAccessReq{})
				require.NoError(t, err)
				require.Equal(t, "", resp.GetErrorMsg())
				require.Equal(t, "oss-cn-hangzhou", resp.GetRegion())
				require.False(t, resp.GetPemChanged())
			},
		},
		"azblob": {
			path:         "/api/v1/azblob/check_access",
			responseBody: `{"errorMsg":""}`,
			call: func(t *testing.T, c *NGRegionalClient) {
				resp, err := c.CheckAzblobAccess(context.Background(), &agentpb.CheckAzblobAccessReq{})
				require.NoError(t, err)
				require.Equal(t, "", resp.GetErrorMsg())
			},
		},
	}
	for name, tc := range tests {
		t.Run(name, func(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, tc.path, r.URL.Path)
				w.Header().Set("Content-Type", "application/json")
				_, err := w.Write([]byte(tc.responseBody))
				require.NoError(t, err)
			}))
			defer server.Close()
			setTestSharedHTTPClient(t, server.Client())
			client := &NGRegionalClient{AgentUrl: server.URL}
			tc.call(t, client)
		})
	}
}

func TestNGRegionalClientCreateSchemasTaskEndpoint(t *testing.T) {
	pem := newAgentPublicKeyForTest(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/crypto/pem":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.GetPublicKeyResponse{Pem: pem}))
			require.NoError(t, err)
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/create-schemas-tasks":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.CreateSchemasTaskResp{TaskId: "schema-task"}))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected route: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

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

func TestNGRegionalClientGetSchemasTaskEndpoint(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-task", r.URL.Path)
		_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.GetSchemasTaskResp{
			TaskId: "schema-task",
			Status: agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_RUNNING,
		}))
		require.NoError(t, err)
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	resp, err := client.GetSchemasTask(context.Background(), "schema-task")
	require.NoError(t, err)
	require.Equal(t, agentpb.CreateSchemasTaskStatus_CREATE_SCHEMAS_TASK_STATUS_RUNNING, resp.GetStatus())
}

func TestNGRegionalClientImportDataSizeEstimateTaskEndpoints(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/estimate-data-size-tasks":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.CreateImportDataSizeEstimateTaskResp{TaskId: "estimate-task"}))
			require.NoError(t, err)
		case r.Method == http.MethodGet && r.URL.Path == "/api/v1/import/estimate-data-size-tasks/estimate-task":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.GetImportDataSizeEstimateTaskResp{
				TaskId: "estimate-task",
				Status: agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_RUNNING,
			}))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected route: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	createResp, err := client.CreateImportDataSizeEstimateTask(
		context.Background(),
		"estimate-task",
		123,
		"s3://bucket/path",
		-1,
		nil,
	)
	require.NoError(t, err)
	require.Equal(t, "estimate-task", createResp.GetTaskId())

	getResp, err := client.GetImportDataSizeEstimateTask(context.Background(), "estimate-task")
	require.NoError(t, err)
	require.Equal(t, agentpb.ImportDataSizeEstimateTaskStatus_IMPORT_DATA_SIZE_ESTIMATE_TASK_STATUS_RUNNING, getResp.GetStatus())
}

func TestNGRegionalClientImportIntoJobTaskEndpoints(t *testing.T) {
	pem := newAgentPublicKeyForTest(t)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/crypto/pem":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.GetPublicKeyResponse{Pem: pem}))
			require.NoError(t, err)
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/import-into-job-tasks":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.CreateImportIntoJobTaskResp{TaskId: "task-1"}))
			require.NoError(t, err)
		case r.Method == http.MethodGet && r.URL.Path == "/api/v1/import/import-into-job-tasks/task-1":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.GetImportIntoJobTaskResp{
				TaskId: "task-1",
				Status: agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING,
			}))
			require.NoError(t, err)
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/import-into-job-tasks:batchCreate":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.BatchCreateImportIntoJobTasksResp{
				Tasks: []*agentpb.CreateImportIntoJobTaskResp{{TaskId: "task-1"}},
			}))
			require.NoError(t, err)
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/import-into-job-tasks:batchGet":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.BatchGetImportIntoJobTasksResp{
				Tasks: []*agentpb.GetImportIntoJobTaskResp{{
					TaskId: "task-1",
					Status: agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING,
				}},
			}))
			require.NoError(t, err)
		case r.Method == http.MethodPost && r.URL.Path == "/api/v1/import/import-into-jobs:batchCancel":
			_, err := w.Write(marshalProtoJSONForTest(t, &agentpb.BatchCancelImportIntoJobsResp{}))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected route: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{AgentUrl: server.URL}
	createResp, err := client.CreateImportIntoJobTask(
		context.Background(),
		"task-1",
		123,
		"root",
		"secret-password",
		"IMPORT INTO db1.t1 FROM 's3://bucket/path'",
		nil,
	)
	require.NoError(t, err)
	require.Equal(t, "task-1", createResp.GetTaskId())

	getResp, err := client.GetImportIntoJobTask(context.Background(), "task-1")
	require.NoError(t, err)
	require.Equal(t, agentpb.ImportIntoJobTaskStatus_IMPORT_INTO_JOB_TASK_STATUS_RUNNING, getResp.GetStatus())

	batchCreateResp, err := client.BatchCreateImportIntoJobTasks(
		context.Background(),
		123,
		"root",
		"secret-password",
		[]*agentpb.BatchCreateImportIntoJobTasksReq_Task{{TaskId: "task-1", Sql: "select 1"}},
	)
	require.NoError(t, err)
	require.Len(t, batchCreateResp.GetTasks(), 1)

	batchGetResp, err := client.BatchGetImportIntoJobTasks(context.Background(), []string{"task-1"})
	require.NoError(t, err)
	require.Len(t, batchGetResp.GetTasks(), 1)

	_, err = client.BatchCancelImportIntoJobs(context.Background(), 123, "root", "secret-password", []uint64{1})
	require.NoError(t, err)
}

func TestNGRegionalClientTiKVCapacityReservationEndpoints(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/api/premium/v1/sharedPools/sp-1/tikvCapacityReservations":
			_, err := w.Write([]byte(`{"sharedPoolId":"sp-1","bizType":"TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT","reservationId":"r-1","status":"TIKV_CAPACITY_RESERVATION_STATUS_READY","needScale":true}`))
			require.NoError(t, err)
		case r.Method == http.MethodGet && r.URL.Path == "/api/premium/v1/sharedPools/sp-1/tikvCapacityReservations/TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT/r-1":
			_, err := w.Write([]byte(`{"sharedPoolId":"sp-1","bizType":"TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT","reservationId":"r-1","status":"TIKV_CAPACITY_RESERVATION_STATUS_READY","needScale":false}`))
			require.NoError(t, err)
		case r.Method == http.MethodDelete && r.URL.Path == "/api/premium/v1/sharedPools/sp-1/tikvCapacityReservations/TIKV_CAPACITY_RESERVATION_BIZ_TYPE_IMPORT/r-1":
			_, err := w.Write([]byte(`{}`))
			require.NoError(t, err)
		default:
			t.Fatalf("unexpected route: %s %s", r.Method, r.URL.Path)
		}
	}))
	defer server.Close()

	setTestSharedHTTPClient(t, server.Client())

	client := &NGRegionalClient{BaseUrl: server.URL}
	createResp, err := client.CreateTiKVCapacityReservation(context.Background(), "sp-1", "r-1", 1024)
	require.NoError(t, err)
	require.Equal(t, TiKVCapacityReservationStatusReady, createResp.Status)
	require.True(t, createResp.NeedScale)

	getResp, err := client.GetTiKVCapacityReservation(context.Background(), "sp-1", "r-1")
	require.NoError(t, err)
	require.Equal(t, TiKVCapacityReservationStatusReady, getResp.Status)
	require.False(t, getResp.NeedScale)

	err = client.DeleteTiKVCapacityReservation(context.Background(), "sp-1", "r-1")
	require.NoError(t, err)
}

func newAgentPublicKeyForTest(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 marshalProtoJSONForTest(t *testing.T, msg proto.Message) []byte {
	t.Helper()

	b, err := protojson.Marshal(msg)
	require.NoError(t, err)
	return b
}