package manager

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"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/config"
	"github.com/pingcap-inc/clinic/pkg/db"
	"github.com/pingcap-inc/clinic/pkg/db/model"
	aws2 "github.com/pingcap-inc/clinic/pkg/storage/aws"
)

func Test_PreUploadFileWithNewFile(t *testing.T) {
	fileUUID := "Test_PreUploadFileWithNewFile1"
	uploadID := "Test_PreUploadFileWithNewFilee_UploadID"
	alias := "myalias1"
	var fileLen int64 = 1

	g := gomega.NewGomegaWithT(t)

	client := aws2.NewAWSStorageWithCli(&aws2.MockS3Client{
		CreateMultipartUploadFunc: func(input *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error) {
			g.Expect(*input.Key).To(gomega.Equal(fmt.Sprintf("%s/%s", aws2.TmpPrefix, fileUUID)))
			return &s3.CreateMultipartUploadOutput{
				UploadId: aws.String(uploadID),
			}, nil
		},
	}, aws2.DefaultBucketName)

	manager := NewUploadFileManager(client)

	uf := &db.UploadFile{
		FileUUID:  fileUUID,
		ClusterID: "1",
		Alias:     alias,
		FileLen:   fileLen,
	}

	serialNum, err := manager.PreUploadFile(uf)

	g.Expect(err).To(gomega.Succeed())

	g.Expect(serialNum).To(gomega.Equal(0))

	uf, derr := model.QueryUploadFile(fileUUID)
	g.Expect(derr).To(gomega.Succeed())

	g.Expect(uf.UploadID).To(gomega.Equal(uploadID))
	g.Expect(uf.FileLen).To(gomega.Equal(fileLen))
	g.Expect(uf.Alias).To(gomega.Equal(alias))
}

func Test_PreUploadFileWithExistedFile(t *testing.T) {
	fileUUID := "Test_PreUploadFileWithExistedFile1"
	uploadID := "Test_PreUploadFileWithExistedFile_UPLOADID"
	alias := "myalias"
	var fileLen int64 = 1

	uf := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   "1",
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
	}

	err := model.InsertUploadFile(uf)

	g := gomega.NewGomegaWithT(t)
	g.Expect(err).To(gomega.Succeed())

	client := aws2.NewAWSStorageWithCli(&aws2.MockS3Client{
		CreateMultipartUploadFunc: func(input *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error) {
			return &s3.CreateMultipartUploadOutput{
				UploadId: aws.String(uploadID),
			}, nil
		},
	}, aws2.DefaultBucketName)

	manager := NewUploadFileManager(client)
	serialNum, err := manager.PreUploadFile(uf)
	g.Expect(err).To(gomega.Succeed())
	g.Expect(serialNum).To(gomega.Equal(0))
}

func Test_PreUploadFileWithExistedFileWithParts(t *testing.T) {
	fileUUID := "Test_PreUploadFileWithExistedFileWithParts1"
	uploadID := "Test_PreUploadFileWithExistedFileWithParts_UPLOADID"
	alias := "myalias"
	var fileLen int64 = 1
	serialCount := 20

	uf := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   "1",
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
	}

	err := model.InsertUploadFile(uf)

	g := gomega.NewGomegaWithT(t)
	g.Expect(err).To(gomega.Succeed())

	for i := 1; i <= serialCount; i++ {
		err = model.UpsertUploadFilePart(fileUUID, i, fmt.Sprintf("etag%d", i), 1)
		g.Expect(err).To(gomega.Succeed())
	}

	client := aws2.NewAWSStorageWithCli(&aws2.MockS3Client{
		CreateMultipartUploadFunc: func(input *s3.CreateMultipartUploadInput) (*s3.CreateMultipartUploadOutput, error) {
			return &s3.CreateMultipartUploadOutput{
				UploadId: aws.String(uploadID),
			}, nil
		},
	}, aws2.DefaultBucketName)

	manager := NewUploadFileManager(client)

	uf2 := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   fileUUID,
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
	}
	serailNum, err := manager.PreUploadFile(uf2)
	g.Expect(err).To(gomega.Succeed())
	g.Expect(serailNum).To(gomega.Equal(serialCount))
}

