package impl

import (
	"context"
	"encoding/json"
	"net/url"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/request"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/iam/iamiface"
	"github.com/stretchr/testify/require"
	restoreaccess "github.com/tidbcloud/dataflow-service/app/models/classic_restore_access"
	"github.com/tidbcloud/dataflow-service/app/repo/entity"
	"github.com/tidbcloud/pkgv2/bizerr"
)

func TestSourcePermissionsPolicyIsScopedToBackupPrefixAndKMSKey(t *testing.T) {
	policy, err := sourcePermissionsPolicy(
		"dedicated-backup", "v1/backups/42", "arn:aws:kms:us-west-2:123456789012:key/source-key",
	)
	require.NoError(t, err)
	var document map[string]any
	require.NoError(t, json.Unmarshal([]byte(policy), &document))
	require.Contains(t, policy, "arn:aws:s3:::dedicated-backup/v1/backups/42/*")
	require.Contains(t, policy, "kms:Decrypt")
	require.Contains(t, policy, "kms:DescribeKey")
	require.NotContains(t, policy, "kms:Encrypt")
	require.NotContains(t, policy, `"Resource":"*"`)
}

func TestAccessIdentifiersAreDeterministicAndBucketScoped(t *testing.T) {
	role1, external1 := accessIdentifiers("restore-1", 10)
	role2, external2 := accessIdentifiers("restore-1", 10)
	role3, external3 := accessIdentifiers("restore-1", 11)
	require.Equal(t, role1, role2)
	require.Equal(t, external1, external2)
	require.NotEqual(t, role1, role3)
	require.NotEqual(t, external1, external3)
	require.LessOrEqual(t, len(role1), 64)
}

func TestValidateRoleARN(t *testing.T) {
	require.NoError(t, validateRoleARN("arn:aws:iam::123456789012:role/premium-worker"))
	require.Error(t, validateRoleARN("arn:aws:iam::123456789012:user/not-a-role"))
	require.Error(t, validateRoleARN("not-an-arn"))
}

func TestEnsureRoleAttachesRequiredPermissionsBoundary(t *testing.T) {
	client := &fakeIAMClient{
		createRoleOutput: &iam.CreateRoleOutput{Role: &iam.Role{Arn: aws.String("arn:aws:iam::123456789012:role/test")}},
	}
	boundaryARN := "arn:aws:iam::123456789012:policy/" + permissionsBoundaryName
	roleARN, err := ensureRole(
		context.Background(), client, "test", `{}`, `{}`, boundaryARN, "request-hash", time.Now(),
	)
	require.NoError(t, err)
	require.Equal(t, "arn:aws:iam::123456789012:role/test", roleARN)
	require.Equal(t, boundaryARN, aws.StringValue(client.createRoleInput.PermissionsBoundary))
	require.Contains(t, client.createRoleInput.Tags, &iam.Tag{
		Key: aws.String(requestHashTag), Value: aws.String("request-hash"),
	})
}

func TestEnsureRoleRejectsExistingRoleParameterDrift(t *testing.T) {
	expectedTrust := `{"Version":"2012-10-17","Statement":[]}`
	client := &fakeIAMClient{
		createRoleErr: awserr.New(iam.ErrCodeEntityAlreadyExistsException, "exists", nil),
		getRoleOutput: &iam.GetRoleOutput{Role: &iam.Role{
			Arn:                      aws.String("arn:aws:iam::123456789012:role/tidb-cloud/classic-restore/test"),
			Path:                     aws.String(rolePath),
			AssumeRolePolicyDocument: aws.String(url.QueryEscape(`{"Version":"2012-10-17","Statement":[{"Effect":"Allow"}]}`)),
			PermissionsBoundary: &iam.AttachedPermissionsBoundary{
				PermissionsBoundaryArn: aws.String("arn:aws:iam::123456789012:policy/" + permissionsBoundaryName),
			},
		}},
	}
	_, err := ensureRole(
		context.Background(), client, "test", expectedTrust, `{}`,
		"arn:aws:iam::123456789012:policy/"+permissionsBoundaryName, "request-hash", time.Now(),
	)
	require.ErrorContains(t, err, "trust policy does not match")
	require.Nil(t, client.tagRoleInput)
}

func TestAccessRequestHashRejectsImmutableParameterDrift(t *testing.T) {
	req := &restoreaccess.PrepareReq{
		BackupID: 1, BucketID: 2, RestoreID: "restore-1",
		TrustedRoleARN: "arn:aws:iam::210987654321:role/restore",
	}
	hash1, err := accessRequestHash(req, "123456789012", "bucket", "backup/1", "kms", "trust", "permissions", "boundary")
	require.NoError(t, err)

	drifted := *req
	drifted.BackupID = 3
	hash2, err := accessRequestHash(&drifted, "123456789012", "bucket", "backup/1", "kms", "trust", "permissions", "boundary")
	require.NoError(t, err)
	require.NotEqual(t, hash1, hash2)

	drifted = *req
	drifted.TrustedRoleARN = "arn:aws:iam::210987654321:role/other"
	hash3, err := accessRequestHash(&drifted, "123456789012", "bucket", "backup/1", "kms", "trust", "permissions", "boundary")
	require.NoError(t, err)
	require.NotEqual(t, hash1, hash3)
}

