package regionalservice

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	v1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/common/configs"
	"github.com/tidbcloud/observability/pkg/common/constants"
	"github.com/tidbcloud/observability/pkg/common/errors"
	"github.com/tidbcloud/observability/pkg/common/message"
	"github.com/tidbcloud/observability/pkg/common/structs"
	"github.com/tidbcloud/observability/pkg/resourcemanager"
	"github.com/tidbcloud/observability/test/mockresourcemanager"
	"go.uber.org/mock/gomock"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func mockCreateIntegrationTaskReq() message.CreateIntegrationTaskReq {
	return message.CreateIntegrationTaskReq{
		DisplayName:     "Test Pipeline",
		DataType:        "slowquery",
		ClusterID:       "cluster-123",
		TenantID:        "tenant-123",
		BizType:         "tidbcloud",
		AccessPrincipal: "arn:aws:iam::123456789012:role/test-role",
		Frequency:       "1h",
		DataFormat:      "parquet",
	}
}

func mockIntegrationTask(req message.CreateIntegrationTaskReq, projectID string) *v1.IntegrationTask {
	taskID := fmt.Sprintf("%s-test-hash", req.DataType)
	namespace := NamespaceName(req.BizType, req.TenantID)

	return &v1.IntegrationTask{
		ObjectMeta: metav1.ObjectMeta{
			Name:      taskID,
			Namespace: namespace,
			Labels: map[string]string{
				constants.LabelApplicationID:       req.ClusterID,
				constants.LabelIntegrationTaskID:   taskID,
				constants.LabelIntegrationTaskType: req.DataType,
				constants.LabelTenantID:            req.TenantID,
				constants.LabelProjectID:           projectID,
				constants.LabelBizType:             req.BizType,
				constants.LabelAlias:               fmt.Sprintf(constants.IntegrationTaskRoleNameTemplate, req.TenantID),
			},
			Annotations: map[string]string{
				constants.AnnotationIntegrationTaskDisplayName: req.DisplayName,
			},
			CreationTimestamp: metav1.Now(),
		},
		Spec: v1.IntegrationTaskSpec{
			TaskType:              v1.TaskType(req.DataType),
			ClusterID:             req.ClusterID,
			AccessRoleARN:         req.AccessPrincipal,
			Frequency:             req.Frequency,
			SourceBucketName:      "test-source-bucket",
			DestinationBucketName: "test-dest-bucket",
			DataFormat:            req.DataFormat,
		},
		Status: v1.IntegrationTaskStatus{
			Phase: constants.ACTIVE,
		},
	}
}

func Test_CreateIntegrationTask(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl)
	resourcemanager.MockClusterReaderWriter(clusterRW)

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"

	t.Run("invalid frequency format", func(t *testing.T) {
		invalidReq := req
		invalidReq.Frequency = "invalid"
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), invalidReq)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("frequency too short", func(t *testing.T) {
		invalidReq := req
		invalidReq.Frequency = "5m"
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), invalidReq)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("frequency too long", func(t *testing.T) {
		invalidReq := req
		invalidReq.Frequency = "25h"
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), invalidReq)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("check existence failed", func(t *testing.T) {
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, errors.New(errors.UnrecognizedError, "")).Times(1)
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.Error(t, err)
	})

	t.Run("task already exists", func(t *testing.T) {
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(true, nil).Times(1)
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("cluster not found", func(t *testing.T) {
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{}, nil).Times(1)
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("list clusters failed", func(t *testing.T) {
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(nil, errors.New(errors.UnrecognizedError, "")).Times(1)
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.Error(t, err)
	})

	t.Run("aws accessPrincipal for alicloud rejects before cluster lookup", func(t *testing.T) {
		originalService := configs.GetServiceInfo()
		defer configs.Init(originalService)

		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AliCloud),
		})

		invalidReq := req
		invalidReq.AccessPrincipal = "arn:aws:iam::123456789012:role/customer-role"
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), invalidReq)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
		assert.Contains(t, err.Error(), "acs:ram::")
	})

	t.Run("create task failed", func(t *testing.T) {
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).Return(errors.New(errors.UnrecognizedError, "")).Times(1)
		_, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.Error(t, err)
	})

	t.Run("normal", func(t *testing.T) {
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.NotEmpty(t, resp.ID)
		assert.Equal(t, req.DisplayName, resp.DisplayName)
		assert.Equal(t, req.DataType, resp.DataType)
		assert.Equal(t, req.ClusterID, resp.ClusterID)
		assert.Equal(t, req.AccessPrincipal, resp.AccessPrincipal)
		assert.Equal(t, req.Frequency, resp.Frequency)
		assert.Equal(t, req.DataFormat, resp.DataFormat)
	})

	t.Run("normal with default frequency", func(t *testing.T) {
		reqNoFreq := req
		reqNoFreq.Frequency = ""
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), reqNoFreq)
		assert.NoError(t, err)
		assert.Equal(t, message.DefaultIntegrationTaskFrequency, resp.Frequency)
	})

	t.Run("normal with default data format", func(t *testing.T) {
		reqNoFormat := req
		reqNoFormat.DataFormat = ""
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), reqNoFormat)
		assert.NoError(t, err)
		assert.Equal(t, message.DefaultIntegrationTaskDataFormat, resp.DataFormat)
	})

	t.Run("normal with default time range", func(t *testing.T) {
		// defaultTimeRange is only applied for internal tasks; external tasks pass StartTime/EndTime as-is
		reqNoTime := req
		reqNoTime.StartTime = ""
		reqNoTime.EndTime = ""

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
					},
				},
			},
		}, nil).Times(1)
		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).Return(nil).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), reqNoTime)
		assert.NoError(t, err)
		// External tasks do not default StartTime/EndTime
		assert.Empty(t, resp.StartTime)
		assert.Empty(t, resp.EndTime)
	})
}

