package stashing

import (
	"context"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/dustin/go-humanize"
	"github.com/tidbcloud/pkgv2/log"
	conf "github.com/tidbcloud/serverless-service/internal/config"
	"go.uber.org/zap"
)

const (
	// one organization can upload up to 256GB in per region
	maxSizeLimit = 256 * 1024 * 1024 * 1024
)

type Feature string

const (
	Export Feature = "export"
	Import Feature = "import"
	DM     Feature = "dm"
)

type StashingArea struct {
	preSignerClient *s3.PresignClient
	s3Client        *s3.Client
	Region          string
	Endpoint        string
	AccessKeyID     string
	SecretAccessKey string
}

// GetBucketName returns the bucket name for the given region,
// bucket is created in https://github.com/tidbcloud/aws-shared-cd/pull/2716/files.
func (s *StashingArea) GetBucketName() string {
	prefix := conf.Global().ServerConfig.StashingAreaBucketPrefix
	if prefix == "" {
		prefix = conf.Stage()
	}
	return fmt.Sprintf("%s-%s-stashing-area", prefix, s.Region)
}

func (s *StashingArea) GetFolderName(orgID string) string {
	return fmt.Sprintf("org-%s", orgID)
}

func (s *StashingArea) GetObjectPath(orgID string, feature Feature, id string) string {
	return fmt.Sprintf("%s/%s/%s", s.GetFolderName(orgID), feature, id)
}

func (s *StashingArea) ListAllObjects(ctx context.Context, input *s3.ListObjectsV2Input) ([]types.Object, error) {
	paginator := s3.NewListObjectsV2Paginator(s.s3Client, input)
	var contents []types.Object
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}

		contents = append(contents, page.Contents...)
	}

	return contents, nil
}

func (s *StashingArea) ListObjects(ctx context.Context, input *s3.ListObjectsV2Input) (*s3.ListObjectsV2Output, error) {
	return s.s3Client.ListObjectsV2(ctx, input)
}

func (s *StashingArea) PresignGetObject(
	ctx context.Context,
	objectKey string,
	dur time.Duration,
) (*v4.PresignedHTTPRequest, error) {
	return s.preSignerClient.PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.GetBucketName()),
		Key:    aws.String(objectKey),
	}, s3.WithPresignExpires(dur))
}

func (s *StashingArea) PresignPutObject(
	ctx context.Context,
	objectKey string,
	dur time.Duration,
) (*v4.PresignedHTTPRequest, error) {
	return s.preSignerClient.PresignPutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(s.GetBucketName()),
		Key:    aws.String(objectKey),
	}, s3.WithPresignExpires(dur))
}

func (s *StashingArea) MultipartUpload(
	ctx context.Context,
	objectKey string,
	partNumber int32,
	dur time.Duration,
) (*string, []string, error) {
	logger := log.WithContextL(ctx)

	upload, err := s.s3Client.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
		Bucket: aws.String(s.GetBucketName()),
		Key:    aws.String(objectKey),
	})
	if err != nil {
		logger.Error("failed to create multipart upload", zap.Error(err))
		return nil, nil, err
	}

	urls := make([]string, 0, partNumber)
	for i := 1; i <= int(partNumber); i++ {
		j := int32(i)
		req, err := s.preSignerClient.PresignUploadPart(ctx, &s3.UploadPartInput{
			Bucket:     aws.String(s.GetBucketName()),
			Key:        aws.String(objectKey),
			PartNumber: aws.Int32(j),
			UploadId:   upload.UploadId,
		}, s3.WithPresignExpires(dur))
		if err != nil {
			logger.Error("failed to presign upload part", zap.Error(err))
			return nil, nil, err
		}
		urls = append(urls, req.URL)
	}
	return upload.UploadId, urls, nil
}

func (s *StashingArea) AbortMultipartUpload(
	ctx context.Context,
	objectKey, uploadID string,
) (*s3.AbortMultipartUploadOutput, error) {
	return s.s3Client.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{
		Bucket:   aws.String(s.GetBucketName()),
		Key:      aws.String(objectKey),
		UploadId: aws.String(uploadID),
	})
}

func (s *StashingArea) CompleteMultipartUpload(
	ctx context.Context,
	objectKey, uploadID string,
	parts []types.CompletedPart,
) (*s3.CompleteMultipartUploadOutput, error) {
	return s.s3Client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(s.GetBucketName()),
		Key:      aws.String(objectKey),
		UploadId: aws.String(uploadID),
		MultipartUpload: &types.CompletedMultipartUpload{
			Parts: parts,
		},
	})
}

// CheckExceedSize checks if the stashing area size exceeds the limit. return true if not exceeded, false if exceeded.
func (s *StashingArea) CheckExceedSize(ctx context.Context, orgID string) (bool, error) {
	logger := log.WithContextL(ctx)

	objs, err := s.ListAllObjects(ctx, &s3.ListObjectsV2Input{
		Bucket: aws.String(s.GetBucketName()),
		Prefix: aws.String(s.GetFolderName(orgID)),
	})
	if err != nil {
		logger.Error("failed to list stashing area objects", zap.Error(err))
		return false, err
	}
	var size int64
	for _, obj := range objs {
		size += *obj.Size
	}
	logger.Info(
		"stashing area size",
		zap.String("folder", s.GetFolderName(orgID)),
		zap.String("size", humanize.IBytes(uint64(size))),
	)
	if size > maxSizeLimit {
		logger.Warn("storage limit for stashing area exceeded")
		return true, nil
	}
	return false, nil
}

func (s *StashingArea) DeleteObject(ctx context.Context, objectKey string) (*s3.DeleteObjectOutput, error) {
	return s.s3Client.DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(s.GetBucketName()),
		Key:    aws.String(objectKey),
	})
}

func (s *StashingArea) GetObject(ctx context.Context, input *s3.GetObjectInput) (*s3.GetObjectOutput, error) {
	return s.s3Client.GetObject(ctx, input)
}

func (s *StashingArea) DeleteFolder(ctx context.Context, folder string) error {
	logger := log.WithContextL(ctx)
	result, err := s.ListAllObjects(ctx, &s3.ListObjectsV2Input{
		Bucket: aws.String(s.GetBucketName()),
		Prefix: aws.String(folder),
	})
	if err != nil {
		return err
	}

	errCount := 0
	for _, obj := range result {
		_, err = s.s3Client.DeleteObject(ctx, &s3.DeleteObjectInput{
			Bucket: aws.String(s.GetBucketName()),
			Key:    obj.Key,
		})
		if err != nil {
			errCount++
			logger.Error("failed to delete", zap.String("key", *obj.Key), zap.Error(err))
		}
	}
	if errCount > 0 {
		return fmt.Errorf("failed to delete all object in the folder, failed number: %d", errCount)
	}
	return nil
}