backend-celestia-trace/handlers/media_handler.go

package handlers

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/celestia-trace/backend/config"
	"github.com/celestia-trace/backend/models"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

const maxAssetUploadBytes int64 = 2 * 1024 * 1024 * 1024 // 2 GiB

type MediaHandler struct {
	DB  *gorm.DB
	Cfg *config.Config
}

// UploadAsset stores a real session attachment. clientId makes retries
// idempotent; the SHA-256 digest allows clients to verify the stored content.
func (h *MediaHandler) UploadAsset(c *gin.Context) {
	userID := c.GetString("userID")
	sessionID := c.Param("id")
	var session models.CelestiaSession
	if err := h.DB.Where("id = ? AND user_id = ? AND deleted_at IS NULL", sessionID, userID).First(&session).Error; err != nil {
		respondError(c, http.StatusNotFound, 404, "record not found")
		return
	}

	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, maxAssetUploadBytes)
	clientID := strings.TrimSpace(c.PostForm("clientId"))
	kind := strings.ToUpper(strings.TrimSpace(c.PostForm("kind")))
	if clientID == "" || (kind != "AUDIO" && kind != "PHOTO") {
		respondError(c, http.StatusBadRequest, 400, "clientId and a valid kind are required")
		return
	}
	var existing models.MediaAsset
	if err := h.DB.Where("session_id = ? AND client_id = ?", sessionID, clientID).First(&existing).Error; err == nil {
		respondSuccess(c, existing)
		return
	}

	header, err := c.FormFile("file")
	if err != nil {
		respondError(c, http.StatusBadRequest, 400, "file is required")
		return
	}

	// Storage quota check against incoming file size
	var usedBytes int64
	h.DB.Model(&models.MediaAsset{}).Where("user_id = ?", userID).Select("COALESCE(SUM(size_bytes), 0)").Scan(&usedBytes)
	var setting models.SyncSetting
	if err := h.DB.First(&setting, "user_id = ?", userID).Error; err != nil {
		setting.CloudStorageMaxMB = 5120 // default 5 GB
	}
	maxBytes := setting.CloudStorageMaxMB * 1024 * 1024
	remainingBytes := maxBytes - usedBytes
	if remainingBytes < header.Size {
		c.JSON(http.StatusRequestEntityTooLarge, gin.H{
			"code": 413, "message": "storage quota exceeded",
			"data": gin.H{
				"totalBytes":     maxBytes,
				"usedBytes":      usedBytes,
				"remainingBytes": remainingBytes,
				"requiredBytes":  header.Size,
			},
		})
		return
	}
	input, err := header.Open()
	if err != nil {
		respondError(c, http.StatusBadRequest, 400, "failed to read uploaded file")
		return
	}
	defer input.Close()

	asset := models.MediaAsset{
		SessionID:   sessionID,
		UserID:      userID,
		ClientID:    clientID,
		Kind:        kind,
		FileName:    filepath.Base(header.Filename),
		MIMEType:    header.Header.Get("Content-Type"),
		SizeBytes:   0,
		SHA256:      "pending",
		StoragePath: "pending",
	}
	if err := h.DB.Create(&asset).Error; err != nil {
		if isDuplicateError(err) && h.DB.Where("session_id = ? AND client_id = ?", sessionID, clientID).First(&existing).Error == nil {
			respondSuccess(c, existing)
			return
		}
		respondError(c, http.StatusInternalServerError, 500, "failed to create asset")
		return
	}

	directory := filepath.Join(h.Cfg.StoragePath, userID, sessionID)
	if err := os.MkdirAll(directory, 0o750); err != nil {
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to prepare asset storage")
		return
	}
	extension := strings.ToLower(filepath.Ext(header.Filename))
	if len(extension) > 12 {
		extension = ""
	}
	storagePath := filepath.Join(directory, asset.ID+extension)
	output, err := os.OpenFile(storagePath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o640)
	if err != nil {
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to create asset file")
		return
	}
	hasher := sha256.New()
	size, copyErr := io.Copy(io.MultiWriter(output, hasher), input)
	closeErr := output.Close()
	if copyErr != nil || closeErr != nil {
		os.Remove(storagePath)
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to store uploaded file")
		return
	}
	asset.SizeBytes = size
	asset.SHA256 = hex.EncodeToString(hasher.Sum(nil))
	asset.StoragePath = storagePath
	if err := h.DB.Save(&asset).Error; err != nil {
		os.Remove(storagePath)
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to finalize uploaded file")
		return
	}
	// Increment session revision so other devices detect the new asset
	h.DB.Model(&session).Updates(map[string]interface{}{
		"revision":   gorm.Expr("revision + 1"),
		"updated_at": gorm.Expr("NOW()"),
	})
	// Reload to get the updated revision
	h.DB.First(&session, "id = ?", session.ID)
	c.JSON(http.StatusOK, gin.H{
		"code": 0, "message": "success",
		"data": gin.H{
			"asset":           asset,
			"sessionRevision": session.Revision,
		},
	})
}

