package test

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"sort"
	"strings"
	"sync"
	"testing"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
)

const (
	accessKeyID     = ""
	secretAccessKey = ""
	sessionToken    = "="
	region          = "us-west-2"
	bucketName      = "export-aurora-test"
	keyName         = "uploaded/bigfile.tar.gz" // uplaod
	// filePath        = "./main.go"
	filePath    = "/data/bigfile.tar.gz"
	partSize    = 100 * 1024 * 1024 // 50MB
	concurrency = 5                 // Number of concurrent upload threads
)

func TestUplaod(t *testing.T) {
	ctx := context.Background()

	// AWS temporary credential configuration
	cfg, err := config.LoadDefaultConfig(ctx,
		config.WithRegion(region),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
			accessKeyID, secretAccessKey, sessionToken,
		)),
	)
	if err != nil {
		log.Fatalf("Failed to load AWS config: %v", err)
	}

	client := s3.NewFromConfig(cfg)

	// Open local large file
	file, err := os.Open(filePath)
	if err != nil {
		log.Fatalf("Failed to open file: %v", err)
	}
	defer file.Close()

	// Initialize Multipart upload
	createResp, err := client.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(keyName),
	})
	if err != nil {
		log.Fatalf("CreateMultipartUpload failed: %v", err)
	}
	uploadID := createResp.UploadId

	var (
		parts      []types.CompletedPart
		mutex      sync.Mutex
		wg         sync.WaitGroup
		partNumber int32 = 1
	)

	// Use buffered channel to control concurrency
	sem := make(chan struct{}, concurrency)
	buffer := make([]byte, partSize)

	for {
		// Read the next partSize chunk of the file
		n, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			log.Fatalf("Failed to read file: %v", err)
		}
		if n == 0 {
			break
		}

		partBuffer := make([]byte, n)
		copy(partBuffer, buffer[:n])

		sem <- struct{}{} // acquire semaphore
		wg.Add(1)
		currentPartNumber := partNumber

		go func(data []byte, pn int32) {
			defer wg.Done()
			defer func() { <-sem }() // release semaphore

			uploadResp, err := client.UploadPart(ctx, &s3.UploadPartInput{
				Bucket:     aws.String(bucketName),
				Key:        aws.String(keyName),
				UploadId:   uploadID,
				PartNumber: aws.Int32(pn),
				Body:       bytes.NewReader(data),
			})
			if err != nil {
				log.Fatalf("UploadPart %d failed: %v", pn, err)
			}

			mutex.Lock()
			parts = append(parts, types.CompletedPart{
				ETag:       uploadResp.ETag,
				PartNumber: aws.Int32(pn),
			})
			mutex.Unlock()

			fmt.Printf("Uploaded part %d (%d bytes)\n", pn, len(data))
		}(partBuffer, currentPartNumber)

		partNumber++
	}

	wg.Wait()

	// ✅ sort parts by PartNumber
	sort.Slice(parts, func(i, j int) bool {
		return *parts[i].PartNumber < *parts[j].PartNumber
	})

	// Sort by PartNumber
	// AWS requires each part to be submitted in order
	// For strictness, you can use sort.Slice(parts, ...) to sort by PartNumber

	// Complete Multipart upload
	_, err = client.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(bucketName),
		Key:      aws.String(keyName),
		UploadId: uploadID,
		MultipartUpload: &types.CompletedMultipartUpload{
			Parts: parts,
		},
	})
	if err != nil {
		log.Fatalf("CompleteMultipartUpload failed: %v", err)
	}

	fmt.Printf("✅ Successfully uploaded %s to s3://%s/%s using multipart upload\n", filePath, bucketName, keyName)
}

func TestDownload(t *testing.T) {
	localPath := "/Users/wlt/Downloads/temp/bigfile.tar.gz" // Change to your desired save location
	err := DownloadFromS3(localPath)
	if err != nil {
		t.Fatalf("Download failed: %v", err)
	}
}

// Concurrently download large files
func TestDownloadConcurrently(t *testing.T) {
	DownloadFromS3Concurrently()
}

func DownloadFromS3(localPath string) error {
	ctx := context.Background()

	// AWS config (use the same temporary credentials as upload)
	cfg, err := config.LoadDefaultConfig(ctx,
		config.WithRegion(region),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
			accessKeyID, secretAccessKey, sessionToken,
		)),
	)
	if err != nil {
		return fmt.Errorf("failed to load AWS config: %w", err)
	}

	client := s3.NewFromConfig(cfg)

	// Create local file for writing
	outFile, err := os.Create(localPath)
	if err != nil {
		return fmt.Errorf("failed to create local file: %w", err)
	}
	defer outFile.Close()

	// Initiate GetObject request (supports large file streaming)
	resp, err := client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(keyName),
	})
	if err != nil {
		return fmt.Errorf("GetObject failed: %w", err)
	}
	defer resp.Body.Close()

	// Write S3 response body directly to local file
	_, err = io.Copy(outFile, resp.Body)
	if err != nil {
		return fmt.Errorf("failed to copy data to local file: %w", err)
	}

	fmt.Printf("✅ Downloaded file from s3://%s/%s to %s\n", bucketName, keyName, localPath)
	return nil
}

// Concurrent download
func DownloadFromS3Concurrently() {
	ctx := context.Background()

	cfg, err := config.LoadDefaultConfig(ctx,
		config.WithRegion(region),
		config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider(
			accessKeyID, secretAccessKey, sessionToken,
		)),
	)
	if err != nil {
		log.Fatalf("Failed to load AWS config: %v", err)
	}

	client := s3.NewFromConfig(cfg)

	// Get object metadata, get total file size
	headResp, err := client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(keyName),
	})
	if err != nil {
		log.Fatalf("Failed to get object metadata: %v", err)
	}
	fileSize := *headResp.ContentLength
	fmt.Printf("Downloading %d bytes...\n", fileSize)

	// Create local file
	outFile, err := os.Create("downloaded_" + keyName[strings.LastIndex(keyName, "/")+1:])
	if err != nil {
		log.Fatalf("Failed to create local file: %v", err)
	}
	defer outFile.Close()

	var wg sync.WaitGroup
	sem := make(chan struct{}, concurrency) // Control concurrency
	partCount := int(fileSize / int64(partSize))
	if fileSize%int64(partSize) != 0 {
		partCount++
	}

	for i := 0; i < partCount; i++ {
		wg.Add(1)
		sem <- struct{}{}

		start := int64(i) * int64(partSize)
		end := start + int64(partSize) - 1
		if end >= fileSize {
			end = fileSize - 1
		}

		go func(start, end int64, index int) {
			defer wg.Done()
			defer func() { <-sem }()

			rangeHeader := fmt.Sprintf("bytes=%d-%d", start, end)
			getResp, err := client.GetObject(ctx, &s3.GetObjectInput{
				Bucket: aws.String(bucketName),
				Key:    aws.String(keyName),
				Range:  aws.String(rangeHeader),
			})
			if err != nil {
				log.Fatalf("Failed to get object part %d: %v", index, err)
			}
			defer getResp.Body.Close()

			// Write to specified position in file
			_, err = outFile.WriteAt(readAll(getResp.Body), start)
			if err != nil {
				log.Fatalf("Failed to write part %d: %v", index, err)
			}
			fmt.Printf("Downloaded part %d (%s)\n", index+1, rangeHeader)
		}(start, end, i)
	}

	wg.Wait()
	fmt.Println("✅ Download complete.")
}

func readAll(r io.Reader) []byte {
	buf := new(bytes.Buffer)
	_, _ = io.Copy(buf, r)
	return buf.Bytes()
}