package handler

import (
	"crypto/md5"
	"crypto/tls"
	"encoding/base64"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"testing"

	"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api"
	s3Errors "git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api/errors"
	"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api/layer"
	"git.frostfs.info/TrueCloudLab/frostfs-s3-gw/api/layer/encryption"
	"github.com/stretchr/testify/require"
)

const (
	partNumberMarkerQuery = "part-number-marker"
)

func TestMultipartUploadInvalidPart(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName, objName := "bucket-to-upload-part", "object-multipart"
	createTestBucket(hc, bktName)
	partSize := 8 // less than min part size

	multipartUpload := createMultipartUpload(hc, bktName, objName, map[string]string{})
	etag1, _ := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 1, partSize)
	etag2, _ := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 2, partSize)
	w := completeMultipartUploadBase(hc, bktName, objName, multipartUpload.UploadID, []string{etag1, etag2})
	assertS3Error(hc.t, w, s3Errors.GetAPIError(s3Errors.ErrEntityTooSmall))
}

func TestMultipartReUploadPart(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName, objName := "bucket-to-upload-part", "object-multipart"
	bktInfo := createTestBucket(hc, bktName)
	partSizeLast := 8 // less than min part size
	partSizeFirst := 5 * 1024 * 1024

	uploadInfo := createMultipartUpload(hc, bktName, objName, map[string]string{})
	etag1, _ := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 1, partSizeLast)
	etag2, _ := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 2, partSizeFirst)

	list := listParts(hc, bktName, objName, uploadInfo.UploadID, "0", http.StatusOK)
	require.Len(t, list.Parts, 2)
	require.Equal(t, etag1, list.Parts[0].ETag)
	require.Equal(t, etag2, list.Parts[1].ETag)

	w := completeMultipartUploadBase(hc, bktName, objName, uploadInfo.UploadID, []string{etag1, etag2})
	assertS3Error(hc.t, w, s3Errors.GetAPIError(s3Errors.ErrEntityTooSmall))

	etag1, data1 := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 1, partSizeFirst)
	etag2, data2 := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 2, partSizeLast)

	list = listParts(hc, bktName, objName, uploadInfo.UploadID, "0", http.StatusOK)
	require.Len(t, list.Parts, 2)
	require.Equal(t, etag1, list.Parts[0].ETag)
	require.Equal(t, etag2, list.Parts[1].ETag)

	innerUploadInfo, err := hc.tree.GetMultipartUpload(hc.context, bktInfo, objName, uploadInfo.UploadID)
	require.NoError(t, err)
	treeParts, err := hc.tree.GetParts(hc.Context(), bktInfo, innerUploadInfo.ID)
	require.NoError(t, err)
	require.Len(t, treeParts, len(list.Parts))

	w = completeMultipartUploadBase(hc, bktName, objName, uploadInfo.UploadID, []string{etag1, etag2})
	assertStatus(hc.t, w, http.StatusOK)

	data, _ := getObject(hc, bktName, objName)
	equalDataSlices(t, append(data1, data2...), data)
}

