package gardener

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	iamtypes "github.com/aws/aws-sdk-go-v2/service/iam/types"
	"github.com/pingcap/errors"

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

func EnsureImportRole(
	ctx context.Context,
	iamClient awsapi.IAMClient,
	accountID string,
	roleName string,
	federated string, // only used for shared provider
	nodeRoleARN string,
	centralRoleARN string,
	skipBanMyBuckets bool,
	tags map[string]string,
) error {
	var awsTags []iamtypes.Tag
	for k, v := range tags {
		awsTags = append(awsTags, iamtypes.Tag{
			Key:   aws.String(k),
			Value: aws.String(v),
		})
	}

	var trusts []string
	if nodeRoleARN != "" {
		trusts = append(trusts, nodeRoleARN)
	}

	if centralRoleARN != "" && centralRoleARN != "TODO" {
		trusts = append(trusts, centralRoleARN)
	}

	p := &awsapi.PolicyDocument{
		Version: "2012-10-17",
		Statement: []awsapi.PolicyStatement{
			{
				Effect: awsapi.StatementEffectAllow,
				Principal: &awsapi.PolicyPrincipal{
					Service:   "ec2.amazonaws.com",
					AWS:       trusts,
					Federated: federated,
				},
				Action: []string{"sts:AssumeRole"},
			},
		},
	}
	if federated != "" {
		p.Statement[0].Action = append(p.Statement[0].Action, "sts:AssumeRoleWithWebIdentity")
	}

	assumePolicy, err := p.ToString()
	if err != nil {
		return err
	}

	roleInput := &iam.CreateRoleInput{
		AssumeRolePolicyDocument: aws.String(assumePolicy),
		Description:              aws.String("role use to import data when running lightning"),
		Path:                     aws.String("/"),
		RoleName:                 aws.String(roleName),
		Tags:                     awsTags,
	}

	err = util.CreateOrUpdateRole(ctx, iamClient, roleInput)
	if err != nil {
		return errors.Annotate(err, "create/update role")
	}

	policyName := roleName

	type Statement struct {
		Sid         *string `json:",omitempty"`
		Effect      string
		Action      string
		Resource    string                 `json:",omitempty"`
		NotResource []string               `json:",omitempty"`
		Condition   map[string]interface{} `json:",omitempty"`
	}

	var statements []*Statement
	statements = append(statements, &Statement{
		Effect:   "Allow",
		Action:   "s3:ListBucket",
		Resource: "*",
	})

	statements = append(statements, &Statement{
		Effect:   "Allow",
		Action:   "s3:GetObject",
		Resource: "*",
	})

	if federated != "" { // for shared provider
		statements = append(statements, &Statement{
			Effect:   "Allow",
			Action:   "sts:AssumeRole",
			Resource: "*",
		})
	}

	if !skipBanMyBuckets {
		statements = append(statements, &Statement{
			Sid:    aws.String("DenyMyBucket"),
			Effect: "Deny",
			Action: "s3:*",
			// We need to skip tidbcloud-samples in internal bucket change
			// Update these codes if there are new sample buckets
			NotResource: []string{
				"arn:aws:s3:::tidbcloud-samples*",
			},
			Condition: map[string]interface{}{
				"StringEquals": map[string]interface{}{
					"s3:ResourceAccount": accountID,
				},
			},
		})
	}

	policy := map[string]interface{}{
		"Version":   "2012-10-17",
		"Statement": statements,
	}

	inlinePolicy, err := json.Marshal(policy)
	if err != nil {
		return errors.AddStack(err)
	}

	inlinePolicyInput := &iam.PutRolePolicyInput{
		PolicyDocument: aws.String(string(inlinePolicy)),
		PolicyName:     aws.String(policyName),
		RoleName:       aws.String(roleName),
	}

	_, err = iamClient.PutRolePolicy(ctx, inlinePolicyInput)
	if err != nil {
		return errors.Annotate(err, "PutRolePolicy")
	}

	if nodeRoleARN != "" {
		err = ensureNodeAssumeImport(ctx, iamClient, roleNameFromARN(nodeRoleARN))
		if err != nil {
			return err
		}
	}

	return nil
}

func ensureNodeAssumeImport(ctx context.Context, client awsapi.IAMClient, nodeRoleName string) error {
	inlinePolicy := `{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": "sts:AssumeRole",
            "Resource": [
                "*"
            ]
        }
    ]
}`

	inlinePolicyInput := &iam.PutRolePolicyInput{
		PolicyDocument: aws.String(inlinePolicy),
		PolicyName:     aws.String("assume-import"),
		RoleName:       aws.String(nodeRoleName),
	}

	_, err := client.PutRolePolicy(ctx, inlinePolicyInput)
	if err != nil {
		return errors.Annotate(err, "PutRolePolicy")
	}

	return nil
}

func roleNameFromARN(arn string) string {
	idx := strings.LastIndex(arn, "/")
	return arn[idx+1:]
}