package aws

import (
	"context"
	"fmt"
	"os"
	"path"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	stsapi "github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/tidbcloud/metering/pkg/common/utils"
)

const (
	sessionName = "dedicated-metering-server"
)

type AwsProviderS3 struct {
	cli       *s3.Client
	cfg       aws.Config
	rolecache map[string]*aws.CredentialsCache
	bucket    string
}

func NewAwsProviderS3(awscfg aws.Config, bucket string) *AwsProviderS3 {
	return &AwsProviderS3{
		cfg:       awscfg,
		rolecache: make(map[string]*aws.CredentialsCache),
		cli:       s3.NewFromConfig(awscfg),
		bucket:    bucket,
	}
}

func NewAwsProviderS3UseHostStyle(awscfg aws.Config, bucket string) *AwsProviderS3 {
	s3Client := s3.NewFromConfig(awscfg, func(o *s3.Options) {
		o.UsePathStyle = false // Access using virtual-host style
	})
	return &AwsProviderS3{
		cfg:       awscfg,
		rolecache: make(map[string]*aws.CredentialsCache),
		cli:       s3Client,
		bucket:    bucket,
	}
}

func (c AwsProviderS3) AssumeRole(ctx context.Context, roleArn string) *AwsProviderS3 {
	//check if the role is already in the cache
	if cred, ok := c.rolecache[roleArn]; ok {
		return &AwsProviderS3{
			cli: s3.NewFromConfig(c.cfg, func(o *s3.Options) {
				o.Credentials = cred
			}),
			cfg:    c.cfg,
			bucket: c.bucket,
		}
	}
	newCred := aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsapi.NewFromConfig(c.cfg), roleArn, func(p *stscreds.AssumeRoleOptions) {
		p.RoleSessionName = sessionName
	}))
	c.rolecache[roleArn] = newCred
	return &AwsProviderS3{
		cli: s3.NewFromConfig(c.cfg, func(o *s3.Options) {
			o.Credentials = newCred
		}),
		cfg:    c.cfg,
		bucket: c.bucket,
	}
}

// ListObjectInTime lists  object which perfix has time range like 20240901-20241001 in the bucket
// the parameter time is the time point  like 20240801. we will build a range like 20240801-20240901
func (a *AwsProviderS3) ListObjectInTime(ctx context.Context, now int64, perfix string) ([]string, error) {
	dateRange := utils.GenerateDateRangeAWS(now)
	objectperfix := path.Join(perfix, dateRange)
	input := &s3.ListObjectsV2Input{
		Bucket: aws.String(a.bucket),
		Prefix: aws.String(objectperfix),
	}
	var result []string
	// Paginate through the results
	paginator := s3.NewListObjectsV2Paginator(a.cli, input)

	// Iterate through the pages
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		// Print the object keys
		for _, object := range page.Contents {
			result = append(result, *object.Key)
		}
	}
	return result, nil
}

// ListObjectByName lists  objects  which giving name  in the bucket
func (a *AwsProviderS3) ListObjectByName(ctx context.Context, perfix string) ([]string, error) {
	input := &s3.ListObjectsV2Input{
		Bucket: aws.String(a.bucket),
		Prefix: aws.String(perfix),
	}
	var result []string
	// Paginate through the results
	paginator := s3.NewListObjectsV2Paginator(a.cli, input)

	// Iterate through the pages
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		// Print the object keys
		for _, object := range page.Contents {
			result = append(result, *object.Key)
		}
	}
	return result, nil
}

// DownloadObjects downloads blobs to destinationDir
func (a *AwsProviderS3) DownloadObjects(ctx context.Context, names []string, destinationDir string) error {
	for _, name := range names {
		input := &s3.GetObjectInput{
			Bucket: aws.String(a.bucket),
			Key:    aws.String(name),
		}
		output, err := a.cli.GetObject(ctx, input)
		if err != nil {
			return err
		}
		defer output.Body.Close()
		// create file
		filePath := path.Join(destinationDir, path.Base(name))
		file, err := os.Create(filePath)
		if err != nil {
			return fmt.Errorf("failed to create file %s, %v", filePath, err)
		}
		defer file.Close()
		// write to file
		_, err = file.ReadFrom(output.Body)
		if err != nil {
			return fmt.Errorf("failed to write to file %s, %v", filePath, err)
		}
	}
	return nil
}

// UploadObject uploads blob to container
func (a *AwsProviderS3) UploadObject(ctx context.Context, objectName string, sourcefile string) error {
	file, err := os.Open(sourcefile)
	if err != nil {
		return fmt.Errorf("failed to open file %s: %w", sourcefile, err)
	}
	defer file.Close()
	input := &s3.PutObjectInput{
		Bucket: aws.String(a.bucket),
		Key:    aws.String(objectName),
		Body:   file,
	}
	_, err = a.cli.PutObject(ctx, input)
	if err != nil {
		return fmt.Errorf("failed to upload file %s: %w", sourcefile, err)
	}
	return nil
}