func TestListMultipartUploads(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName := "bucket-to-list-uploads"
	createTestBucket(hc, bktName)

	objName1 := "/my/object/name"
	uploadInfo1 := createMultipartUpload(hc, bktName, objName1, map[string]string{})
	objName2 := "/my/object2"
	uploadInfo2 := createMultipartUpload(hc, bktName, objName2, map[string]string{})
	objName3 := "/zzz/object/name3"
	uploadInfo3 := createMultipartUpload(hc, bktName, objName3, map[string]string{})

	t.Run("check upload key", func(t *testing.T) {
		listUploads := listAllMultipartUploads(hc, bktName)
		require.Len(t, listUploads.Uploads, 3)
		for i, upload := range []*InitiateMultipartUploadResponse{uploadInfo1, uploadInfo2, uploadInfo3} {
			require.Equal(t, upload.UploadID, listUploads.Uploads[i].UploadID)
			require.Equal(t, upload.Key, listUploads.Uploads[i].Key)
		}
	})

	t.Run("check max uploads", func(t *testing.T) {
		listUploads := listMultipartUploadsBase(hc, bktName, "", "", "", "", 2)
		require.Len(t, listUploads.Uploads, 2)
		require.Equal(t, uploadInfo1.UploadID, listUploads.Uploads[0].UploadID)
		require.Equal(t, uploadInfo2.UploadID, listUploads.Uploads[1].UploadID)
	})

	t.Run("check prefix", func(t *testing.T) {
		listUploads := listMultipartUploadsBase(hc, bktName, "/my", "", "", "", -1)
		require.Len(t, listUploads.Uploads, 2)
		require.Equal(t, uploadInfo1.UploadID, listUploads.Uploads[0].UploadID)
		require.Equal(t, uploadInfo2.UploadID, listUploads.Uploads[1].UploadID)
	})

	t.Run("check markers", func(t *testing.T) {
		t.Run("check only key-marker", func(t *testing.T) {
			listUploads := listMultipartUploadsBase(hc, bktName, "", "", "", objName2, -1)
			require.Len(t, listUploads.Uploads, 1)
			// If upload-id-marker is not specified, only the keys lexicographically greater than the specified key-marker will be included in the list.
			require.Equal(t, uploadInfo3.UploadID, listUploads.Uploads[0].UploadID)
		})

		t.Run("check only upload-id-marker", func(t *testing.T) {
			uploadIDMarker := uploadInfo1.UploadID
			if uploadIDMarker > uploadInfo2.UploadID {
				uploadIDMarker = uploadInfo2.UploadID
			}
			listUploads := listMultipartUploadsBase(hc, bktName, "", "", uploadIDMarker, "", -1)
			// If key-marker is not specified, the upload-id-marker parameter is ignored.
			require.Len(t, listUploads.Uploads, 3)
		})

		t.Run("check key-marker along with upload-id-marker", func(t *testing.T) {
			uploadIDMarker := "00000000-0000-0000-0000-000000000000"

			listUploads := listMultipartUploadsBase(hc, bktName, "", "", uploadIDMarker, objName3, -1)
			require.Len(t, listUploads.Uploads, 1)
			// If upload-id-marker is specified, any multipart uploads for a key equal to the key-marker might also be included,
			// provided those multipart uploads have upload IDs lexicographically greater than the specified upload-id-marker.
			require.Equal(t, uploadInfo3.UploadID, listUploads.Uploads[0].UploadID)
		})
	})
}

func TestMultipartUploadSize(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName, objName := "bucket-for-test-multipart-size", "object-multipart"
	createTestBucket(hc, bktName)

	partSize := layer.UploadMinSize
	objLen := 2 * partSize
	headers := map[string]string{}

	data := multipartUpload(hc, bktName, objName, headers, objLen, partSize)
	require.Equal(t, objLen, len(data))

	t.Run("check correct size in list v1", func(t *testing.T) {
		listV1 := listObjectsV1(hc, bktName, "", "", "", -1)
		require.Len(t, listV1.Contents, 1)
		require.Equal(t, objLen, int(listV1.Contents[0].Size))
		require.Equal(t, objName, listV1.Contents[0].Key)
	})

	t.Run("check correct size in list v2", func(t *testing.T) {
		listV2 := listObjectsV2(hc, bktName, "", "", "", "", -1)
		require.Len(t, listV2.Contents, 1)
		require.Equal(t, objLen, int(listV2.Contents[0].Size))
		require.Equal(t, objName, listV2.Contents[0].Key)
	})

	t.Run("check correct get", func(t *testing.T) {
		_, hdr := getObject(hc, bktName, objName)
		require.Equal(t, strconv.Itoa(objLen), hdr.Get(api.ContentLength))

		part := getObjectRange(t, hc, bktName, objName, partSize, objLen-1)
		equalDataSlices(t, data[partSize:], part)
	})

	t.Run("check correct size when part copy", func(t *testing.T) {
		objName2 := "obj2"
		uploadInfo := createMultipartUpload(hc, bktName, objName2, headers)
		sourceCopy := bktName + "/" + objName
		uploadPartCopy(hc, bktName, objName2, uploadInfo.UploadID, 1, sourceCopy, 0, 0)
		uploadPartCopy(hc, bktName, objName2, uploadInfo.UploadID, 2, sourceCopy, 0, partSize)
	})

	t.Run("check correct size when copy part from encrypted source", func(t *testing.T) {
		newBucket, newObjName := "new-bucket", "new-object-multipart"
		bktInfo := createTestBucket(hc, newBucket)

		srcObjName := "source-object"
		key := []byte("firstencriptionkeyofsourceobject")
		keyMd5 := md5.Sum(key)
		srcEnc, err := encryption.NewParams(key)
		require.NoError(t, err)
		srcObjInfo := createTestObject(hc, bktInfo, srcObjName, *srcEnc)

		multipartInfo := createMultipartUpload(hc, newBucket, newObjName, headers)

		sourceCopy := newBucket + "/" + srcObjName

		query := make(url.Values)
		query.Set(uploadIDQuery, multipartInfo.UploadID)
		query.Set(partNumberQuery, "1")

		// empty copy-source-sse headers
		w, r := prepareTestRequestWithQuery(hc, newBucket, newObjName, query, nil)
		r.TLS = &tls.ConnectionState{}
		r.Header.Set(api.AmzCopySource, sourceCopy)
		hc.Handler().UploadPartCopy(w, r)

		assertStatus(t, w, http.StatusBadRequest)

		// success copy
		w, r = prepareTestRequestWithQuery(hc, newBucket, newObjName, query, nil)
		r.TLS = &tls.ConnectionState{}
		r.Header.Set(api.AmzCopySource, sourceCopy)
		r.Header.Set(api.AmzCopySourceServerSideEncryptionCustomerAlgorithm, layer.AESEncryptionAlgorithm)
		r.Header.Set(api.AmzCopySourceServerSideEncryptionCustomerKey, base64.StdEncoding.EncodeToString(key))
		r.Header.Set(api.AmzCopySourceServerSideEncryptionCustomerKeyMD5, base64.StdEncoding.EncodeToString(keyMd5[:]))
		hc.Handler().UploadPartCopy(w, r)

		uploadPartCopyResponse := &UploadPartCopyResponse{}
		readResponse(hc.t, w, http.StatusOK, uploadPartCopyResponse)

		completeMultipartUpload(hc, newBucket, newObjName, multipartInfo.UploadID, []string{uploadPartCopyResponse.ETag})
		attr := getObjectAttributes(hc, newBucket, newObjName, objectParts)
		require.Equal(t, 1, attr.ObjectParts.PartsCount)
		require.Equal(t, srcObjInfo.Headers[layer.AttributeDecryptedSize], strconv.Itoa(attr.ObjectParts.Parts[0].Size))
	})
}

