package aws

import (
	"bytes"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/onsi/gomega"

	"github.com/pingcap-inc/clinic/pkg/storage/iface"
)

func Test_PreUpload(t *testing.T) {
	g := gomega.NewGomegaWithT(t)

	fileUUID := "Test_PreUpload"
	uploadID := "testUploadID"

	client := &MockS3Client{
		CreateMultipartUploadFunc: func(input *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error) {
			g.Expect(*input.Bucket).To(gomega.Equal(DefaultBucketName))
			g.Expect(*input.Key).To(gomega.Equal(*tmpKey(fileUUID)))
			return &s3.CreateMultipartUploadOutput{
				UploadId: aws.String(uploadID),
			}, nil
		},
	}

	s := NewAWSStorageWithCli(client, DefaultBucketName)

	uid, err := s.PreUpload(fileUUID)
	g.Expect(err).To(gomega.Succeed())
	g.Expect(*uid).To(gomega.Equal(uploadID))
}

func Test_Upload(t *testing.T) {
	g := gomega.NewGomegaWithT(t)

	etag := "etag"
	fileUUID := "Test_Upload"
	uploadID := "testUploadID"
	serialNum := 1
	body := []byte("body")

	client := &MockS3Client{
		UploadPartFunc: func(input *s3.UploadPartInput) (*s3.UploadPartOutput, error) {
			g.Expect(*input.Bucket).To(gomega.Equal(DefaultBucketName))
			g.Expect(*input.Key).To(gomega.Equal(*tmpKey(fileUUID)))
			g.Expect(*input.UploadId).To(gomega.Equal(uploadID))
			g.Expect(int(*input.PartNumber)).To(gomega.Equal(serialNum))
			// g.Expect(int(*input.ContentLength)).To(gomega.Equal(len(body)))

			return &s3.UploadPartOutput{
				ETag: aws.String(etag),
			}, nil
		},
	}

	s := NewAWSStorageWithCli(client, DefaultBucketName)

	aetag, err := s.Upload(fileUUID, bytes.NewReader(body), iface.UploadConfigRequest{
		SerialNumber: serialNum,
		UploadId:     uploadID,
	})
	g.Expect(err).To(gomega.Succeed())
	g.Expect(*aetag).To(gomega.Equal(etag))
}

func Test_Flush(t *testing.T) {
	g := gomega.NewGomegaWithT(t)

	fileUUID := "Test_Flush"
	uploadID := "testUploadID"

	etags := []iface.ETag{
		{
			ETag:         "1",
			SerialNumber: 1,
		},
		{
			ETag:         "2",
			SerialNumber: 2,
		},
		{
			ETag:         "3",
			SerialNumber: 3,
		},
	}

	c := &MockS3Client{
		CompleteMultipartUploadFunc: func(input *s3.CompleteMultipartUploadInput) (*s3.CompleteMultipartUploadOutput, error) {
			g.Expect(*input.Bucket).To(gomega.Equal(DefaultBucketName))
			g.Expect(*input.Key).To(gomega.Equal(*tmpKey(fileUUID)))
			g.Expect(*input.UploadId).To(gomega.Equal(uploadID))
			g.Expect(len(input.MultipartUpload.Parts)).To(gomega.Equal(len(etags)))

			aetags := make(map[int]string)
			for _, part := range input.MultipartUpload.Parts {
				aetags[int(*part.PartNumber)] = *part.ETag
			}

			for _, t := range etags {
				g.Expect(aetags[t.SerialNumber]).To(gomega.Equal(t.ETag))
			}

			return &s3.CompleteMultipartUploadOutput{}, nil
		},
	}

	s := NewAWSStorageWithCli(c, DefaultBucketName)

	err := s.Flush(fileUUID, iface.FlushConfigRequest{
		UploadId: uploadID,
		Etags:    etags,
	})
	g.Expect(err).To(gomega.Succeed())
}