at-container-registry/pkg/hold/oci/multipart.go

package oci

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"log/slog"
	"sort"
	"strings"
	"sync"
	"time"

	"atcr.io/pkg/atproto"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/google/uuid"
)

// MultipartMode indicates how multipart uploads are handled
type MultipartMode int

const (
	// S3Native uses S3's native multipart API with presigned URLs
	S3Native MultipartMode = iota
	// Buffered buffers parts in memory and assembles them in the hold service
	Buffered
)

// PartInfo represents an uploaded part with its ETag
type PartInfo struct {
	PartNumber int    `json:"part_number"`
	ETag       string `json:"etag"`
}

// PartUploadInfo contains structured information for uploading a part
// Used for both S3 presigned URLs and buffered mode with headers
type PartUploadInfo struct {
	URL     string            `json:"url"`               // URL to PUT the part to
	Method  string            `json:"method,omitempty"`  // HTTP method (usually "PUT")
	Headers map[string]string `json:"headers,omitempty"` // Additional headers required for the request
}

// MultipartSession tracks an in-progress multipart upload
type MultipartSession struct {
	UploadID     string                 // Unique upload ID
	Digest       string                 // Target digest path
	Mode         MultipartMode          // Upload mode (S3Native or Buffered)
	S3UploadID   string                 // S3 upload ID (for S3Native mode)
	Parts        map[int]*MultipartPart // Buffered parts (for Buffered mode)
	CreatedAt    time.Time              // When upload started
	LastActivity time.Time              // Last part upload
	mu           sync.RWMutex           // Protects Parts map
}

// MultipartPart represents a single part in a multipart upload
type MultipartPart struct {
	PartNumber int       // Part number (1-indexed)
	Data       []byte    // Part data (for Buffered mode)
	ETag       string    // ETag from S3 or computed hash
	Size       int64     // Part size in bytes
	UploadedAt time.Time // When part was uploaded
}

// MultipartManager manages multipart upload sessions
type MultipartManager struct {
	sessions map[string]*MultipartSession // uploadID -> session
	mu       sync.RWMutex                 // Protects sessions map
}

// NewMultipartManager creates a new multipart manager
func NewMultipartManager() *MultipartManager {
	mgr := &MultipartManager{
		sessions: make(map[string]*MultipartSession),
	}

	// Start cleanup goroutine for abandoned uploads
	go mgr.cleanupLoop()

	return mgr
}

// cleanupLoop periodically cleans up expired sessions
func (m *MultipartManager) cleanupLoop() {
	ticker := time.NewTicker(15 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		m.cleanupExpiredSessions()
	}
}

// cleanupExpiredSessions removes sessions inactive for >24 hours
func (m *MultipartManager) cleanupExpiredSessions() {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	for uploadID, session := range m.sessions {
		if now.Sub(session.LastActivity) > 24*time.Hour {
			slog.Debug("Cleaning up expired multipart session",
				"uploadID", uploadID,
				"age", now.Sub(session.CreatedAt))
			delete(m.sessions, uploadID)
		}
	}
}

// CreateSession creates a new multipart upload session
func (m *MultipartManager) CreateSession(digest string, mode MultipartMode, s3UploadID string) *MultipartSession {
	uploadID := uuid.New().String()

	session := &MultipartSession{
		UploadID:     uploadID,
		Digest:       digest,
		Mode:         mode,
		S3UploadID:   s3UploadID,
		Parts:        make(map[int]*MultipartPart),
		CreatedAt:    time.Now(),
		LastActivity: time.Now(),
	}

	m.mu.Lock()
	m.sessions[uploadID] = session
	m.mu.Unlock()

	slog.Debug("Created multipart session",
		"uploadID", uploadID,
		"digest", digest,
		"mode", mode)
	return session
}

// GetSession retrieves a multipart session by upload ID
func (m *MultipartManager) GetSession(uploadID string) (*MultipartSession, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()

	session, ok := m.sessions[uploadID]
	if !ok {
		return nil, fmt.Errorf("multipart session not found: %s", uploadID)
	}

	return session, nil
}

// DeleteSession removes a multipart session
func (m *MultipartManager) DeleteSession(uploadID string) {
	m.mu.Lock()
	defer m.mu.Unlock()

	delete(m.sessions, uploadID)
	slog.Debug("Deleted multipart session", "uploadID", uploadID)
}

