package aws

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	aws_config "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/s3"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/aws/smithy-go/logging"
	"github.com/tidbcloud/serverless-service/internal/config"
)

type S3ClientProvider struct {
	s3Clients     map[string]*s3.Client
	o11yS3Clients map[string]*s3.Client
}

const (
	O11yRole        = "arn:aws:iam::557537366020:role/autonomous-o11y-s3-readonly"
	O11ySessionName = "serverless_download_s3_logs"
)

func NewS3ClientProvider() *S3ClientProvider {
	s3Clients := make(map[string]*s3.Client)
	o11yS3Clients := make(map[string]*s3.Client)

	for _, regionConfig := range config.Global().WorkerConfig.RegionalConfig {
		key := getS3ClientKey(regionConfig.Provider, regionConfig.Region)
		switch regionConfig.Provider {
		case config.ProviderAWS:
			// config credential expiration time and cache options
			awsCfg, err := aws_config.LoadDefaultConfig(
				context.Background(),
				aws_config.WithWebIdentityRoleCredentialOptions(func(opts *stscreds.WebIdentityRoleOptions) {
					opts.Duration = 12 * time.Hour
				}),
				aws_config.WithCredentialsCacheOptions(func(opts *aws.CredentialsCacheOptions) {
					opts.ExpiryWindow = 8 * time.Hour
				}),
				aws_config.WithRegion(regionConfig.Region),
			)
			if err != nil {
				panic(fmt.Sprintf("load default config for aws failed: %s", err))
			}
			s3Clients[key] = s3.NewFromConfig(awsCfg)

			// o11y s3client
			awsCfg, err = aws_config.LoadDefaultConfig(context.Background(), aws_config.WithRegion(regionConfig.Region))
			if err != nil {
				panic(fmt.Sprintf("load default config for aws failed: %s", err))
			}
			stsSvc := sts.NewFromConfig(awsCfg)
			assumeRoleProvider := stscreds.NewAssumeRoleProvider(
				stsSvc,
				O11yRole,
				func(opts *stscreds.AssumeRoleOptions) {
					opts.RoleSessionName = O11ySessionName
				},
			)
			newCfg, err := aws_config.LoadDefaultConfig(
				context.Background(),
				aws_config.WithRegion(regionConfig.Region),
				aws_config.WithCredentialsProvider(assumeRoleProvider),
				// awsConfig.WithClientLogMode(aws.LogRetries|aws.LogRequestWithBody|aws.LogResponseWithBody),
			)
			if err != nil {
				panic(fmt.Sprintf("load default config for aws failed: %s", err))
			}
			o11yS3Clients[regionConfig.Region] = s3.NewFromConfig(newCfg)
		case config.ProviderAliCloud:
			region := regionConfig.Region
			endpoint := fmt.Sprintf("https://oss-%s.aliyuncs.com", region)
			cfg := aws.Config{
				Region:        region,
				ClientLogMode: aws.LogRequestWithBody | aws.LogResponseWithBody,
				BaseEndpoint:  aws.String(endpoint),
				Logger:        logging.NewStandardLogger(os.Stderr),
				Credentials: credentials.NewStaticCredentialsProvider(
					config.Global().OssConfig.AccessKeyID,
					config.Global().OssConfig.SecretAccessKey,
					"",
				),
			}
			s3Clients[key] = s3.NewFromConfig(cfg)
		}
	}

	return &S3ClientProvider{
		s3Clients:     s3Clients,
		o11yS3Clients: o11yS3Clients,
	}
}

func (s *S3ClientProvider) GetS3Client(provider, region string) (*s3.Client, error) {
	key := getS3ClientKey(provider, region)
	if s.s3Clients[key] == nil {
		return nil, fmt.Errorf("s3 client for provider %s and region %s not found", provider, region)
	}
	return s.s3Clients[key], nil
}

func getS3ClientKey(provider, region string) string {
	return fmt.Sprintf("%s/%s", provider, region)
}

func (s *S3ClientProvider) GetO11yS3Client(region string) (*s3.Client, error) {
	if s.o11yS3Clients[region] == nil {
		return nil, fmt.Errorf("o11y s3 client for region %s not found", region)
	}
	return s.o11yS3Clients[region], nil
}