func Test_defaultTimeRange(t *testing.T) {
	t.Run("invalid endTime format", func(t *testing.T) {
		req := &message.CreateIntegrationTaskReq{
			EndTime: "invalid-time",
		}

		err := defaultTimeRange(req)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("invalid startTime format", func(t *testing.T) {
		req := &message.CreateIntegrationTaskReq{
			StartTime: "invalid-time",
			EndTime:   "2026-02-04T12:00:00Z",
		}

		err := defaultTimeRange(req)
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("default endTime when startTime is provided", func(t *testing.T) {
		// When StartTime is provided but EndTime is empty, EndTime defaults to now
		req := &message.CreateIntegrationTaskReq{
			StartTime: "2026-02-04T00:00:00Z",
		}

		before := time.Now().UTC()
		err := defaultTimeRange(req)
		assert.NoError(t, err)
		assert.Equal(t, "2026-02-04T00:00:00Z", req.StartTime)
		parsedEnd, parseErr := time.Parse(time.RFC3339, req.EndTime)
		assert.NoError(t, parseErr)
		assert.WithinDuration(t, before, parsedEnd, 5*time.Second)
	})
}

func Test_validateAccessPrincipalForVendor_AllowsRootPrincipal(t *testing.T) {
	testCases := []struct {
		name            string
		vendor          string
		accessPrincipal string
	}{
		{
			name:            "aws root principal",
			vendor:          string(constants.AWS),
			accessPrincipal: "arn:aws:iam::123456789012:root",
		},
		{
			name:            "alicloud root principal",
			vendor:          string(constants.AliCloud),
			accessPrincipal: "acs:ram::123456789012:root",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateAccessPrincipalForVendor(tc.vendor, tc.accessPrincipal)
			assert.NoError(t, err)
		})
	}
}

func Test_validateAccessPrincipalForVendor_RejectsRoleWithRootSuffix(t *testing.T) {
	testCases := []struct {
		name            string
		vendor          string
		accessPrincipal string
	}{
		{
			name:            "aws role with root suffix",
			vendor:          string(constants.AWS),
			accessPrincipal: "arn:aws:iam::123456789012:role/customer-role:root",
		},
		{
			name:            "alicloud role with root suffix",
			vendor:          string(constants.AliCloud),
			accessPrincipal: "acs:ram::123456789012:role/customer-role:root",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateAccessPrincipalForVendor(tc.vendor, tc.accessPrincipal)
			if assert.Error(t, err) {
				assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
			}
		})
	}
}

func Test_DetailIntegrationTask(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task := mockIntegrationTask(req, projectID)

	t.Run("empty id", func(t *testing.T) {
		_, err := GetIntegrationTaskHandler().DetailIntegrationTask(context.TODO(), message.DetailIntegrationTaskReq{
			ID: "",
		})
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("task not found", func(t *testing.T) {
		integrationTaskRW.EXPECT().GetIntegrationTaskByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New(errors.IntegrationTaskNotFound, "")).Times(1)
		_, err := GetIntegrationTaskHandler().DetailIntegrationTask(context.TODO(), message.DetailIntegrationTaskReq{
			ID:        "not-found",
			ClusterID: req.ClusterID,
		})
		assert.Error(t, err)
	})

	t.Run("normal", func(t *testing.T) {
		integrationTaskRW.EXPECT().GetIntegrationTaskByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(task, nil).Times(1)
		resp, err := GetIntegrationTaskHandler().DetailIntegrationTask(context.TODO(), message.DetailIntegrationTaskReq{
			ID:        task.Name,
			ClusterID: req.ClusterID,
		})
		assert.NoError(t, err)
		assert.Equal(t, task.Name, resp.ID)
		assert.Equal(t, req.DisplayName, resp.DisplayName)
		assert.Equal(t, req.DataType, resp.DataType)
		assert.Equal(t, req.ClusterID, resp.ClusterID)
		assert.Equal(t, string(constants.ACTIVE), resp.Phase)
	})
}

func Test_DeleteIntegrationTask(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task := mockIntegrationTask(req, projectID)

	t.Run("empty id", func(t *testing.T) {
		_, err := GetIntegrationTaskHandler().DeleteIntegrationTask(context.TODO(), message.DeleteIntegrationTaskReq{
			ID: "",
		})
		assert.Error(t, err)
		assert.Equal(t, errors.ParameterInvalid, err.(errors.Error).Code())
	})

	t.Run("task not found", func(t *testing.T) {
		integrationTaskRW.EXPECT().GetIntegrationTaskByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New(errors.IntegrationTaskNotFound, "")).Times(1)
		_, err := GetIntegrationTaskHandler().DeleteIntegrationTask(context.TODO(), message.DeleteIntegrationTaskReq{
			ID:        "not-found",
			ClusterID: req.ClusterID,
		})
		assert.Error(t, err)
	})

	t.Run("delete failed", func(t *testing.T) {
		integrationTaskRW.EXPECT().GetIntegrationTaskByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(task, nil).Times(1)
		integrationTaskRW.EXPECT().DeleteIntegrationTask(gomock.Any(), gomock.Any()).Return(errors.New(errors.UnrecognizedError, "")).Times(1)
		_, err := GetIntegrationTaskHandler().DeleteIntegrationTask(context.TODO(), message.DeleteIntegrationTaskReq{
			ID:        task.Name,
			ClusterID: req.ClusterID,
		})
		assert.Error(t, err)
	})

	t.Run("normal", func(t *testing.T) {
		integrationTaskRW.EXPECT().GetIntegrationTaskByID(gomock.Any(), gomock.Any(), gomock.Any()).Return(task, nil).Times(1)
		integrationTaskRW.EXPECT().DeleteIntegrationTask(gomock.Any(), gomock.Any()).Return(nil).Times(1)
		resp, err := GetIntegrationTaskHandler().DeleteIntegrationTask(context.TODO(), message.DeleteIntegrationTaskReq{
			ID:        task.Name,
			ClusterID: req.ClusterID,
		})
		assert.NoError(t, err)
		assert.Equal(t, task.Name, resp.ID)
		assert.NotEmpty(t, resp.Message)
	})
}

func Test_QueryIntegrationTasks(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task1 := mockIntegrationTask(req, projectID)

	req2 := req
	req2.DataType = "slowquery"
	task2 := mockIntegrationTask(req2, projectID)

	t.Run("query failed", func(t *testing.T) {
		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New(errors.UnrecognizedError, "")).Times(1)
		_, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
		})
		assert.Error(t, err)
	})

	t.Run("normal without filter", func(t *testing.T) {
		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(&v1.IntegrationTaskList{
			Items: []v1.IntegrationTask{*task1, *task2},
		}, nil).Times(1)

		resp, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
		})
		assert.NoError(t, err)
		assert.Equal(t, 2, len(resp.Pipelines))
		assert.Equal(t, 2, resp.TotalSize)
	})

	t.Run("normal with data type filter", func(t *testing.T) {
		dataType := "slowquery"
		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(&v1.IntegrationTaskList{
			Items: []v1.IntegrationTask{*task1},
		}, nil).Times(1)

		resp, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
			DataType:  &dataType,
		})
		assert.NoError(t, err)
		assert.Equal(t, 1, len(resp.Pipelines))
		assert.Equal(t, dataType, resp.Pipelines[0].DataType)
	})

	t.Run("normal with pagination", func(t *testing.T) {
		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(&v1.IntegrationTaskList{
			ListMeta: metav1.ListMeta{
				Continue: "next-token",
			},
			Items: []v1.IntegrationTask{*task1},
		}, nil).Times(1)

		resp, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
			PageSize:  1,
			PageToken: "",
		})
		assert.NoError(t, err)
		assert.Equal(t, 1, len(resp.Pipelines))
		assert.Equal(t, "next-token", resp.NextPageToken)
	})

	t.Run("normal with custom page size", func(t *testing.T) {
		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Any(), gomock.Any()).Return(&v1.IntegrationTaskList{
			Items: []v1.IntegrationTask{*task1, *task2},
		}, nil).Times(1)

		resp, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
			PageSize:  10,
		})
		assert.NoError(t, err)
		assert.Equal(t, 2, len(resp.Pipelines))
	})

	t.Run("external query excludes internal tasks", func(t *testing.T) {
		externalTask := mockIntegrationTask(req2, projectID)
		externalTask.Labels[constants.LabelInternal] = "false"

		integrationTaskRW.EXPECT().QueryIntegrationTasks(gomock.Any(), gomock.Eq(map[string]string{
			constants.LabelApplicationID: req.ClusterID,
			constants.LabelTenantID:      req.TenantID,
			constants.LabelInternal:      "false",
		}), gomock.Any()).Return(&v1.IntegrationTaskList{
			Items: []v1.IntegrationTask{*externalTask},
		}, nil).Times(1)

		resp, err := GetIntegrationTaskHandler().QueryIntegrationTasks(context.TODO(), message.QueryIntegrationTasksReq{
			ClusterID: req.ClusterID,
			TenantID:  req.TenantID,
			BizType:   req.BizType,
		})
		assert.NoError(t, err)
		assert.Equal(t, 1, len(resp.Pipelines))
		assert.Equal(t, externalTask.Name, resp.Pipelines[0].ID)
	})
}

