package main

import (
	"log"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

func listAndDeleteObjects(svc *s3.S3, bucketName, folderName string, workerCount int) error {
	if folderName == "" {
		// IMP: if folder name is empty, then we are deleting the entire bucket
		return fmt.Errorf("folder name is required")
	}

	var wg sync.WaitGroup
	objChan := make(chan []*s3.ObjectIdentifier, workerCount)

	// Worker function to process object deletion
	deleteWorker := func() {
		defer wg.Done()
		for objects := range objChan {
			_, err := svc.DeleteObjects(&s3.DeleteObjectsInput{
				Bucket: aws.String(bucketName),
				Delete: &s3.Delete{
					Objects: objects,
					Quiet:   aws.Bool(true),
				},
			})
			if err != nil {
				log.Printf("Failed to delete objects: %v", err)
			}
		}
	}

	// Start worker goroutines
	for i := 0; i < workerCount; i++ {
		wg.Add(1)
		go deleteWorker()
	}

	// List and enqueue objects for deletion
	params := &s3.ListObjectsV2Input{
		Bucket: aws.String(bucketName),
		Prefix: aws.String(folderName),
	}
	for {
		resp, err := svc.ListObjectsV2(params)
		if err != nil {
			return err
		}

		var objectsToDelete []*s3.ObjectIdentifier
		for _, item := range resp.Contents {
			objectsToDelete = append(objectsToDelete, &s3.ObjectIdentifier{Key: item.Key})
		}

		// Send objects to deletion workers in batches
		if len(objectsToDelete) > 0 {
			objChan <- objectsToDelete
		}

		// Break if no more objects
		if !aws.BoolValue(resp.IsTruncated) {
			break
		}
		params.ContinuationToken = resp.NextContinuationToken
	}

	close(objChan)
	wg.Wait()

	return nil
}

func main() {
	// Initialize S3 session
	sess, err := session.NewSession(&aws.Config{
		Region: aws.String("us-east-1"), // Replace with your region
	})
	if err != nil {
		log.Fatalf("Failed to create session: %v", err)
	}

	svc := s3.New(sess)

	// Specify bucket and folder name
	bucketName := "your-bucket-name"
	folderName := "your-folder-name/" // Ensure the folder name ends with "/"

	// Delete the folder with concurrency
	workerCount := 5 // Adjust number of concurrent workers as needed
	err = listAndDeleteObjects(svc, bucketName, folderName, workerCount)
	if err != nil {
		log.Fatalf("Failed to delete folder: %v", err)
	}

	log.Println("Folder deleted successfully")
}