package central

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

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/pingcap/errors"
)

const (
	defaultRegionHint = "us-west-2"
)

type S3URL struct {
	URL    string
	Prefix string
	Bucket string
	region string
}

func checkS3Access(
	ctx context.Context,
	s3URL, accessKeyId, secretAccessKey string,
	roleArn, externalID string,
) error {
	u, err := buildS3URL(s3URL)
	if err != nil {
		return err
	}
	var region string
	if r, err := getBucketRegion(ctx, u); err != nil {
		return err
	} else {
		region = r
	}

	var client *s3.S3
	if len(accessKeyId) == 0 {
		client, err = createClientWithAssumeRole(region, roleArn, externalID)
	} else {
		client, err = createClientWithStaticCredentials(region, accessKeyId, secretAccessKey)
	}
	if err != nil {
		return errors.AddStack(err)
	}

	// test s3:ListBucket
	if _, err := client.ListObjectsV2WithContext(ctx, &s3.ListObjectsV2Input{
		Bucket:  aws.String(u.Bucket),
		Prefix:  aws.String(u.Prefix),
		MaxKeys: aws.Int64(1),
	}); err != nil {
		return errors.Annotate(err, "s3:ListBucket error")
	}
	return nil
}

func buildS3URL(s3URL string) (*S3URL, error) {
	s3URL = strings.TrimSpace(s3URL)
	u, err := url.Parse(s3URL)
	if err != nil {
		return nil, errors.AddStack(err)
	}

	if u.Scheme != "s3" {
		return nil, errors.New("an invalid S3 endpoint URL")
	}
	bucket := u.Host
	prefix := u.Path
	prefix = strings.TrimLeft(prefix, "/")
	if prefix != "" && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	return &S3URL{
		URL:    s3URL,
		Prefix: prefix,
		Bucket: bucket,
		region: "",
	}, nil
}

// Only url starting with S3:// need to get the region
// see more here https://docs.aws.amazon.com/AmazonS3/latest/userguide/access-bucket-intro.html
func getBucketRegion(ctx context.Context, s *S3URL) (string, error) {
	if s.region != "" {
		return s.region, nil
	}
	se := session.Must(session.NewSession())
	region, err := s3manager.GetBucketRegion(ctx, se, s.Bucket, defaultRegionHint)
	if err != nil {
		return "", err
	}
	s.region = region
	return s.region, nil
}

func createClientWithAssumeRole(region string, roleARN string, externalID string) (*s3.S3, error) {
	awsSessionOpts := session.Options{
		Config: *aws.NewConfig().WithRegion(region).WithCredentialsChainVerboseErrors(true),
	}

	ses, err := session.NewSessionWithOptions(awsSessionOpts)
	if err != nil {
		return nil, errors.AddStack(err)
	}

	if roleARN == "" {
		return s3.New(ses), nil
	}
	creds := stscreds.NewCredentials(ses, roleARN, func(p *stscreds.AssumeRoleProvider) {
		p.RoleARN = roleARN
		if externalID != "" {
			p.ExternalID = &externalID
		}
	})
	client := s3.New(ses, &aws.Config{Credentials: creds})

	return client, nil
}

func createClientWithStaticCredentials(region string, accessKeyId string, secretAccessKey string) (*s3.S3, error) {
	awsSessionOpts := session.Options{
		Config: *aws.NewConfig().WithRegion(region).WithCredentialsChainVerboseErrors(true),
	}

	ses, err := session.NewSessionWithOptions(awsSessionOpts)
	if err != nil {
		return nil, errors.AddStack(err)
	}

	creds := credentials.NewStaticCredentials(accessKeyId, secretAccessKey, "")

	client := s3.New(ses, &aws.Config{Credentials: creds})

	return client, nil
}