func Test_dataBucketName(t *testing.T) {
	originalService := configs.GetServiceInfo()
	defer func() {
		configs.Init(originalService)
	}()

	t.Run("with specific bucket config", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvDev,
			LogicalRegionName: "us-west-2",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {
					SpecificBucket: "custom-source-bucket",
					ExternalBucket: "custom-dest-bucket",
				},
			},
		})

		src, dst := dataBucketName("tidbcloud", false, "", "", "")
		assert.Equal(t, "custom-source-bucket", src)
		assert.Equal(t, "custom-dest-bucket", dst)
	})

	t.Run("with default bucket config", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvDev,
			LogicalRegionName: "us-west-2",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {},
			},
		})

		src, dst := dataBucketName("tidbcloud", false, "", "", "")
		assert.Contains(t, src, "dev")
		assert.Contains(t, src, "us-west-2")
		assert.Contains(t, dst, "dev")
		assert.Contains(t, dst, "us-west-2")
	})

	t.Run("internal task uses env region and user bucket", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvDev,
			Region:            "eu-central-1",
			LogicalRegionName: "eu-central-1-aaa",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {},
			},
		})

		src, dst := dataBucketName("tidbcloud", true, "user-123", "tenant-xyz", "")
		assert.Contains(t, src, "dev")
		assert.Contains(t, src, "eu-central-1")
		assert.Equal(t, "o11y-user-123-dev-eu-central-1", dst)
	})

	t.Run("aws external task uses per-tenant destination bucket", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvProd,
			Vendor:            string(constants.AWS),
			Region:            "us-west-2",
			LogicalRegionName: "us-west-2-staging",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				// Even with a configured shared external bucket, AWS external
				// tasks must fall through to a per-tenant bucket.
				"tidbcloud": {ExternalBucket: "o11y-external-prod-shared-us-west-2-staging"},
			},
		})

		_, dst := dataBucketName("tidbcloud", false, "", "tenant123", "")
		assert.Equal(t, "o11y-external-prod-tenant123-us-west-2-staging", dst)
	})

	t.Run("alicloud external task defaults to per-tenant destination bucket", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvProd,
			Vendor:            string(constants.AliCloud),
			LogicalRegionName: "cn-hangzhou",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {ExternalBucket: "o11y-external-prod-shared-cn-hangzhou"},
			},
		})

		_, dst := dataBucketName("tidbcloud", false, "", "tenant123", "")
		assert.Equal(t, "o11y-external-prod-tenant123-cn-hangzhou", dst)
	})

	t.Run("alicloud configured shared tenant uses shared destination bucket", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvProd,
			Vendor:            string(constants.AliCloud),
			LogicalRegionName: "cn-hangzhou",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {ExternalBucket: "o11y-external-prod-shared-cn-hangzhou"},
			},
			ControllerConfig: structs.ControllerConfig{
				IntegrationTaskConfig: structs.IntegrationTaskConfig{
					SharedBucketTenantIDs: []string{"tenant123"},
				},
			},
		})

		_, dst := dataBucketName("tidbcloud", false, "", "tenant123", "")
		assert.Equal(t, "o11y-external-prod-shared-cn-hangzhou", dst)
	})

	t.Run("explicit exclusive policy wins over shared tenant config", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Env:               constants.EnvProd,
			Vendor:            string(constants.AliCloud),
			LogicalRegionName: "cn-hangzhou",
			BizTypeConfigs: map[string]structs.BizTypeConfig{
				"tidbcloud": {ExternalBucket: "o11y-external-prod-shared-cn-hangzhou"},
			},
			ControllerConfig: structs.ControllerConfig{
				IntegrationTaskConfig: structs.IntegrationTaskConfig{
					SharedBucketTenantIDs: []string{"tenant123"},
				},
			},
		})

		_, dst := dataBucketName("tidbcloud", false, "", "tenant123", v1.BucketSharePolicyExclusive)
		assert.Equal(t, "o11y-external-prod-tenant123-cn-hangzhou", dst)
	})
}

