package util

import (
	"context"
	"fmt"
	"net/url"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/kms"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

// AssumeRole contains AWS assume role configuration.
type AssumeRole struct {
	RoleARN    string
	ExternalID string
}

// ParsedS3URL holds parsed S3 URL components.
type ParsedS3URL struct {
	Bucket string
	Prefix string
}

// ParseS3URL parses an S3 URL and extracts bucket and prefix.
func ParseS3URL(u *url.URL) *ParsedS3URL {
	prefix := strings.TrimLeft(u.Path, "/")
	return &ParsedS3URL{
		Bucket: u.Host,
		Prefix: prefix,
	}
}

// NewS3ClientWithAssumeRole creates an S3 client with assume role credentials.
// Extra optFns are applied to the s3 client options, e.g. to set a custom
// BaseEndpoint for S3-compatible storages such as Alibaba Cloud OSS.
func NewS3ClientWithAssumeRole(
	ctx context.Context, region string, roles []AssumeRole, optFns ...func(*s3.Options),
) (*s3.Client, error) {
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, fmt.Errorf("failed to load default aws config: %w", err)
	}
	cfg.Credentials = NewAssumeRoleChainProvider(cfg, roles)
	return s3.NewFromConfig(cfg, optFns...), nil
}

// NewS3ClientWithStaticCredentials creates an S3 client with static credentials.
// Extra optFns are applied to the s3 client options, e.g. to set a custom
// BaseEndpoint for S3-compatible storages such as Alibaba Cloud OSS.
func NewS3ClientWithStaticCredentials(
	ctx context.Context, region, accessKeyID, secretAccessKey string, optFns ...func(*s3.Options),
) (*s3.Client, error) {
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, fmt.Errorf("failed to load default aws config: %w", err)
	}
	cfg.Credentials = credentials.NewStaticCredentialsProvider(accessKeyID, secretAccessKey, "")
	return s3.NewFromConfig(cfg, optFns...), nil
}

// NewAssumeRoleChainProvider creates a chained assume role provider.
func NewAssumeRoleChainProvider(cfg aws.Config, roles []AssumeRole) *stscreds.AssumeRoleProvider {
	var chain []*stscreds.AssumeRoleProvider
	for idx, role := range roles {
		opts := func(opts *stscreds.AssumeRoleOptions) {
			if role.ExternalID != "" {
				opts.ExternalID = aws.String(role.ExternalID)
			}
		}
		if idx == 0 {
			p := stscreds.NewAssumeRoleProvider(sts.NewFromConfig(cfg), role.RoleARN, opts)
			chain = append(chain, p)
			continue
		}
		p := stscreds.NewAssumeRoleProvider(sts.NewFromConfig(cfg, func(o *sts.Options) {
			o.Credentials = aws.NewCredentialsCache(chain[idx-1])
		}), role.RoleARN, opts)
		chain = append(chain, p)
	}
	return chain[len(chain)-1]
}

// NewKMSClientWithAssumeRole creates a KMS client with assume role credentials.
func NewKMSClientWithAssumeRole(ctx context.Context, region string, roles []AssumeRole) (*kms.Client, error) {
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, fmt.Errorf("failed to load default aws config: %w", err)
	}
	if len(roles) > 0 && roles[0].RoleARN != "" {
		cfg.Credentials = NewAssumeRoleChainProvider(cfg, roles)
	}
	return kms.NewFromConfig(cfg), nil
}