package aws

import (
	"context"
	"errors"

	"github.com/tidbcloud/observability/pkg/common/metrics"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	iamapi "github.com/aws/aws-sdk-go-v2/service/iam"
	"github.com/aws/aws-sdk-go-v2/service/iam/types"
	stsapi "github.com/aws/aws-sdk-go-v2/service/sts"
)

const (
	iamClientLabel                  = "iam-client"
	sessionName                     = "observability"
	errCreateRoleFailed             = "create new iam role failed"
	errGetRoleFailed                = "get iam role failed"
	errListRolePoliciesFailed       = "list role policies failed"
	errPutRolePolicyFailed          = "put role policy failed"
	errDeleteRolePolicyFailed       = "delete role policy failed"
	errDeleteRoleFailed             = "delete role failed"
	errUpdateAssumeRolePolicyFailed = "update assume role policy failed"
)

type IAMClient struct {
	cli       iamClient
	cfg       aws.Config
	rolecache map[string]*aws.CredentialsCache
}

func NewIAM(cfg aws.Config) IAMClient {
	c := iamapi.NewFromConfig(cfg)
	return IAMClient{
		cli:       c,
		cfg:       cfg,
		rolecache: make(map[string]*aws.CredentialsCache),
	}
}

func (c IAMClient) AssumeRole(ctx context.Context, roleArn string) *IAMClient {
	//check if the role is already in the cache
	if cred, ok := c.rolecache[roleArn]; ok {
		return &IAMClient{
			cli: iamapi.NewFromConfig(c.cfg, func(o *iamapi.Options) {
				o.Credentials = cred
			}),
			cfg: c.cfg,
		}
	}
	newCred := aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsapi.NewFromConfig(c.cfg), roleArn, func(p *stscreds.AssumeRoleOptions) {
		p.RoleSessionName = sessionName
	}))
	c.rolecache[roleArn] = newCred
	return &IAMClient{
		cli: iamapi.NewFromConfig(c.cfg, func(o *iamapi.Options) {
			o.Credentials = newCred
		}),
		cfg: c.cfg,
	}
}

// CreateRole used to create a new IAM role in IAM
func (iam IAMClient) CreateRole(ctx context.Context, name string, trustDoc string) error {
	_, err := iam.cli.CreateRole(ctx, &iamapi.CreateRoleInput{
		RoleName:                 aws.String(name),
		AssumeRolePolicyDocument: aws.String(trustDoc),
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "CreateRole", name, errCreateRoleFailed)
	}
	return err
}

// GetRole used to get  IAM role in IAM
func (iam IAMClient) GetRole(ctx context.Context, name string) (*iamapi.GetRoleOutput, bool, error) {
	role, err := iam.cli.GetRole(ctx, &iamapi.GetRoleInput{
		RoleName: aws.String(name),
	})
	if err != nil {
		var nf *types.NoSuchEntityException
		if errors.As(err, &nf) {
			return nil, false, nil
		}
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "GetRole", name, errGetRoleFailed)
		return nil, false, err
	}
	return role, true, nil
}

// ListRolePolicies used to get role binded inline police lists in IAM
func (iam IAMClient) ListRolePolicies(ctx context.Context, name string, marker *string) (*iamapi.ListRolePoliciesOutput, error) {
	output, err := iam.cli.ListRolePolicies(ctx, &iamapi.ListRolePoliciesInput{
		RoleName: aws.String(name),
		Marker:   marker,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "ListRolePolicies", name, errListRolePoliciesFailed)
	}
	return output, err
}

// PutRolePolicy used to create or modify inline policies in  IAM role
func (iam IAMClient) PutRolePolicy(ctx context.Context, name string, policyName string, policyDoc string) error {
	_, err := iam.cli.PutRolePolicy(ctx, &iamapi.PutRolePolicyInput{
		RoleName:       aws.String(name),
		PolicyName:     aws.String(policyName),
		PolicyDocument: aws.String(policyDoc),
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "PutRolePolicy", name, errPutRolePolicyFailed)
	}
	return err
}

// DeleteRolePolicy used to delete inline policies in  IAM role
func (iam IAMClient) DeleteRolePolicy(ctx context.Context, name string, policyName string) error {
	_, err := iam.cli.DeleteRolePolicy(ctx, &iamapi.DeleteRolePolicyInput{
		RoleName:   aws.String(name),
		PolicyName: aws.String(policyName),
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "DeleteRolePolicy", name, errDeleteRolePolicyFailed)
	}
	return err
}

// DeleteRole used to delete IAM role
func (iam IAMClient) DeleteRole(ctx context.Context, name string) error {
	_, err := iam.cli.DeleteRole(ctx, &iamapi.DeleteRoleInput{
		RoleName: aws.String(name),
	})
	if err != nil {
		var nf *types.NoSuchEntityException
		if errors.As(err, &nf) {
			return nil
		}
		return err
	}
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "DeleteRole", name, errDeleteRoleFailed)
	}
	return err
}

// UpdateAssumeRolePolicy used to update assume policy in IAM role
func (iam IAMClient) UpdateAssumeRolePolicy(ctx context.Context, name string, assumeDocument string) error {
	_, err := iam.cli.UpdateAssumeRolePolicy(ctx, &iamapi.UpdateAssumeRolePolicyInput{
		RoleName:       aws.String(name),
		PolicyDocument: aws.String(assumeDocument),
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(iam.cfg.Region, iamClientLabel, "UpdateAssumeRolePolicy", name, errUpdateAssumeRolePolicyFailed)
	}
	return err
}