func Test_buildDestinationInfo(t *testing.T) {
	originalService := configs.GetServiceInfo()
	defer func() {
		configs.Init(originalService)
	}()

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task := mockIntegrationTask(req, projectID)

	t.Run("AWS", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AWS),
			AccountInfo: structs.AccountInfoConfig{
				Account: "123456789012",
			},
		})

		info := buildDestinationInfo(task)
		assert.Contains(t, info.AccessPoint, "s3://")
		assert.Contains(t, info.AccessPoint, task.Spec.DestinationBucketName)
		assert.Contains(t, info.AssumeRoleARN, "arn:aws:iam::")
		assert.Contains(t, info.AssumeRoleARN, "123456789012")
	})

	t.Run("AliCloud", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AliCloud),
			AccountInfo: structs.AccountInfoConfig{
				Account: "123456789012",
			},
		})

		info := buildDestinationInfo(task)
		assert.Contains(t, info.AccessPoint, "oss://")
		assert.Contains(t, info.AccessPoint, task.Spec.DestinationBucketName)
		assert.Contains(t, info.AssumeRoleARN, "acs:ram::")
		assert.Contains(t, info.AssumeRoleARN, "123456789012")
	})

	t.Run("GCP", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.GCP),
		})

		info := buildDestinationInfo(task)
		assert.Empty(t, info.AccessPoint)
		assert.Empty(t, info.AssumeRoleARN)
	})

	t.Run("Internal task uses userId bucket directly", func(t *testing.T) {
		internalTask := mockIntegrationTask(req, projectID)
		internalTask.Labels[constants.LabelInternal] = "true"
		internalTask.Spec.DestinationBucketName = "o11y-task-user-456"
		internalTask.Spec.AccessRoleARN = "arn:aws:iam::123456789012:role/vector-internal-role"

		info := buildDestinationInfo(internalTask)
		assert.Equal(t, "s3://o11y-task-user-456", info.AccessPoint)
		// AssumeRoleARN is not set for internal tasks by buildDestinationInfo
		assert.Empty(t, info.AssumeRoleARN)
	})

	t.Run("Internal task ignores vendor config", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AliCloud),
			AccountInfo: structs.AccountInfoConfig{
				Account: "999999999999",
			},
		})

		internalTask := mockIntegrationTask(req, projectID)
		internalTask.Labels[constants.LabelInternal] = "true"
		internalTask.Spec.DestinationBucketName = "o11y-task-user-789"
		internalTask.Spec.AccessRoleARN = "arn:aws:iam::123456789012:role/vector-internal-role"

		info := buildDestinationInfo(internalTask)
		// Internal tasks always use s3://, regardless of vendor
		assert.Equal(t, "s3://o11y-task-user-789", info.AccessPoint)
		assert.Empty(t, info.AssumeRoleARN)
	})
}

