file-system/internal/infrastructure/s3/file_repository_impl.go

package s3

import (
	"context"
	"rag/file-system/internal/common"
	"rag/file-system/internal/domain/repository"
	"io"
	"sort"
	"time"

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

const maxContentPreviewSize = 10 * 1024 * 1024 // 10MB

type S3FileRepository struct {
	client *RustFSClient
}

func NewS3FileRepository(client *RustFSClient) repository.FileRepository {
	return &S3FileRepository{client: client}
}

func (r *S3FileRepository) UploadFile(ctx context.Context, bucketName string, objectKey string, data io.Reader) error {
	_, err := r.client.S3Client().PutObject(ctx, &s3.PutObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
		Body:   data,
	})
	return common.WrapS3Error(err)
}

func (r *S3FileRepository) DownloadFile(ctx context.Context, bucketName string, objectKey string) (io.ReadCloser, error) {
	resp, err := r.client.S3Client().GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	})
	if err != nil {
		return nil, common.WrapS3Error(err)
	}
	return resp.Body, nil
}

func (r *S3FileRepository) ListBuckets(ctx context.Context) ([]string, error) {
	resp, err := r.client.S3Client().ListBuckets(ctx, &s3.ListBucketsInput{})
	if err != nil {
		return nil, common.WrapS3Error(err)
	}
	var buckets []string
	for _, b := range resp.Buckets {
		if b.Name != nil {
			buckets = append(buckets, *b.Name)
		}
	}
	return buckets, nil
}

func (r *S3FileRepository) CreateBucket(ctx context.Context, bucketName string) error {
	_, err := r.client.S3Client().CreateBucket(ctx, &s3.CreateBucketInput{
		Bucket: aws.String(bucketName),
	})
	return common.WrapS3Error(err)
}

func (r *S3FileRepository) DeleteBucket(ctx context.Context, bucketName string) error {
	_, err := r.client.S3Client().DeleteBucket(ctx, &s3.DeleteBucketInput{
		Bucket: aws.String(bucketName),
	})
	return common.WrapS3Error(err)
}

// GetFileContent retrieves text file content for preview (e.g., Markdown files)
func (r *S3FileRepository) GetFileContent(ctx context.Context, bucketName string, objectKey string) (string, error) {
	resp, err := r.client.S3Client().GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	})
	if err != nil {
		return "", common.WrapS3Error(err)
	}
	defer resp.Body.Close()

	data, err := io.ReadAll(io.LimitReader(resp.Body, maxContentPreviewSize))
	if err != nil {
		return "", err
	}
	if int64(len(data)) >= maxContentPreviewSize {
		return "", common.NewBusinessException("file too large for content preview (max 10MB)")
	}
	return string(data), nil
}

// DeleteFile removes a file from the bucket
func (r *S3FileRepository) DeleteFile(ctx context.Context, bucketName string, objectKey string) error {
	_, err := r.client.S3Client().DeleteObject(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	})
	return common.WrapS3Error(err)
}

// ListObjectsV2 lists files with pagination support
func (r *S3FileRepository) ListObjectsV2(ctx context.Context, bucketName string, prefix string, maxKeys int32, continuationToken *string) (*repository.ListFilesResult, error) {
	input := &s3.ListObjectsV2Input{
		Bucket:  aws.String(bucketName),
		Prefix:  aws.String(prefix),
		MaxKeys: aws.Int32(maxKeys),
	}
	if continuationToken != nil && *continuationToken != "" {
		input.ContinuationToken = continuationToken
	}

	resp, err := r.client.S3Client().ListObjectsV2(ctx, input)
	if err != nil {
		return nil, common.WrapS3Error(err)
	}

	files := make([]repository.FileInfo, 0, len(resp.Contents))
	for _, obj := range resp.Contents {
		if obj.Key == nil || obj.Size == nil || obj.LastModified == nil || obj.ETag == nil {
			continue
		}
		files = append(files, repository.FileInfo{
			Key:          *obj.Key,
			Size:         *obj.Size,
			LastModified: *obj.LastModified,
			ETag:         *obj.ETag,
		})
	}

	return &repository.ListFilesResult{
		Files:                 files,
		NextContinuationToken: resp.NextContinuationToken,
	}, nil
}

// GeneratePresignedURL generates a presigned URL for temporary file access
func (r *S3FileRepository) GeneratePresignedURL(ctx context.Context, bucketName string, objectKey string, expiry time.Duration) (string, error) {
	presignResult, err := r.client.PresignClient().PresignGetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	}, func(opts *s3.PresignOptions) {
		opts.Expires = expiry
	})
	if err != nil {
		return "", common.WrapS3Error(err)
	}
	return presignResult.URL, nil
}

// CreateMultipartUpload initializes a multipart upload session
func (r *S3FileRepository) CreateMultipartUpload(ctx context.Context, bucketName string, objectKey string) (string, error) {
	resp, err := r.client.S3Client().CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectKey),
	})
	if err != nil {
		return "", common.WrapS3Error(err)
	}
	if resp.UploadId == nil {
		return "", common.NewBusinessException("failed to initialize multipart upload")
	}
	return *resp.UploadId, nil
}

// UploadPart uploads a single part of a multipart upload
func (r *S3FileRepository) UploadPart(ctx context.Context, bucketName string, objectKey string, uploadId string, partNumber int32, data io.Reader) (string, error) {
	resp, err := r.client.S3Client().UploadPart(ctx, &s3.UploadPartInput{
		Bucket:     aws.String(bucketName),
		Key:        aws.String(objectKey),
		UploadId:   aws.String(uploadId),
		PartNumber: aws.Int32(partNumber),
		Body:       data,
	})
	if err != nil {
		return "", common.WrapS3Error(err)
	}
	if resp.ETag == nil {
		return "", common.NewBusinessException("failed to upload part")
	}
	return *resp.ETag, nil
}

// CompleteMultipartUpload assembles all parts to complete the upload
func (r *S3FileRepository) CompleteMultipartUpload(ctx context.Context, bucketName string, objectKey string, uploadId string, parts []common.Part) (string, error) {
	sort.Slice(parts, func(i, j int) bool {
		return parts[i].PartNumber < parts[j].PartNumber
	})

	completedParts := make([]types.CompletedPart, len(parts))
	for i, p := range parts {
		completedParts[i] = types.CompletedPart{
			ETag:       aws.String(p.ETag),
			PartNumber: aws.Int32(p.PartNumber),
		}
	}

	resp, err := r.client.S3Client().CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(bucketName),
		Key:      aws.String(objectKey),
		UploadId: aws.String(uploadId),
		MultipartUpload: &types.CompletedMultipartUpload{
			Parts: completedParts,
		},
	})
	if err != nil {
		return "", common.WrapS3Error(err)
	}
	if resp.Location == nil {
		return "", common.NewBusinessException("failed to complete multipart upload")
	}
	return *resp.Location, nil
}

// AbortMultipartUpload cancels an in-progress multipart upload
func (r *S3FileRepository) AbortMultipartUpload(ctx context.Context, bucketName string, objectKey string, uploadId string) error {
	_, err := r.client.S3Client().AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{
		Bucket:   aws.String(bucketName),
		Key:      aws.String(objectKey),
		UploadId: aws.String(uploadId),
	})
	return common.WrapS3Error(err)
}