pipely/internal/download/handler.go

package download

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"

	"onixbyte.com/message-converter-manifest/internal/database"
	"onixbyte.com/message-converter-manifest/internal/middleware"
)

// Handler manages artefact downloads with bandwidth throttling and Range support.
type Handler struct {
	db          *gorm.DB
	artefactDir string
	semaphore   *middleware.GatedSemaphore
	bytesPerSec int
}

func NewHandler(db *gorm.DB, artefactDir string, sem *middleware.GatedSemaphore, bytesPerSec int) *Handler {
	return &Handler{
		db:          db,
		artefactDir: artefactDir,
		semaphore:   sem,
		bytesPerSec: bytesPerSec,
	}
}

// HandleDownload serves artefact files with resumable download support.
func (h *Handler) HandleDownload(c *gin.Context) {
	artefactID, err := strconv.ParseInt(c.Param("id"), 10, 64)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "Invalid artefact id"})
		return
	}

	deviceID := c.Query("deviceId")
	deploymentID, _ := strconv.ParseInt(c.Query("deploymentId"), 10, 64)

	artefact, err := database.GetArtefactByID(h.db, artefactID)
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			c.AbortWithStatusJSON(http.StatusNotFound, gin.H{"error": "Artefact not found"})
			return
		}
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	file, err := os.Open(artefact.FilePath)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{"error": "Package file not found on disk"})
		return
	}
	defer file.Close()

	fileInfo, err := file.Stat()
	if err != nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": "Failed to stat file"})
		return
	}

	// Update device status to downloading
	if deviceID != "" && deploymentID > 0 {
		if err := database.UpdateDeploymentDeviceStatus(h.db, deploymentID, deviceID, "downloading", ""); err != nil {
			log.Printf("[download] Failed to update device status: %v", err)
		}
	}

	// Apply per-connection bandwidth throttling
	throttledWriter := &ThrottledResponseWriter{
		ResponseWriter: c.Writer,
		bytesPerSecond: h.bytesPerSec,
	}
	c.Writer = throttledWriter

	c.Header("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, artefact.FileName))
	c.Header("X-Checksum-SHA256", artefact.SHA256Hash)

	http.ServeContent(c.Writer, c.Request, artefact.FileName, fileInfo.ModTime(), file)

	// Mark as downloaded on completion (only for non-Range requests)
	rangeHeader := c.Request.Header.Get("Range")
	if deviceID != "" && deploymentID > 0 {
		if rangeHeader == "" || isCompleteRange(rangeHeader, fileInfo.Size()) {
			if err := database.UpdateDeploymentDeviceStatus(h.db, deploymentID, deviceID, "downloaded", ""); err != nil {
				log.Printf("[download] Failed to mark device downloaded: %v", err)
			}
		}
	}
}

// ThrottledResponseWriter limits write throughput per connection.
type ThrottledResponseWriter struct {
	gin.ResponseWriter
	bytesPerSecond int
}

func (w *ThrottledResponseWriter) Write(b []byte) (int, error) {
	chunkSize := w.bytesPerSecond / 10
	if chunkSize < 4096 {
		chunkSize = 4096
	}

	totalWritten := 0
	for totalWritten < len(b) {
		end := totalWritten + chunkSize
		if end > len(b) {
			end = len(b)
		}

		start := time.Now()
		n, err := w.ResponseWriter.Write(b[totalWritten:end])
		totalWritten += n

		if err != nil {
			return totalWritten, err
		}

		elapsed := time.Since(start)
		expected := time.Duration(n) * time.Second / time.Duration(w.bytesPerSecond)
		if expected > elapsed {
			time.Sleep(expected - elapsed)
		}
	}

	return totalWritten, nil
}

func isCompleteRange(rangeHeader string, fileSize int64) bool {
	// Simple check: if range starts at 0 and covers the entire file, it's "complete"
	// This avoids double-counting full-file requests that use Range: bytes=0-
	return false
}