func Test_UploadFile(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fileUUID := "Test_Manager_UploadFile1"
	uploadID := "Test_Manager_UploadFile_UPLOADID"
	serialNum := 20
	etag := "etag"
	alias := "myalias"
	var fileLen int64 = 1
	content := []byte("Test_UploadFile")

	uf := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   "1",
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
		UploadID:    uploadID,
	}

	err := model.InsertUploadFile(uf)
	g.Expect(err).To(gomega.Succeed())

	client := &aws2.MockS3Client{
		UploadPartFunc: func(input *s3.UploadPartInput) (*s3.UploadPartOutput, error) {
			g.Expect(*input.Key).To(gomega.Equal(fmt.Sprintf("%s/%s", aws2.TmpPrefix, 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(content)))
			return &s3.UploadPartOutput{
				ETag: aws.String(etag),
			}, nil
		},
	}

	manager := NewUploadFileManager(aws2.NewAWSStorageWithCli(client, aws2.DefaultBucketName))

	err = manager.UploadFile(fileUUID, bytes.NewReader(content), serialNum, 1)
	g.Expect(err).To(gomega.Succeed())

	ufp, err := model.QueryUploadFilePart(fileUUID, serialNum)
	g.Expect(err).To(gomega.Succeed())

	g.Expect(ufp.Etag).To(gomega.Equal(etag))
}

func Test_UploadFileFailed(t *testing.T) {
	fileUUID := "Test_Manager_UploadFileFailed1"
	manager := NewUploadFileManager(aws2.NewAWSStorageWithCli(nil, aws2.DefaultBucketName))
	err := manager.UploadFile(fileUUID, nil, 0, 1)

	g := gomega.NewGomegaWithT(t)
	g.Expect(err).To(gomega.MatchError(fmt.Sprintf("the file record is not existed, file=%s", fileUUID)))
}

func Test_UploadCompletedFile(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fileUUID := "Test_Manager_UploadCompletedFile1"
	uploadID := "Test_Manager_UploadCompletedFile_UPLOADID"
	serialNum := 1
	etag := "etag"
	content := []byte("Test_UploadCompletedFile")
	alias := "myalias"
	var fileLen int64 = 1

	uf := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   "1",
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
		UploadID:    uploadID,
	}

	err := model.InsertUploadFile(uf)
	g.Expect(err).To(gomega.Succeed())

	err = model.UploadFileCompleted(fileUUID)
	g.Expect(err).To(gomega.Succeed())

	client := &aws2.MockS3Client{
		UploadPartFunc: func(input *s3.UploadPartInput) (*s3.UploadPartOutput, error) {
			g.Expect(*input.Key).To(gomega.Equal(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(content)))
			return &s3.UploadPartOutput{
				ETag: aws.String(etag),
			}, nil
		},
	}

	manager := NewUploadFileManager(aws2.NewAWSStorageWithCli(client, aws2.DefaultBucketName))
	err = manager.UploadFile(fileUUID, bytes.NewReader(content), serialNum, 1)

	g.Expect(err).To(gomega.MatchError(fmt.Sprintf("the file is completed, file=%s", fileUUID)))
}

func Test_UploadFileCompleteFailed(t *testing.T) {
	fileUUID := "Test_Manager_UploadFileCompleteFailed"
	manager := NewUploadFileManager(aws2.NewAWSStorageWithCli(nil, aws2.DefaultBucketName))
	_, err := manager.UploadFileComplete(fileUUID)

	g := gomega.NewGomegaWithT(t)
	g.Expect(err).To(gomega.MatchError(fmt.Sprintf("the file record is not existed, file=%s", fileUUID)))
}

func Test_UploadFileComplete(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	fileUUID := "Test_Manager_UploadFileComplete1"
	uploadID := "Test_Manager_UploadFileComplete_UPLOADID"
	serialCount := 20
	alias := "myalias"
	var fileLen int64 = 1
	startTime := time.Now().Add(-24 * time.Hour)
	endTime := time.Now()

	uf := &db.UploadFile{
		FileUUID:    fileUUID,
		ClusterID:   "1",
		Alias:       alias,
		FileLen:     fileLen,
		StorageType: string(config.AWSStorageType),
		UploadID:    uploadID,
		BeginTime:   &startTime,
		EndTime:     &endTime,
	}
	err := model.InsertUploadFile(uf)
	g.Expect(err).To(gomega.Succeed())

	etags := make(map[int]string)
	for i := 1; i <= serialCount; i++ {
		err = model.UpsertUploadFilePart(fileUUID, i, fmt.Sprintf("etag%d", i), 1)
		g.Expect(err).To(gomega.Succeed())
		etags[i] = fmt.Sprintf("etag%d", i)
	}

	err = model.UploadFileCompleted(fileUUID)
	g.Expect(err).To(gomega.Succeed())

	c := &aws2.MockS3Client{
		CompleteMultipartUploadFunc: func(input *s3.CompleteMultipartUploadInput) (*s3.CompleteMultipartUploadOutput, error) {
			g.Expect(*input.Key).To(gomega.Equal(fmt.Sprintf("%s/%s", aws2.TmpPrefix, 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 k, v := range etags {
				g.Expect(aetags[k]).To(gomega.Equal(v))
			}

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

	manager := NewUploadFileManager(aws2.NewAWSStorageWithCli(c, aws2.DefaultBucketName))
	_, err = manager.UploadFileComplete(fileUUID)
	g.Expect(err).To(gomega.Succeed())
}