// StorePart stores a part in the session (for Buffered mode)
func (s *MultipartSession) StorePart(partNumber int, data []byte) string {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Compute ETag as SHA256 hash of part data
	hash := sha256.Sum256(data)
	etag := hex.EncodeToString(hash[:])

	part := &MultipartPart{
		PartNumber: partNumber,
		Data:       data,
		ETag:       etag,
		Size:       int64(len(data)),
		UploadedAt: time.Now(),
	}

	s.Parts[partNumber] = part
	s.LastActivity = time.Now()

	slog.Debug("Stored part",
		"uploadID", s.UploadID,
		"part", partNumber,
		"size", len(data),
		"etag", etag)
	return etag
}

// AssembleBufferedParts assembles all buffered parts into a single blob
// Returns the complete data and total size
func (s *MultipartSession) AssembleBufferedParts() ([]byte, int64, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if s.Mode != Buffered {
		return nil, 0, fmt.Errorf("session is not in buffered mode")
	}

	// Calculate total size
	var totalSize int64
	maxPart := 0
	for partNum, part := range s.Parts {
		totalSize += part.Size
		if partNum > maxPart {
			maxPart = partNum
		}
	}

	// Check for missing parts
	for i := 1; i <= maxPart; i++ {
		if _, ok := s.Parts[i]; !ok {
			return nil, 0, fmt.Errorf("missing part %d", i)
		}
	}

	// Assemble parts in order
	assembled := make([]byte, 0, totalSize)
	for i := 1; i <= maxPart; i++ {
		part := s.Parts[i]
		assembled = append(assembled, part.Data...)
	}

	slog.Debug("Assembled buffered parts",
		"uploadID", s.UploadID,
		"parts", maxPart,
		"totalSize", totalSize)
	return assembled, totalSize, nil
}

// StartMultipartUploadWithManager initiates a multipart upload using the manager
// Returns uploadID and mode
func (h *XRPCHandler) StartMultipartUploadWithManager(ctx context.Context, digest string) (string, MultipartMode, error) {
	// Check if presigned URLs are disabled for testing
	if h.disablePresignedURLs {
		slog.Debug("Presigned URLs disabled, using buffered mode", "reason", "DISABLE_PRESIGNED_URLS=true")
		session := h.MultipartMgr.CreateSession(digest, Buffered, "")
		slog.Debug("Started buffered multipart", "uploadID", session.UploadID)
		return session.UploadID, Buffered, nil
	}

	// Try S3 native multipart first
	if h.s3Service.Client != nil {
		if h.s3Service.Client == nil {
			return "", S3Native, fmt.Errorf("S3 not configured")
		}
		path := blobPath(digest)
		s3Key := strings.TrimPrefix(path, "/")
		if h.s3Service.PathPrefix != "" {
			s3Key = h.s3Service.PathPrefix + "/" + s3Key
		}

		result, err := h.s3Service.Client.CreateMultipartUploadWithContext(ctx, &s3.CreateMultipartUploadInput{
			Bucket: &h.s3Service.Bucket,
			Key:    &s3Key,
		})
		if err == nil {
			s3UploadID := *result.UploadId
			// S3 native multipart succeeded
			session := h.MultipartMgr.CreateSession(digest, S3Native, s3UploadID)
			slog.Debug("Started S3 native multipart",
				"digest", digest,
				"uploadID", session.UploadID,
				"s3UploadID", s3UploadID)
			return session.UploadID, S3Native, nil
		}
		slog.Warn("S3 native multipart failed, falling back to buffered mode", "error", err)
	}

	// Fallback to buffered mode
	session := h.MultipartMgr.CreateSession(digest, Buffered, "")
	slog.Debug("Started buffered multipart", "uploadID", session.UploadID)
	return session.UploadID, Buffered, nil
}