func Test_buildDestinationInfoWithoutProject(t *testing.T) {
	originalService := configs.GetServiceInfo()
	defer func() {
		configs.Init(originalService)
	}()

	req := mockCreateIntegrationTaskReq()
	task := mockIntegrationTask(req, "")
	task.Labels[constants.LabelProjectID] = ""

	t.Run("AWS", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AWS),
			AccountInfo: structs.AccountInfoConfig{
				Account: "123456789012",
			},
		})

		info := buildDestinationInfo(task)
		assert.Equal(t, "s3://test-dest-bucket/log_data/data/tenant-123/cluster-123", info.AccessPoint)
	})

	t.Run("AliCloud", func(t *testing.T) {
		configs.Init(structs.ServiceInfo{
			Vendor: string(constants.AliCloud),
			AccountInfo: structs.AccountInfoConfig{
				Account: "123456789012",
			},
		})

		info := buildDestinationInfo(task)
		assert.Equal(t, "oss://test-dest-bucket/log_data/data/tenant-123/cluster-123", info.AccessPoint)
	})
}

func Test_convertToDetailResp(t *testing.T) {
	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task := mockIntegrationTask(req, projectID)

	resp := convertToDetailResp(task)

	assert.Equal(t, task.Name, resp.ID)
	assert.Contains(t, resp.Name, task.Spec.ClusterID)
	assert.Contains(t, resp.Name, task.Name)
	assert.Equal(t, req.DisplayName, resp.DisplayName)
	assert.Equal(t, req.DataType, resp.DataType)
	assert.Equal(t, req.ClusterID, resp.ClusterID)
	assert.Equal(t, req.AccessPrincipal, resp.AccessPrincipal)
	assert.Equal(t, req.Frequency, resp.Frequency)
	assert.Equal(t, req.DataFormat, resp.DataFormat)
	assert.Equal(t, string(constants.ACTIVE), resp.Phase)
	assert.NotEmpty(t, resp.CreateTime)
}