func (h *MediaHandler) DownloadAsset(c *gin.Context) {
	userID := c.GetString("userID")
	var asset models.MediaAsset
	if err := h.DB.Where("id = ? AND session_id = ? AND user_id = ?", c.Param("assetId"), c.Param("id"), userID).First(&asset).Error; err != nil {
		respondError(c, http.StatusNotFound, 404, "asset not found")
		return
	}
	if _, err := os.Stat(asset.StoragePath); err != nil {
		respondError(c, http.StatusNotFound, 404, "asset file is unavailable")
		return
	}
	if asset.MIMEType != "" {
		c.Header("Content-Type", asset.MIMEType)
	}
	c.FileAttachment(asset.StoragePath, asset.FileName)
}

// InitChunkedUpload creates a new chunked upload session for large files.
func (h *MediaHandler) InitChunkedUpload(c *gin.Context) {
	userID := c.GetString("userID")
	sessionID := c.Param("id")
	var session models.CelestiaSession
	if err := h.DB.Where("id = ? AND user_id = ? AND deleted_at IS NULL", sessionID, userID).First(&session).Error; err != nil {
		respondError(c, http.StatusNotFound, 404, "record not found")
		return
	}

	var req struct {
		ClientID  string `json:"clientId" binding:"required"`
		Kind      string `json:"kind" binding:"required"`
		FileName  string `json:"fileName" binding:"required"`
		MIMEType  string `json:"mimeType"`
		FileSize  int64  `json:"fileSize" binding:"required,gt=0"`
		ChunkSize int    `json:"chunkSize" binding:"required,gt=0"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		respondError(c, http.StatusBadRequest, 400, "invalid request: "+err.Error())
		return
	}
	req.Kind = strings.ToUpper(strings.TrimSpace(req.Kind))
	if req.Kind != "AUDIO" && req.Kind != "PHOTO" {
		respondError(c, http.StatusBadRequest, 400, "kind must be AUDIO or PHOTO")
		return
	}

	// Storage quota pre-check against requested FileSize
	var usedBytes int64
	h.DB.Model(&models.MediaAsset{}).Where("user_id = ?", userID).Select("COALESCE(SUM(size_bytes), 0)").Scan(&usedBytes)
	var setting models.SyncSetting
	if err := h.DB.First(&setting, "user_id = ?", userID).Error; err != nil {
		setting.CloudStorageMaxMB = 5120
	}
	maxBytes := setting.CloudStorageMaxMB * 1024 * 1024
	remainingBytes := maxBytes - usedBytes
	if remainingBytes < req.FileSize {
		c.JSON(http.StatusRequestEntityTooLarge, gin.H{
			"code": 413, "message": "storage quota exceeded",
			"data": gin.H{
				"totalBytes":     maxBytes,
				"usedBytes":      usedBytes,
				"remainingBytes": remainingBytes,
				"requiredBytes":  req.FileSize,
			},
		})
		return
	}

	totalChunks := int(math.Ceil(float64(req.FileSize) / float64(req.ChunkSize)))

	storageDir := filepath.Join(h.Cfg.StoragePath, userID, sessionID, "chunks")
	if err := os.MkdirAll(storageDir, 0o750); err != nil {
		respondError(c, http.StatusInternalServerError, 500, "failed to prepare chunk storage")
		return
	}

	upload := models.ChunkedUpload{
		SessionID:      sessionID,
		UserID:         userID,
		ClientID:       strings.TrimSpace(req.ClientID),
		Kind:           req.Kind,
		FileName:       filepath.Base(req.FileName),
		MIMEType:       req.MIMEType,
		TotalSize:      req.FileSize,
		ChunkSize:      req.ChunkSize,
		TotalChunks:    totalChunks,
		UploadedChunks: 0,
		StorageDir:     storageDir,
		Status:         "pending",
		ExpiresAt:      time.Now().Add(24 * time.Hour),
	}
	if err := h.DB.Create(&upload).Error; err != nil {
		respondError(c, http.StatusInternalServerError, 500, "failed to create chunked upload")
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"code": 0, "message": "success",
		"data": gin.H{
			"uploadId":    upload.ID,
			"totalChunks": totalChunks,
			"chunkSize":   req.ChunkSize,
			"expiresAt":   upload.ExpiresAt,
		},
	})
}

// UploadChunk receives a single chunk of a chunked upload.
func (h *MediaHandler) UploadChunk(c *gin.Context) {
	userID := c.GetString("userID")
	sessionID := c.Param("id")

	uploadID := strings.TrimSpace(c.PostForm("uploadId"))
	chunkIndexStr := strings.TrimSpace(c.PostForm("chunkIndex"))
	if uploadID == "" || chunkIndexStr == "" {
		respondError(c, http.StatusBadRequest, 400, "uploadId and chunkIndex are required")
		return
	}
	chunkIndex, err := strconv.Atoi(chunkIndexStr)
	if err != nil || chunkIndex < 0 {
		respondError(c, http.StatusBadRequest, 400, "chunkIndex must be a non-negative integer")
		return
	}

	var upload models.ChunkedUpload
	if err := h.DB.Where("id = ? AND session_id = ? AND user_id = ? AND status = 'pending'", uploadID, sessionID, userID).First(&upload).Error; err != nil {
		respondError(c, http.StatusNotFound, 404, "chunked upload not found")
		return
	}
	if time.Now().After(upload.ExpiresAt) {
		respondError(c, http.StatusGone, 410, "chunked upload has expired")
		return
	}
	if chunkIndex >= upload.TotalChunks {
		respondError(c, http.StatusBadRequest, 400, "chunkIndex out of range")
		return
	}

	header, err := c.FormFile("file")
	if err != nil {
		respondError(c, http.StatusBadRequest, 400, "file is required")
		return
	}
	input, err := header.Open()
	if err != nil {
		respondError(c, http.StatusBadRequest, 400, "failed to read uploaded chunk")
		return
	}
	defer input.Close()

	chunkPath := filepath.Join(upload.StorageDir, fmt.Sprintf("chunk_%04d", chunkIndex))
	output, err := os.OpenFile(chunkPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o640)
	if err != nil {
		respondError(c, http.StatusInternalServerError, 500, "failed to create chunk file")
		return
	}
	if _, err := io.Copy(output, input); err != nil {
		output.Close()
		os.Remove(chunkPath)
		respondError(c, http.StatusInternalServerError, 500, "failed to write chunk")
		return
	}
	output.Close()

	h.DB.Model(&upload).Update("uploaded_chunks", gorm.Expr("uploaded_chunks + 1"))

	c.JSON(http.StatusOK, gin.H{
		"code": 0, "message": "success",
		"data": gin.H{
			"uploadId":       upload.ID,
			"chunkIndex":     chunkIndex,
			"uploadedChunks": upload.UploadedChunks + 1,
			"totalChunks":    upload.TotalChunks,
		},
	})
}

// CompleteChunkedUpload merges all chunks into a single file and creates the
// MediaAsset record.
func (h *MediaHandler) CompleteChunkedUpload(c *gin.Context) {
	userID := c.GetString("userID")
	sessionID := c.Param("id")

	var req struct {
		UploadID    string `json:"uploadId" binding:"required"`
		TotalChunks int    `json:"totalChunks" binding:"required"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		respondError(c, http.StatusBadRequest, 400, "invalid request: "+err.Error())
		return
	}

	var upload models.ChunkedUpload
	if err := h.DB.Where("id = ? AND session_id = ? AND user_id = ? AND status = 'pending'", req.UploadID, sessionID, userID).First(&upload).Error; err != nil {
		respondError(c, http.StatusNotFound, 404, "chunked upload not found")
		return
	}
	if time.Now().After(upload.ExpiresAt) {
		respondError(c, http.StatusGone, 410, "chunked upload has expired")
		return
	}
	if req.TotalChunks != upload.TotalChunks {
		respondError(c, http.StatusBadRequest, 400, "totalChunks mismatch")
		return
	}

	// Verify all chunks exist
	for i := 0; i < upload.TotalChunks; i++ {
		chunkPath := filepath.Join(upload.StorageDir, fmt.Sprintf("chunk_%04d", i))
		if _, err := os.Stat(chunkPath); err != nil {
			respondError(c, http.StatusBadRequest, 400, fmt.Sprintf("chunk %d is missing", i))
			return
		}
	}

	// Merge chunks into final file
	directory := filepath.Join(h.Cfg.StoragePath, userID, sessionID)
	if err := os.MkdirAll(directory, 0o750); err != nil {
		respondError(c, http.StatusInternalServerError, 500, "failed to prepare asset storage")
		return
	}
	extension := strings.ToLower(filepath.Ext(upload.FileName))
	if len(extension) > 12 {
		extension = ""
	}

	// Create MediaAsset first to get the ID for the filename
	asset := models.MediaAsset{
		SessionID:   sessionID,
		UserID:      userID,
		ClientID:    upload.ClientID,
		Kind:        upload.Kind,
		FileName:    upload.FileName,
		MIMEType:    upload.MIMEType,
		SizeBytes:   0,
		SHA256:      "pending",
		StoragePath: "pending",
	}
	if err := h.DB.Create(&asset).Error; err != nil {
		if isDuplicateError(err) {
			var existing models.MediaAsset
			if h.DB.Where("session_id = ? AND client_id = ?", sessionID, upload.ClientID).First(&existing).Error == nil {
				respondSuccess(c, existing)
				return
			}
		}
		respondError(c, http.StatusInternalServerError, 500, "failed to create asset")
		return
	}

	storagePath := filepath.Join(directory, asset.ID+extension)
	output, err := os.OpenFile(storagePath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o640)
	if err != nil {
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to create asset file")
		return
	}

	hasher := sha256.New()
	var totalSize int64
	for i := 0; i < upload.TotalChunks; i++ {
		chunkPath := filepath.Join(upload.StorageDir, fmt.Sprintf("chunk_%04d", i))
		chunkFile, err := os.Open(chunkPath)
		if err != nil {
			output.Close()
			os.Remove(storagePath)
			h.DB.Delete(&asset)
			respondError(c, http.StatusInternalServerError, 500, "failed to read chunk")
			return
		}
		n, copyErr := io.Copy(io.MultiWriter(output, hasher), chunkFile)
		chunkFile.Close()
		if copyErr != nil {
			output.Close()
			os.Remove(storagePath)
			h.DB.Delete(&asset)
			respondError(c, http.StatusInternalServerError, 500, "failed to merge chunk")
			return
		}
		totalSize += n
	}
	if err := output.Close(); err != nil {
		os.Remove(storagePath)
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to finalize merged file")
		return
	}

	asset.SizeBytes = totalSize
	asset.SHA256 = hex.EncodeToString(hasher.Sum(nil))
	asset.StoragePath = storagePath
	if err := h.DB.Save(&asset).Error; err != nil {
		os.Remove(storagePath)
		h.DB.Delete(&asset)
		respondError(c, http.StatusInternalServerError, 500, "failed to finalize asset")
		return
	}

	// Clean up chunk files
	for i := 0; i < upload.TotalChunks; i++ {
		os.Remove(filepath.Join(upload.StorageDir, fmt.Sprintf("chunk_%04d", i)))
	}
	os.Remove(upload.StorageDir)

	// Mark upload as completed
	h.DB.Model(&upload).Update("status", "completed")

	// Increment session revision
	var session models.CelestiaSession
	h.DB.First(&session, "id = ?", sessionID)
	h.DB.Model(&session).Updates(map[string]interface{}{
		"revision":   gorm.Expr("revision + 1"),
		"updated_at": gorm.Expr("NOW()"),
	})
	h.DB.First(&session, "id = ?", session.ID)

	c.JSON(http.StatusOK, gin.H{
		"code": 0, "message": "success",
		"data": gin.H{
			"asset":           asset,
			"sessionRevision": session.Revision,
		},
	})
}