package premium_dm

import (
	"context"
	"encoding/json"
	"testing"

	dm "github.com/pingcap/tiflow/dm/openapi"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
)

func TestCreatePremiumDMTaskWithMetricLabels(t *testing.T) {
	original := getTiDBHostFromAgentFunc
	getTiDBHostFromAgentFunc = func(context.Context, *agent.NGRegionalClient, string) (string, error) {
		return "tidb.internal", nil
	}
	t.Cleanup(func() { getTiDBHostFromAgentFunc = original })
	tests := map[string]struct {
		mode       dm.TaskTaskMode
		importMode dm.TaskFullMigrateConfImportMode
	}{
		"full logical":  {mode: dm.TaskTaskModeAll, importMode: dm.TaskFullMigrateConfImportModeLogical},
		"full physical": {mode: dm.TaskTaskModeAll, importMode: dm.TaskFullMigrateConfImportModeImportInto},
		"incremental":   {mode: dm.TaskTaskModeIncremental},
	}
	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			transport := &mockDMOpenAPIClient{status: 200, respBody: []byte(`{}`)}
			proxy := &PremiumDMProxy{useServiceProxy: true, dmOpenAPI: transport}
			req := &dm.CreateTaskRequest{Task: dm.Task{Name: "dm-task-1", TaskMode: tc.mode}}
			if tc.importMode != "" {
				req.Task.SourceConfig.FullMigrateConf = &dm.TaskFullMigrateConf{ImportMode: &tc.importMode}
			}
			_, err := createPremiumDMTaskWithMetricLabels(context.Background(), proxy, req, "10001")
			require.NoError(t, err)
			require.Equal(t, "POST", transport.lastMethod)
			var payload struct {
				Task struct {
					Name         string            `json:"name"`
					MetricLabels map[string]string `json:"metric_labels"`
				} `json:"task"`
			}
			require.NoError(t, json.Unmarshal(transport.lastBody, &payload))
			require.Equal(t, "dm-task-1", payload.Task.Name)
			require.Equal(t, map[string]string{"keyspace_name": "10001"}, payload.Task.MetricLabels)
		})
	}
}