func Test_convertToSummary(t *testing.T) {
	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"
	task := mockIntegrationTask(req, projectID)

	summary := convertToSummary(task)

	assert.Equal(t, task.Name, summary.ID)
	assert.Equal(t, req.DataType, summary.DataType)
	assert.Equal(t, req.AccessPrincipal, summary.AccessPrincipal)
	assert.Equal(t, req.Frequency, summary.Frequency)
	assert.Equal(t, req.DataFormat, summary.DataFormat)
	assert.NotEmpty(t, summary.CreateTime)
}

func Test_GetIntegrationTaskHandler(t *testing.T) {
	handler1 := GetIntegrationTaskHandler()
	handler2 := GetIntegrationTaskHandler()

	assert.NotNil(t, handler1)
	assert.Equal(t, handler1, handler2, "should return same singleton instance")
}

func Test_CreateIntegrationTaskInternal_AddsLabelInternal(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	originalService := configs.GetServiceInfo()
	defer func() {
		configs.Init(originalService)
	}()

	configs.Init(structs.ServiceInfo{
		ControllerConfig: structs.ControllerConfig{
			IntegrationTaskConfig: structs.IntegrationTaskConfig{
				VectorRoleARN: "arn:aws:iam::123456789012:role/VectorAgentInternalRole",
			},
		},
		BizTypeConfigs: map[string]structs.BizTypeConfig{
			"tidbcloud": {},
		},
	})

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl)
	resourcemanager.MockClusterReaderWriter(clusterRW)

	req := mockCreateIntegrationTaskReq()
	// Internal tasks require StartTime since defaultTimeRange validates it
	req.StartTime = "2026-02-04T00:00:00Z"
	req.EndTime = "2026-02-04T12:00:00Z"
	projectID := "project-123"

	t.Run("internal flag adds LabelInternal", func(t *testing.T) {
		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, "true", task.Labels[constants.LabelInternal])
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTaskInternal(context.TODO(), req)
		assert.NoError(t, err)
		assert.NotEmpty(t, resp.ID)
	})
}

