package premium_dm

import (
	"context"
	"errors"
	"net/url"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service-ng/app/internal/feature_gate"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	"github.com/tidbcloud/dataflow-service-ng/pkg/agent"
	"go.uber.org/zap"
)

type recordingFeatureGate struct {
	enabled bool
	req     *feature_gate.QueryFeatureGateReq
	feature string
}

func (g *recordingFeatureGate) QueryEnabledFeatures(_ *feature_gate.QueryFeatureGateReq) feature_gate.QueryFeatureGateResp {
	return feature_gate.QueryFeatureGateResp{}
}

func (g *recordingFeatureGate) IsFeatureEnabled(req *feature_gate.QueryFeatureGateReq, feature string) bool {
	g.req = req
	g.feature = feature
	return g.enabled && feature == feature_gate.ExplicitImportRole
}

type fakeImportRoleConfigClient struct {
	roleConfig *agent.ImportRoleConfig
	err        error
	calls      int
	clusterID  uint64
}

func (c *fakeImportRoleConfigClient) GetImportRoleConfig(_ context.Context, clusterID uint64) (*agent.ImportRoleConfig, error) {
	c.calls++
	c.clusterID = clusterID
	if c.err != nil {
		return nil, c.err
	}
	return c.roleConfig, nil
}

func TestApplyImportRoleConfigToDataDirIfNeeded(t *testing.T) {
	roleConfig := &agent.ImportRoleConfig{
		RoleArn:    "arn:aws:iam::123456789012:role/import-role",
		ExternalID: "42",
	}

	tests := map[string]struct {
		task         *entity.PremiumDMTask
		enabled      bool
		wantRole     bool
		wantCalls    int
		wantErr      string
		clientErr    error
		inputDataDir string
		wantDataDir  string
		wantRegion   string
		wantProvider string
	}{
		"feature disabled keeps AWS physical data dir unchanged": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "us-west-2",
				CloudProvider: string(entity.CloudProviderTypeAWS),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			wantDataDir: "s3://bucket/path",
		},
		"feature enabled applies role auth to AWS physical data dir": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "us-west-2",
				CloudProvider: string(entity.CloudProviderTypeAWS),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			enabled:   true,
			wantRole:  true,
			wantCalls: 1,
		},
		"feature enabled applies role auth to AliCloud physical data dir": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "ap-southeast-1",
				CloudProvider: string(entity.CloudProviderTypeAliCloud),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			enabled:      true,
			wantRole:     true,
			wantCalls:    1,
			inputDataDir: "s3://bucket/path?endpoint=https%3A%2F%2Foss-ap-southeast-1-internal.aliyuncs.com&provider=alibaba&region=ap-southeast-1",
			wantRegion:   "ap-southeast-1",
			wantProvider: string(entity.CloudProviderTypeAliCloud),
		},
		"feature disabled keeps AliCloud physical data dir unchanged": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "ap-southeast-1",
				CloudProvider: string(entity.CloudProviderTypeAliCloud),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			inputDataDir: "s3://bucket/path?endpoint=https%3A%2F%2Foss-ap-southeast-1-internal.aliyuncs.com&provider=alibaba&region=ap-southeast-1",
			wantDataDir:  "s3://bucket/path?endpoint=https%3A%2F%2Foss-ap-southeast-1-internal.aliyuncs.com&provider=alibaba&region=ap-southeast-1",
		},
		"logical import mode keeps data dir unchanged": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "us-west-2",
				CloudProvider: string(entity.CloudProviderTypeAWS),
				ImportMode:    entity.PremiumDMTaskImportModeLogical,
			},
			enabled:     true,
			wantDataDir: "s3://bucket/path",
		},
		"non-AWS provider keeps data dir unchanged": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "us-west-2",
				CloudProvider: string(entity.CloudProviderTypeGCP),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			enabled:     true,
			wantDataDir: "gcs://bucket/path",
		},
		"role config lookup failure is returned": {
			task: &entity.PremiumDMTask{
				TiDBID:        "42",
				TenantID:      "1",
				Region:        "us-west-2",
				CloudProvider: string(entity.CloudProviderTypeAWS),
				ImportMode:    entity.PremiumDMTaskImportModePhysical,
			},
			enabled:   true,
			wantCalls: 1,
			clientErr: errors.New("csng unavailable"),
			wantErr:   "get Premium DM import role config: csng unavailable",
		},
	}

	for name, tc := range tests {
		t.Run(name, func(t *testing.T) {
			gate := &recordingFeatureGate{enabled: tc.enabled}
			worker := &PremiumDMTaskWorker{FeatureGate: gate}
			client := &fakeImportRoleConfigClient{
				roleConfig: roleConfig,
				err:        tc.clientErr,
			}
			inputDataDir := "s3://bucket/path"
			if tc.inputDataDir != "" {
				inputDataDir = tc.inputDataDir
			} else if tc.wantDataDir != "" {
				inputDataDir = tc.wantDataDir
			}

			got, err := worker.applyImportRoleConfigToDataDirIfNeeded(context.Background(), tc.task, 42, inputDataDir, client)
			if tc.wantErr != "" {
				require.EqualError(t, err, tc.wantErr)
				require.Equal(t, tc.wantCalls, client.calls)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tc.wantCalls, client.calls)

			if tc.wantRole {
				parsed, err := url.Parse(got)
				require.NoError(t, err)
				require.Equal(t, "s3", parsed.Scheme)
				require.Equal(t, roleConfig.RoleArn, parsed.Query().Get("role-arn"))
				require.Equal(t, roleConfig.ExternalID, parsed.Query().Get("external-id"))
				require.Equal(t, uint64(42), client.clusterID)
				require.NotNil(t, gate.req)
				require.Equal(t, uint64(1), gate.req.TenantID)
				require.Equal(t, uint64(42), gate.req.ClusterID)
				wantRegion := "us-west-2"
				if tc.wantRegion != "" {
					wantRegion = tc.wantRegion
				}
				require.Equal(t, wantRegion, gate.req.Region)
				wantProvider := string(entity.CloudProviderTypeAWS)
				if tc.wantProvider != "" {
					wantProvider = tc.wantProvider
				}
				require.Equal(t, wantProvider, gate.req.InfraProvider)
				require.Equal(t, feature_gate.ExplicitImportRole, gate.feature)
				if tc.wantProvider == string(entity.CloudProviderTypeAliCloud) {
					require.Equal(t, "alibaba", parsed.Query().Get("provider"))
					require.Equal(t, "ap-southeast-1", parsed.Query().Get("region"))
					require.Equal(t, "https://oss-ap-southeast-1-internal.aliyuncs.com", parsed.Query().Get("endpoint"))
				}
			} else {
				require.Equal(t, inputDataDir, got)
			}
		})
	}
}

func TestModifyURIForAliCloudOSSPreservesImportRole(t *testing.T) {
	task := &entity.PremiumDMTask{
		CloudProvider: string(entity.CloudProviderTypeAliCloud),
		Region:        "ap-southeast-1",
	}
	uri := "s3://bucket/path?role-arn=arn%3Aaws%3Aiam%3A%3A123456789012%3Arole%2Fimport-role&external-id=42"

	got, err := (&PremiumDMTaskWorker{}).modifyURIForAliCloudOSS(task, &uri, zap.NewNop())
	require.NoError(t, err)

	parsed, err := url.Parse(*got)
	require.NoError(t, err)
	require.Equal(t, "s3", parsed.Scheme)
	require.Equal(t, "alibaba", parsed.Query().Get("provider"))
	require.Equal(t, "ap-southeast-1", parsed.Query().Get("region"))
	require.Equal(t, "https://oss-ap-southeast-1-internal.aliyuncs.com", parsed.Query().Get("endpoint"))
	require.Equal(t, "arn:aws:iam::123456789012:role/import-role", parsed.Query().Get("role-arn"))
	require.Equal(t, "42", parsed.Query().Get("external-id"))
}