func TestListParts(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName, objName := "bucket-for-test-list-parts", "object-multipart"
	_ = createTestBucket(hc, bktName)
	partSize := 5 * 1024 * 1024

	uploadInfo := createMultipartUpload(hc, bktName, objName, map[string]string{})
	etag1, _ := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 1, partSize)
	etag2, _ := uploadPart(hc, bktName, objName, uploadInfo.UploadID, 2, partSize)

	list := listParts(hc, bktName, objName, uploadInfo.UploadID, "0", http.StatusOK)
	require.Len(t, list.Parts, 2)
	require.Equal(t, etag1, list.Parts[0].ETag)
	require.Equal(t, etag2, list.Parts[1].ETag)

	list = listParts(hc, bktName, objName, uploadInfo.UploadID, "1", http.StatusOK)
	require.Len(t, list.Parts, 1)
	require.Equal(t, etag2, list.Parts[0].ETag)

	list = listParts(hc, bktName, objName, uploadInfo.UploadID, "2", http.StatusOK)
	require.Len(t, list.Parts, 0)

	list = listParts(hc, bktName, objName, uploadInfo.UploadID, "7", http.StatusOK)
	require.Len(t, list.Parts, 0)

	list = listParts(hc, bktName, objName, uploadInfo.UploadID, "-1", http.StatusInternalServerError)
	require.Len(t, list.Parts, 0)
}

func TestMultipartUploadWithContentLanguage(t *testing.T) {
	hc := prepareHandlerContext(t)

	bktName, objName := "bucket-1", "object-1"
	createTestBucket(hc, bktName)

	partSize := 5 * 1024 * 1024
	exceptedContentLanguage := "en"
	headers := map[string]string{
		api.ContentLanguage: exceptedContentLanguage,
	}

	multipartUpload := createMultipartUpload(hc, bktName, objName, headers)
	etag1, _ := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 1, partSize)
	etag2, _ := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 2, partSize)
	w := completeMultipartUploadBase(hc, bktName, objName, multipartUpload.UploadID, []string{etag1, etag2})
	assertStatus(t, w, http.StatusOK)

	w, r := prepareTestRequest(hc, bktName, objName, nil)
	hc.Handler().HeadObjectHandler(w, r)
	require.Equal(t, exceptedContentLanguage, w.Header().Get(api.ContentLanguage))
}

