package aws

import (
	"context"
	"os"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"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/pingcap-inc/clinic/clinic_data_service/pkg/log"
	"go.uber.org/zap"
)

const (
	TargetRoleArn = "arn:aws:iam::552185537300:role/clinic-job-role" // AWS Account ARN: 552185537300
	BucketName    = "delicated-healthcheck"
	Region        = "us-west-2"
	// FilePath      = "./output/sql_diff.json"
	// KeyName       = "raffic-comparison-report/sql_diff.json"
)

func UploadFileToS3(filePath string, keyName string) error {
	ctx := context.TODO()

	// 1. load default config with region
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(Region))
	if err != nil {
		log.Error(ctx, "unable to load SDK config", zap.Error(err))
	}

	// 2. Create STS client
	stsClient := sts.NewFromConfig(cfg)

	// 3. Create an AssumeRole credential provider
	// This automatically performs sts:AssumeRole and manages token refresh
	creds := stscreds.NewAssumeRoleProvider(stsClient, TargetRoleArn)

	// 4. Reload AWS config with temporary credentials from AssumeRole
	// Note: the loaded cfg now assumes the role in account B
	assumedCfg, err := config.LoadDefaultConfig(ctx,
		config.WithRegion(Region),
		config.WithCredentialsProvider(creds),
	)
	if err != nil {
		// log.Fatalf("unable to load assumed config, %v", err)
		log.Error(ctx, "unable to load assumed config", zap.Error(err))

	}

	// 5. Create S3 client with assumed credentials
	s3Client := s3.NewFromConfig(assumedCfg)

	// 6. Upload file
	file, err := os.Open(filePath)
	if err != nil {
		log.Error(ctx, "open file failed", zap.Error(err))
	}
	defer file.Close()

	// fmt.Printf("Uploading file using role %s...\n", TargetRoleArn)
	log.Info(ctx, "upload file to s3", zap.String("filePath", filePath), zap.String("keyName", keyName), zap.String("TargetRoleArn", TargetRoleArn))
	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(BucketName),
		Key:    aws.String(keyName),
		Body:   file,
	})

	if err != nil {
		log.Error(ctx, "upload file to s3 failed", zap.Error(err))
	}

	log.Info(ctx, "upload file to s3 success", zap.String("keyName", keyName))
	return nil
}