// GetPartUploadURL generates a presigned URL for uploading a part
// Only used for S3Native mode - Buffered mode is handled by blobstore adapter
func (h *XRPCHandler) GetPartUploadURL(ctx context.Context, uploadID string, partNumber int) (*PartUploadInfo, error) {
	session, err := h.MultipartMgr.GetSession(uploadID)
	if err != nil {
		return nil, err
	}

	// For S3Native mode: return presigned URL
	if session.Mode == S3Native {
		if h.s3Service.Client == nil {
			return nil, fmt.Errorf("S3 not configured")
		}

		path := blobPath(session.Digest)
		s3Key := strings.TrimPrefix(path, "/")
		if h.s3Service.PathPrefix != "" {
			s3Key = h.s3Service.PathPrefix + "/" + s3Key
		}
		pnum := int64(partNumber)
		req, _ := h.s3Service.Client.UploadPartRequest(&s3.UploadPartInput{
			Bucket:     &h.s3Service.Bucket,
			Key:        &s3Key,
			UploadId:   &session.S3UploadID,
			PartNumber: &pnum,
		})

		url, err := req.Presign(15 * time.Minute)
		if err != nil {
			return nil, err
		}

		slog.Debug("Generated part presigned URL",
			"digest", session.Digest,
			"uploadID", uploadID,
			"part", partNumber)

		return &PartUploadInfo{
			URL:    url,
			Method: "PUT",
		}, nil
	}

	// Buffered mode: return XRPC endpoint with headers
	return &PartUploadInfo{
		URL:    fmt.Sprintf("%s%s", h.pds.PublicURL, atproto.HoldUploadPart),
		Method: "PUT",
		Headers: map[string]string{
			"X-Upload-Id":   uploadID,
			"X-Part-Number": fmt.Sprintf("%d", partNumber),
		},
	}, nil
}

// CompleteMultipartUploadWithManager completes a multipart upload and moves to final location
// finalDigest is the real digest (e.g., "sha256:abc123...") for the final storage location
// session.Digest is the temp location (e.g., "uploads/temp-<uuid>")
func (h *XRPCHandler) CompleteMultipartUploadWithManager(ctx context.Context, uploadID string, finalDigest string, parts []PartInfo) error {
	session, err := h.MultipartMgr.GetSession(uploadID)
	defer h.MultipartMgr.DeleteSession(uploadID)
	if err != nil {
		return err
	}

	if session.Mode == S3Native {
		if h.s3Service.Client == nil {
			return fmt.Errorf("S3 not configured")
		}

		// Sort parts by part number (S3 requires ascending order)
		sort.Slice(parts, func(i, j int) bool {
			return parts[i].PartNumber < parts[j].PartNumber
		})

		// Convert to S3 CompletedPart format
		// IMPORTANT: S3 requires ETags to be quoted in the CompleteMultipartUpload XML
		s3Parts := make([]*s3.CompletedPart, len(parts))
		for i, p := range parts {
			etag := normalizeETag(p.ETag)
			pnum := int64(p.PartNumber)
			s3Parts[i] = &s3.CompletedPart{
				PartNumber: &pnum,
				ETag:       &etag,
			}
		}
		sourcePath := blobPath(session.Digest)
		s3Key := strings.TrimPrefix(sourcePath, "/")
		if h.s3Service.PathPrefix != "" {
			s3Key = h.s3Service.PathPrefix + "/" + s3Key
		}

		_, err = h.s3Service.Client.CompleteMultipartUploadWithContext(ctx, &s3.CompleteMultipartUploadInput{
			Bucket:   &h.s3Service.Bucket,
			Key:      &s3Key,
			UploadId: &session.S3UploadID,
			MultipartUpload: &s3.CompletedMultipartUpload{
				Parts: s3Parts,
			},
		})
		if err != nil {
			return fmt.Errorf("failed to complete multipart upload: digest=%s, uploadID=%s, err=%v", session.Digest, uploadID, err)
		}
		slog.Info("Completed S3 native multipart at temp location",
			"digest", session.Digest,
			"uploadID", session.UploadID,
			"parts", len(s3Parts))

		// Verify the blob exists at temp location before moving
		destPath := blobPath(finalDigest)
		slog.Debug("About to move blob",
			"source", sourcePath,
			"dest", destPath)

		if _, err := h.driver.Stat(ctx, sourcePath); err != nil {
			slog.Error("Source blob not found after multipart complete",
				"path", sourcePath,
				"error", err)
			return fmt.Errorf("source blob not found after multipart complete: %w", err)
		}
		slog.Debug("Source blob verified", "path", sourcePath)

		// Move from temp to final digest location using driver
		// Driver handles path management correctly (including S3 prefix)
		if err := h.driver.Move(ctx, sourcePath, destPath); err != nil {
			slog.Error("Failed to move blob",
				"source", sourcePath,
				"dest", destPath,
				"error", err)
			return fmt.Errorf("failed to move blob to final location: %w", err)
		}

		slog.Info("Moved blob to final location",
			"from", session.Digest,
			"to", finalDigest,
			"sourcePath", sourcePath,
			"destPath", destPath)
		return nil
	}

	// Buffered mode: assemble parts and write directly to final location
	data, size, err := session.AssembleBufferedParts()
	if err != nil {
		return fmt.Errorf("failed to assemble parts: %w", err)
	}

	// Write assembled blob to final digest location (not temp)
	path := blobPath(finalDigest)
	writer, err := h.driver.Writer(ctx, path, false)
	if err != nil {
		return fmt.Errorf("failed to create writer: %w", err)
	}

	written, err := writer.Write(data)
	if err != nil {
		writer.Cancel(ctx)
		return fmt.Errorf("failed to write blob: %w", err)
	}

	if err := writer.Commit(ctx); err != nil {
		return fmt.Errorf("failed to commit blob: %w", err)
	}

	slog.Info("Completed buffered multipart",
		"uploadID", session.UploadID,
		"finalDigest", finalDigest,
		"size", size,
		"written", written)
	return nil
}

