package util

import (
	"context"
	"errors"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/iam"
	iamtypes "github.com/aws/aws-sdk-go-v2/service/iam/types"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/smithy-go"
	"github.com/golang/mock/gomock"
	. "github.com/onsi/gomega"

	"github.com/tidbcloud/infra-api/pkg/providers/aws/fake"
)

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

	m := fake.NewMockIAMClient(ctrl)

	createParams := &iam.CreateRoleInput{
		RoleName:                 aws.String("rname"),
		AssumeRolePolicyDocument: aws.String("assumerole"),
	}

	updateParams := &iam.UpdateAssumeRolePolicyInput{
		PolicyDocument: createParams.AssumeRolePolicyDocument,
		RoleName:       createParams.RoleName,
	}

	// first time call CreateOrUpdateRole
	m.EXPECT().CreateRole(gomock.Any(), gomock.Eq(createParams)).Return(nil, nil)
	// second time call CreateOrUpdateRole
	m.EXPECT().CreateRole(gomock.Any(), gomock.Eq(createParams)).Return(nil, &iamtypes.EntityAlreadyExistsException{})
	m.EXPECT().UpdateAssumeRolePolicy(gomock.Any(), gomock.Eq(updateParams)).Return(nil, nil)

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := CreateOrUpdateRole(ctx, m, createParams)
	g.Expect(err).Should(Succeed())
	err = CreateOrUpdateRole(ctx, m, createParams)
	g.Expect(err).Should(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)

	createParams := &iam.CreateRoleInput{
		RoleName:                 aws.String("rname"),
		AssumeRolePolicyDocument: aws.String("assumerole"),
	}

	putParams := &iam.PutRolePolicyInput{
		RoleName:       createParams.RoleName,
		PolicyName:     aws.String("name"),
		PolicyDocument: aws.String("document"),
	}

	m.EXPECT().CreateRole(gomock.Any(), gomock.Eq(createParams)).Return(nil, nil)
	m.EXPECT().PutRolePolicy(gomock.Any(), gomock.Eq(putParams)).Return(nil, nil)

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := EnsureRole(ctx, m, createParams, putParams)
	g.Expect(err).Should(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)

	createParams := &iam.CreateUserInput{
		UserName: aws.String("username"),
	}

	m.EXPECT().CreateUser(gomock.Any(), gomock.Eq(createParams)).Return(nil, nil)
	m.EXPECT().CreateUser(gomock.Any(), gomock.Eq(createParams)).Return(nil, &iamtypes.EntityAlreadyExistsException{})
	m.EXPECT().CreateUser(gomock.Any(), gomock.Eq(createParams)).Return(nil, &smithy.GenericAPIError{Code: "EntityAlreadyExists"})
	m.EXPECT().CreateUser(gomock.Any(), gomock.Eq(createParams)).Return(nil, errors.New("some error"))

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := EnsureUser(ctx, m, createParams)
	g.Expect(err).Should(Succeed())
	err = EnsureUser(ctx, m, createParams)
	g.Expect(err).Should(Succeed())
	err = EnsureUser(ctx, m, createParams)
	g.Expect(err).Should(Succeed())
	err = EnsureUser(ctx, m, createParams)
	g.Expect(err).ShouldNot(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)

	roleName := "rolename"
	deleteParams := &iam.DeleteRoleInput{
		RoleName: aws.String(roleName),
	}

	listOutput := &iam.ListRolePoliciesOutput{}
	m.EXPECT().ListRolePolicies(gomock.All(), gomock.All()).Return(listOutput, nil).AnyTimes()
	listAttachedOutput := &iam.ListAttachedRolePoliciesOutput{}
	m.EXPECT().ListAttachedRolePolicies(gomock.All(), gomock.Any()).Return(listAttachedOutput, nil).AnyTimes()

	m.EXPECT().DeleteRole(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, nil)
	m.EXPECT().DeleteRole(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, &iamtypes.NoSuchEntityException{})
	m.EXPECT().DeleteRole(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, &smithy.GenericAPIError{Code: "NoSuchEntity"})
	m.EXPECT().DeleteRole(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, errors.New("some error"))

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := DeleteRoleIgnoreNotFound(ctx, m, roleName)
	g.Expect(err).Should(Succeed())
	err = DeleteRoleIgnoreNotFound(ctx, m, roleName)
	g.Expect(err).Should(Succeed())
	err = DeleteRoleIgnoreNotFound(ctx, m, roleName)
	g.Expect(err).Should(Succeed())
	err = DeleteRoleIgnoreNotFound(ctx, m, roleName)
	g.Expect(err).ShouldNot(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)

	userName := "username"
	deleteParams := &iam.DeleteUserInput{
		UserName: aws.String(userName),
	}

	m.EXPECT().DeleteUser(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, nil)
	m.EXPECT().DeleteUser(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, &iamtypes.NoSuchEntityException{})
	m.EXPECT().DeleteUser(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, &smithy.GenericAPIError{Code: "NoSuchEntity"})
	m.EXPECT().DeleteUser(gomock.Any(), gomock.Eq(deleteParams)).Return(nil, errors.New("some error"))

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := DeleteUserIgnoreNotFound(ctx, m, userName)
	g.Expect(err).Should(Succeed())
	err = DeleteUserIgnoreNotFound(ctx, m, userName)
	g.Expect(err).Should(Succeed())
	err = DeleteUserIgnoreNotFound(ctx, m, userName)
	g.Expect(err).Should(Succeed())
	err = DeleteUserIgnoreNotFound(ctx, m, userName)
	g.Expect(err).ShouldNot(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)
	roleName := "rolename"

	listParams := &iam.ListRolePoliciesInput{
		RoleName: aws.String(roleName),
	}
	listOutput := &iam.ListRolePoliciesOutput{
		PolicyNames: []string{"pname"},
	}
	m.EXPECT().ListRolePolicies(gomock.All(), gomock.Eq(listParams)).Return(listOutput, nil)
	m.EXPECT().DeleteRolePolicy(gomock.Any(), &iam.DeleteRolePolicyInput{
		RoleName:   aws.String(roleName),
		PolicyName: aws.String("pname"),
	}).Return(nil, nil)

	listAttachedParams := &iam.ListAttachedRolePoliciesInput{
		RoleName: aws.String(roleName),
	}
	listAttachedOutput := &iam.ListAttachedRolePoliciesOutput{
		AttachedPolicies: []iamtypes.AttachedPolicy{
			{
				PolicyArn: aws.String("arn"),
			},
		},
	}
	m.EXPECT().ListAttachedRolePolicies(gomock.All(), gomock.Eq(listAttachedParams)).Return(listAttachedOutput, nil)
	m.EXPECT().DetachRolePolicy(gomock.Any(), &iam.DetachRolePolicyInput{
		RoleName:  aws.String(roleName),
		PolicyArn: aws.String("arn"),
	}).Return(nil, nil)

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := DeleteAndDetachRolePolicies(ctx, m, roleName)
	g.Expect(err).Should(Succeed())
}

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

	m := fake.NewMockIAMClient(ctrl)
	name := "username"

	listInput := &iam.ListAccessKeysInput{
		UserName: aws.String(name),
	}

	listOutput := &iam.ListAccessKeysOutput{
		AccessKeyMetadata: []iamtypes.AccessKeyMetadata{
			{
				AccessKeyId: aws.String("keyID"),
			},
		},
	}

	deleteInput := &iam.DeleteAccessKeyInput{
		UserName:    aws.String(name),
		AccessKeyId: aws.String("keyID"),
	}

	m.EXPECT().ListAccessKeys(gomock.Any(), listInput).Return(listOutput, nil)
	m.EXPECT().DeleteAccessKey(gomock.Any(), deleteInput).Return(&iam.DeleteAccessKeyOutput{}, nil)

	ctx := context.Background()
	g := NewGomegaWithT(t)
	err := DeleteUserAccessKeys(ctx, m, name)
	g.Expect(err).Should(Succeed())
}