package byoc

import (
	"context"
	"fmt"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/acmpca"
	"github.com/aws/aws-sdk-go-v2/service/cloudformation"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ecr"
	"github.com/aws/aws-sdk-go-v2/service/iam"
	"github.com/aws/aws-sdk-go-v2/service/resourcegroupstaggingapi"
	"github.com/aws/aws-sdk-go-v2/service/route53"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/ssm"
	"github.com/aws/aws-sdk-go-v2/service/sts"

	cpaws "github.com/tidbcloud/controlplane9/internal/aws"
)

var (
	loadBYOCAWSConfigQuiet                   = cpaws.LoadAWSConfigQuiet
	loadBYOCAWSConfigFromExportedCredentials = func(ctx context.Context, profile, region string) (awssdk.Config, error) {
		return cpaws.LoadAWSConfigFromExportedProfileCredentials(ctx, profile, region, nil)
	}
)

type byocAWSClients struct {
	acmpca  *acmpca.Client
	cfn     *cloudformation.Client
	ec2     *ec2.Client
	ecr     *ecr.Client
	iam     *iam.Client
	tagging *resourcegroupstaggingapi.Client
	route53 *route53.Client
	s3      *s3.Client
	ssm     *ssm.Client
	sts     *sts.Client
}

func (d *AWSDriver) clients(ctx context.Context, profile, region string) (*byocAWSClients, error) {
	key := profile + "\x00" + region
	d.clientsMu.Lock()
	if d.clientsByProfileRegion == nil {
		d.clientsByProfileRegion = make(map[string]*byocAWSClients)
	}
	clients := d.clientsByProfileRegion[key]
	d.clientsMu.Unlock()
	if clients != nil {
		return clients, nil
	}

	config, err := loadBYOCAWSConfig(ctx, profile, region)
	if err != nil {
		return nil, fmt.Errorf("load AWS SDK profile %s in %s: %w", profile, region, err)
	}
	clients = &byocAWSClients{
		acmpca:  acmpca.NewFromConfig(config),
		cfn:     cloudformation.NewFromConfig(config),
		ec2:     ec2.NewFromConfig(config),
		ecr:     ecr.NewFromConfig(config),
		iam:     iam.NewFromConfig(config),
		tagging: resourcegroupstaggingapi.NewFromConfig(config),
		route53: route53.NewFromConfig(config),
		s3:      s3.NewFromConfig(config),
		ssm:     ssm.NewFromConfig(config),
		sts:     sts.NewFromConfig(config),
	}

	d.clientsMu.Lock()
	if existing := d.clientsByProfileRegion[key]; existing != nil {
		clients = existing
	} else {
		d.clientsByProfileRegion[key] = clients
	}
	d.clientsMu.Unlock()
	return clients, nil
}

func loadBYOCAWSConfig(ctx context.Context, profile, region string) (awssdk.Config, error) {
	config, err := loadBYOCAWSConfigQuiet(ctx, profile, region)
	if err != nil {
		return awssdk.Config{}, err
	}
	if config.Credentials == nil {
		return awssdk.Config{}, fmt.Errorf("AWS profile %s has no credentials provider", profile)
	}
	_, credentialErr := config.Credentials.Retrieve(ctx)
	if credentialErr == nil {
		return config, nil
	}
	if !cpaws.IsSSOCredentialRefreshError(credentialErr) {
		return awssdk.Config{}, fmt.Errorf("resolve AWS credentials for profile %s: %w", profile, credentialErr)
	}
	exported, exportErr := loadBYOCAWSConfigFromExportedCredentials(ctx, profile, region)
	if exportErr != nil {
		return awssdk.Config{}, fmt.Errorf("%w; AWS CLI credential export fallback failed: %w", credentialErr, exportErr)
	}
	return exported, nil
}