func Test_CreateIntegrationTask_WithNewSpecFields(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl)
	resourcemanager.MockClusterReaderWriter(clusterRW)

	projectID := "project-123"

	t.Run("maps DataPrefix field", func(t *testing.T) {
		req := mockCreateIntegrationTaskReq()
		req.DataPrefix = "my-data-prefix"

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, "my-data-prefix", task.Spec.DataPrefix)
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.Equal(t, "my-data-prefix", resp.DataPrefix)
	})

	t.Run("maps JobType field", func(t *testing.T) {
		req := mockCreateIntegrationTaskReq()
		req.JobType = "Once"

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, v1.JobType("Once"), task.Spec.JobType)
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.Equal(t, "Once", resp.JobType)
	})

	t.Run("maps StartTime and EndTime fields", func(t *testing.T) {
		req := mockCreateIntegrationTaskReq()
		req.StartTime = "2026-02-04T00:00:00Z"
		req.EndTime = "2026-02-04T23:59:59Z"

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, "2026-02-04T00:00:00Z", task.Spec.StartTime)
				assert.Equal(t, "2026-02-04T23:59:59Z", task.Spec.EndTime)
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.Equal(t, "2026-02-04T00:00:00Z", resp.StartTime)
		assert.Equal(t, "2026-02-04T23:59:59Z", resp.EndTime)
	})

	t.Run("maps RawLogComponents field", func(t *testing.T) {
		req := mockCreateIntegrationTaskReq()
		req.RawLogComponents = []string{"loki", "operator", "tidb"}

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, []string{"loki", "operator", "tidb"}, task.Spec.RawLogComponents)
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.Equal(t, []string{"loki", "operator", "tidb"}, resp.RawLogComponents)
	})

	t.Run("handles empty RawLogComponents", func(t *testing.T) {
		req := mockCreateIntegrationTaskReq()
		req.RawLogComponents = []string{}

		integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
		clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
			Items: []v1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Labels: map[string]string{
							constants.LabelProjectID: projectID,
						},
						Namespace: "test-namespace",
					},
				},
			},
		}, nil).Times(1)

		integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
			func(ctx context.Context, task *v1.IntegrationTask) error {
				assert.Equal(t, []string{}, task.Spec.RawLogComponents)
				return nil
			}).Times(1)

		resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
		assert.NoError(t, err)
		assert.Equal(t, []string{}, resp.RawLogComponents)
	})
}

func Test_CreateIntegrationTask_RegularSetsLabelInternalFalse(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	integrationTaskRW := mockresourcemanager.NewMockIntegrationTaskReaderWriter(ctrl)
	resourcemanager.MockIntegrationTaskReaderWriter(integrationTaskRW)

	clusterRW := mockresourcemanager.NewMockClusterReaderWriter(ctrl)
	resourcemanager.MockClusterReaderWriter(clusterRW)

	req := mockCreateIntegrationTaskReq()
	projectID := "project-123"

	integrationTaskRW.EXPECT().IsIntegrationTaskExists(gomock.Any(), gomock.Any()).Return(false, nil).Times(1)
	clusterRW.EXPECT().ListClusters(gomock.Any(), gomock.Any()).Return(&v1.ClusterList{
		Items: []v1.Cluster{
			{
				ObjectMeta: metav1.ObjectMeta{
					Labels: map[string]string{
						constants.LabelProjectID: projectID,
					},
					Namespace: "test-namespace",
				},
			},
		},
	}, nil).Times(1)

	integrationTaskRW.EXPECT().CreateIntegrationTask(gomock.Any(), gomock.Any()).DoAndReturn(
		func(ctx context.Context, task *v1.IntegrationTask) error {
			val, exists := task.Labels[constants.LabelInternal]
			assert.True(t, exists)
			assert.Equal(t, "false", val)
			return nil
		}).Times(1)

	resp, err := GetIntegrationTaskHandler().CreateIntegrationTask(context.TODO(), req)
	assert.NoError(t, err)
	assert.NotEmpty(t, resp.ID)
}