func TestMultipartUploadEnabledMD5(t *testing.T) {
	hc := prepareHandlerContext(t)
	hc.config.md5Enabled = true
	hc.layerFeatures.SetMD5Enabled(true)

	bktName, objName := "bucket-md5", "object-md5"
	createTestBucket(hc, bktName)

	partSize := 5 * 1024 * 1024
	multipartUpload := createMultipartUpload(hc, bktName, objName, map[string]string{})
	etag1, partBody1 := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 1, partSize)
	md5Sum1 := md5.Sum(partBody1)
	require.Equal(t, hex.EncodeToString(md5Sum1[:]), etag1)

	etag2, partBody2 := uploadPart(hc, bktName, objName, multipartUpload.UploadID, 2, partSize)
	md5Sum2 := md5.Sum(partBody2)
	require.Equal(t, hex.EncodeToString(md5Sum2[:]), etag2)

	w := completeMultipartUploadBase(hc, bktName, objName, multipartUpload.UploadID, []string{etag1, etag2})
	assertStatus(t, w, http.StatusOK)
	resp := &CompleteMultipartUploadResponse{}
	err := xml.NewDecoder(w.Result().Body).Decode(resp)
	require.NoError(t, err)
	completeMD5Sum := md5.Sum(append(md5Sum1[:], md5Sum2[:]...))
	require.Equal(t, hex.EncodeToString(completeMD5Sum[:])+"-2", resp.ETag)
}

func uploadPartCopy(hc *handlerContext, bktName, objName, uploadID string, num int, srcObj string, start, end int) *UploadPartCopyResponse {
	return uploadPartCopyBase(hc, bktName, objName, false, uploadID, num, srcObj, start, end)
}

func uploadPartCopyBase(hc *handlerContext, bktName, objName string, encrypted bool, uploadID string, num int, srcObj string, start, end int) *UploadPartCopyResponse {
	query := make(url.Values)
	query.Set(uploadIDQuery, uploadID)
	query.Set(partNumberQuery, strconv.Itoa(num))

	w, r := prepareTestRequestWithQuery(hc, bktName, objName, query, nil)
	if encrypted {
		setEncryptHeaders(r)
	}
	r.Header.Set(api.AmzCopySource, srcObj)
	if start+end > 0 {
		r.Header.Set(api.AmzCopySourceRange, fmt.Sprintf("bytes=%d-%d", start, end))
	}

	hc.Handler().UploadPartCopy(w, r)
	uploadPartCopyResponse := &UploadPartCopyResponse{}
	readResponse(hc.t, w, http.StatusOK, uploadPartCopyResponse)

	return uploadPartCopyResponse
}

func listAllMultipartUploads(hc *handlerContext, bktName string) *ListMultipartUploadsResponse {
	return listMultipartUploadsBase(hc, bktName, "", "", "", "", -1)
}

func listMultipartUploadsBase(hc *handlerContext, bktName, prefix, delimiter, uploadIDMarker, keyMarker string, maxUploads int) *ListMultipartUploadsResponse {
	query := make(url.Values)
	query.Set(prefixQueryName, prefix)
	query.Set(delimiterQueryName, delimiter)
	query.Set(uploadIDMarkerQueryName, uploadIDMarker)
	query.Set(keyMarkerQueryName, keyMarker)
	if maxUploads != -1 {
		query.Set(maxUploadsQueryName, strconv.Itoa(maxUploads))
	}

	w, r := prepareTestRequestWithQuery(hc, bktName, "", query, nil)

	hc.Handler().ListMultipartUploadsHandler(w, r)
	listPartsResponse := &ListMultipartUploadsResponse{}
	readResponse(hc.t, w, http.StatusOK, listPartsResponse)

	return listPartsResponse
}

func listParts(hc *handlerContext, bktName, objName string, uploadID, partNumberMarker string, status int) *ListPartsResponse {
	return listPartsBase(hc, bktName, objName, false, uploadID, partNumberMarker, status)
}

func listPartsBase(hc *handlerContext, bktName, objName string, encrypted bool, uploadID, partNumberMarker string, status int) *ListPartsResponse {
	query := make(url.Values)
	query.Set(uploadIDQuery, uploadID)
	query.Set(partNumberMarkerQuery, partNumberMarker)

	w, r := prepareTestRequestWithQuery(hc, bktName, objName, query, nil)
	if encrypted {
		setEncryptHeaders(r)
	}

	hc.Handler().ListPartsHandler(w, r)
	listPartsResponse := &ListPartsResponse{}
	readResponse(hc.t, w, status, listPartsResponse)

	return listPartsResponse
}