package awshandler

import (
	"bytes"
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/pingcap/dbaas/billing2/config"
	"github.com/pingcap/dbaas/billing2/utils/log"
	"go.uber.org/zap"
)

type S3MultipartUploader struct {
	s3Cli      *s3.S3
	bucketName string
	key        string
	uploadID   string
	parts      []*s3.CompletedPart
	partNum    int64
}

func NewS3MultipartUploader(key string) (*S3MultipartUploader, error) {
	bucketName := config.Get().GetString(config.S3BillsDetailBucketName)
	region := config.Get().GetString(config.S3BillsDetailRegion)
	s3Cli, err := NewS3ClientWithRegion(region)
	if err != nil {
		return nil, err
	}
	resp, err := s3Cli.CreateMultipartUpload(&s3.CreateMultipartUploadInput{
		Bucket:      &bucketName,
		Key:         &key,
		ContentType: aws.String("application/x-ndjson"),
	})
	if err != nil {
		return nil, err
	}
	return &S3MultipartUploader{
		s3Cli:      s3Cli,
		bucketName: bucketName,
		key:        key,
		uploadID:   *resp.UploadId,
	}, nil
}

func (u *S3MultipartUploader) WritePart(data []byte) error {
	u.partNum++
	resp, err := u.s3Cli.UploadPart(&s3.UploadPartInput{
		Bucket:     &u.bucketName,
		Key:        &u.key,
		UploadId:   &u.uploadID,
		PartNumber: &u.partNum,
		Body:       bytes.NewReader(data),
	})
	if err != nil {
		return err
	}
	u.parts = append(u.parts, &s3.CompletedPart{
		ETag:       resp.ETag,
		PartNumber: &u.partNum,
	})
	return nil
}

func (u *S3MultipartUploader) Abort() error {
	_, err := u.s3Cli.AbortMultipartUpload(&s3.AbortMultipartUploadInput{
		Bucket:   &u.bucketName,
		Key:      &u.key,
		UploadId: &u.uploadID,
	})
	return err
}

func (u *S3MultipartUploader) Complete() error {
	_, err := u.s3Cli.CompleteMultipartUpload(&s3.CompleteMultipartUploadInput{
		Bucket:   &u.bucketName,
		Key:      &u.key,
		UploadId: &u.uploadID,
		MultipartUpload: &s3.CompletedMultipartUpload{
			Parts: u.parts,
		},
	})
	return err
}

func PutBillsDetailToS3(key string, data []byte) error {
	if !config.Get().GetBool(config.EnablePutBillsDetailToS3) {
		log.Info("not enable put bills detail to s3")
		return nil
	}
	bucketName := config.Get().GetString(config.S3BillsDetailBucketName)
	region := config.Get().GetString(config.S3BillsDetailRegion)
	s3Cli, err := NewS3ClientWithRegion(region)
	if err != nil {
		return err
	}
	_, err = s3Cli.PutObject(&s3.PutObjectInput{
		Bucket:      &bucketName,
		Key:         &key,
		Body:        bytes.NewReader(data),
		ContentType: aws.String("application/json"),
	})
	if err != nil {
		log.Error("failed to put bills detail to S3", zap.Error(err), zap.String("key", key))
		return err
	}
	log.Info("success to put bills detail to S3", zap.String("key", key))
	return nil
}

func WriteBillsDetailJSONLToS3(key string, jsonlBytes []byte) error {
	if !config.Get().GetBool(config.EnablePutBillsDetailToS3) {
		log.Info("not enable put bills detail to s3")
		return nil
	}
	if len(jsonlBytes) == 0 {
		return nil
	}
	uploader, err := NewS3MultipartUploader(key)
	if err != nil {
		return err
	}
	chunkSize := 5 * 1024 * 1024
	for offset := 0; offset < len(jsonlBytes); offset += chunkSize {
		end := offset + chunkSize
		if end > len(jsonlBytes) {
			end = len(jsonlBytes)
		}
		if err := uploader.WritePart(jsonlBytes[offset:end]); err != nil {
			uploader.Abort()
			return err
		}
	}
	if err := uploader.Complete(); err != nil {
		return err
	}
	log.Info("success to put bills detail JSONL to S3", zap.String("key", key))
	return nil
}

func GetBillsDetailPresignURL(key string) (string, error) {
	bucketName := config.Get().GetString(config.S3BillsDetailBucketName)
	region := config.Get().GetString(config.S3BillsDetailRegion)
	expireMinutes := config.Get().GetInt(config.S3BillsDetailPresignExpireMinutes)
	s3Cli, err := NewS3ClientWithRegion(region)
	if err != nil {
		return "", err
	}
	req, _ := s3Cli.GetObjectRequest(&s3.GetObjectInput{
		Bucket: &bucketName,
		Key:    &key,
	})
	urlStr, err := req.Presign(time.Duration(expireMinutes) * time.Minute)
	if err != nil {
		return "", err
	}
	return urlStr, nil
}

func CheckS3BillsDetailExists(key string) (bool, error) {
	bucketName := config.Get().GetString(config.S3BillsDetailBucketName)
	region := config.Get().GetString(config.S3BillsDetailRegion)
	s3Cli, err := NewS3ClientWithRegion(region)
	if err != nil {
		return false, err
	}
	_, err = s3Cli.HeadObject(&s3.HeadObjectInput{
		Bucket: &bucketName,
		Key:    &key,
	})
	if err != nil {
		if isS3NotFound(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

func isS3NotFound(err error) bool {
	if aerr, ok := err.(s3.RequestFailure); ok {
		return aerr.StatusCode() == 404
	}
	return false
}

func BuildS3BillsDetailKey(tenantID uint, date string) string {
	env := config.Get().GetString(config.ENV)
	return fmt.Sprintf("%s/bills_detail/%d/%s.json", env, tenantID, date)
}

func BuildS3BillsDetailMonthPrefix(tenantID uint, month string) string {
	env := config.Get().GetString(config.ENV)
	return fmt.Sprintf("%s/bills_detail/%d/%s", env, tenantID, month)
}