func TestCleanupUsesPersistedAccountWithoutBucketLookup(t *testing.T) {
	accessRepo := &fakeAccessRepo{access: &entity.ClassicRestoreAccess{
		Updatable: entity.Updatable{Base: entity.Base{ID: 9}},
		RestoreID: "restore-1", BucketID: 2, AWSAccountID: "123456789012", RoleName: "temporary-role",
	}}
	client := &fakeIAMClient{}
	proxy := &fakeIAMProxy{client: client}
	m := &model{accessRepo: accessRepo, iamProxy: proxy}

	err := m.Cleanup(context.Background(), &restoreaccess.CleanupReq{RestoreID: "restore-1", BucketID: 2})
	require.NoError(t, err)
	require.Equal(t, "123456789012", proxy.accountID)
	require.Equal(t, "temporary-role", aws.StringValue(client.deleteRoleInput.RoleName))
	require.Equal(t, entity.ID(9), accessRepo.deletedID)
}

func TestReconcileExpiredUsesPersistedAccounts(t *testing.T) {
	accessRepo := &fakeAccessRepo{expired: []entity.ClassicRestoreAccess{{
		Updatable:    entity.Updatable{Base: entity.Base{ID: 10}},
		AWSAccountID: "123456789012", RoleName: "expired-role",
	}}}
	client := &fakeIAMClient{}
	proxy := &fakeIAMProxy{client: client}
	m := &model{accessRepo: accessRepo, iamProxy: proxy}

	err := m.ReconcileExpired(context.Background())
	require.NoError(t, err)
	require.Equal(t, "123456789012", proxy.accountID)
	require.Equal(t, "expired-role", aws.StringValue(client.deleteRoleInput.RoleName))
	require.Equal(t, entity.ID(10), accessRepo.deletedID)
}

type fakeIAMClient struct {
	iamiface.IAMAPI
	createRoleInput  *iam.CreateRoleInput
	createRoleOutput *iam.CreateRoleOutput
	createRoleErr    error
	getRoleOutput    *iam.GetRoleOutput
	getRoleErr       error
	getPolicyOutput  *iam.GetRolePolicyOutput
	getPolicyErr     error
	tagRoleInput     *iam.TagRoleInput
	deleteRoleInput  *iam.DeleteRoleInput
}

func (f *fakeIAMClient) CreateRoleWithContext(
	_ aws.Context, input *iam.CreateRoleInput, _ ...request.Option,
) (*iam.CreateRoleOutput, error) {
	f.createRoleInput = input
	return f.createRoleOutput, f.createRoleErr
}

func (f *fakeIAMClient) GetRoleWithContext(
	_ aws.Context, _ *iam.GetRoleInput, _ ...request.Option,
) (*iam.GetRoleOutput, error) {
	return f.getRoleOutput, f.getRoleErr
}

func (f *fakeIAMClient) GetRolePolicyWithContext(
	_ aws.Context, _ *iam.GetRolePolicyInput, _ ...request.Option,
) (*iam.GetRolePolicyOutput, error) {
	return f.getPolicyOutput, f.getPolicyErr
}

func (f *fakeIAMClient) TagRoleWithContext(
	_ aws.Context, input *iam.TagRoleInput, _ ...request.Option,
) (*iam.TagRoleOutput, error) {
	f.tagRoleInput = input
	return &iam.TagRoleOutput{}, nil
}

func (f *fakeIAMClient) DeleteRolePolicyWithContext(
	_ aws.Context, _ *iam.DeleteRolePolicyInput, _ ...request.Option,
) (*iam.DeleteRolePolicyOutput, error) {
	return &iam.DeleteRolePolicyOutput{}, nil
}

func (f *fakeIAMClient) DeleteRoleWithContext(
	_ aws.Context, input *iam.DeleteRoleInput, _ ...request.Option,
) (*iam.DeleteRoleOutput, error) {
	f.deleteRoleInput = input
	return &iam.DeleteRoleOutput{}, nil
}

type fakeIAMProxy struct {
	accountID string
	client    iamiface.IAMAPI
}

func (f *fakeIAMProxy) GetClient(_ context.Context, accountID string) (iamiface.IAMAPI, error) {
	f.accountID = accountID
	return f.client, nil
}

type fakeAccessRepo struct {
	access    *entity.ClassicRestoreAccess
	expired   []entity.ClassicRestoreAccess
	deletedID entity.ID
}

func (f *fakeAccessRepo) GetOrCreate(
	_ context.Context, access *entity.ClassicRestoreAccess,
) (*entity.ClassicRestoreAccess, bizerr.BizError) {
	return access, nil
}

func (f *fakeAccessRepo) Get(
	_ context.Context, _ string, _ uint64,
) (*entity.ClassicRestoreAccess, bizerr.BizError) {
	return f.access, nil
}

func (f *fakeAccessRepo) ListExpired(
	_ context.Context, _ time.Time,
) ([]entity.ClassicRestoreAccess, bizerr.BizError) {
	return f.expired, nil
}

func (f *fakeAccessRepo) Delete(_ context.Context, id entity.ID) bizerr.BizError {
	f.deletedID = id
	return nil
}