// AbortMultipartUploadWithManager aborts a multipart upload
func (h *XRPCHandler) AbortMultipartUploadWithManager(ctx context.Context, uploadID string) error {
	session, err := h.MultipartMgr.GetSession(uploadID)
	defer h.MultipartMgr.DeleteSession(uploadID)
	if err != nil {
		return err
	}

	if session.Mode == S3Native {
		if h.s3Service.Client == nil {
			return fmt.Errorf("S3 not configured")
		}
		path := blobPath(session.Digest)
		s3Key := strings.TrimPrefix(path, "/")
		if h.s3Service.PathPrefix != "" {
			s3Key = h.s3Service.PathPrefix + "/" + s3Key
		}

		_, err := h.s3Service.Client.AbortMultipartUploadWithContext(ctx, &s3.AbortMultipartUploadInput{
			Bucket:   &h.s3Service.Bucket,
			Key:      &s3Key,
			UploadId: &session.S3UploadID,
		})
		// Abort S3 multipart upload
		if err != nil {
			return fmt.Errorf("failed  to abort multipart upload: digest=%s, uploadID=%s, err=%v", session.Digest, uploadID, err)
		}
		slog.Debug("Aborted S3 native multipart",
			"digest", session.Digest,
			"uploadID", session.UploadID)
		return nil
	}

	// Buffered mode: just delete the session (parts are in memory)
	slog.Debug("Aborted buffered multipart", "uploadID", session.UploadID)
	return nil
}

// HandleBufferedPartUpload handles uploading a part in buffered mode
func (h *XRPCHandler) HandleBufferedPartUpload(ctx context.Context, uploadID string, partNumber int, data []byte) (string, error) {
	session, err := h.MultipartMgr.GetSession(uploadID)
	if err != nil {
		return "", err
	}

	if session.Mode != Buffered {
		return "", fmt.Errorf("session is not in buffered mode")
	}

	etag := session.StorePart(partNumber, data)
	return etag, nil
}

// normalizeETag ensures an ETag has quotes (required by S3 CompleteMultipartUpload)
// S3 returns ETags with quotes, but HTTP clients may strip them
func normalizeETag(etag string) string {
	// Already has quotes
	if strings.HasPrefix(etag, "\"") && strings.HasSuffix(etag, "\"") {
		return etag
	}
	// Add quotes
	return fmt.Sprintf("\"%s\"", etag)
}

// blobPath converts a digest (e.g., "sha256:abc123...") or temp path to a storage path
// Distribution stores blobs as: /docker/registry/v2/blobs/{algorithm}/{xx}/{hash}/data
// where xx is the first 2 characters of the hash for directory sharding
// NOTE: Path must start with / for filesystem driver
// This is used for OCI container layers (content-addressed, globally deduplicated)
func blobPath(digest string) string {
	// Handle temp paths (start with uploads/temp-)
	if strings.HasPrefix(digest, "uploads/temp-") {
		return fmt.Sprintf("/docker/registry/v2/%s/data", digest)
	}

	// Split digest into algorithm and hash
	parts := strings.SplitN(digest, ":", 2)
	if len(parts) != 2 {
		// Fallback for malformed digest
		return fmt.Sprintf("/docker/registry/v2/blobs/%s/data", digest)
	}

	algorithm := parts[0]
	hash := parts[1]

	// Use first 2 characters for sharding
	if len(hash) < 2 {
		return fmt.Sprintf("/docker/registry/v2/blobs/%s/%s/data", algorithm, hash)
	}

	return fmt.Sprintf("/docker/registry/v2/blobs/%s/%s/%s/data", algorithm, hash[:2], hash)
}