portal_page/handlers/upload.go

package handlers

import (
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/disintegration/imaging"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
)

// allowedMIMETypes defines the permitted upload file MIME types.
var allowedMIMETypes = map[string]string{
	"image/jpeg": ".jpg",
	"image/png":  ".png",
	"image/gif":  ".gif",
}

// maxUploadSize is the maximum allowed file size in bytes (5 MB).
const maxUploadSize = 5 << 20

// uploadDir is the directory where uploaded files are stored.
const uploadDir = "./data/uploads"

// thumbSuffix is the suffix appended to compressed image filenames.
const thumbSuffix = "_thumb"

// UploadResponse is the JSON response returned after a successful upload.
type UploadResponse struct {
	URL string `json:"url"`
}

// UploadHandler handles image upload requests.
// POST /admin/upload
// Accepts multipart form with "file" field and optional "type" parameter ("icon" or "background").
// Returns JSON: {"url": "/uploads/{uuid}.{ext}"}
func UploadHandler(c *gin.Context) {
	// Parse multipart form with size limit
	if err := c.Request.ParseMultipartForm(maxUploadSize); err != nil {
		c.JSON(400, gin.H{"error": "文件太大，最大允许 5MB"})
		return
	}

	file, header, err := c.Request.FormFile("file")
	if err != nil {
		c.JSON(400, gin.H{"error": "请选择要上传的文件"})
		return
	}
	defer file.Close()

	// Validate file size
	if header.Size > maxUploadSize {
		c.JSON(400, gin.H{"error": "文件太大，最大允许 5MB"})
		return
	}

	// Validate MIME type by reading first 512 bytes
	buf := make([]byte, 512)
	n, err := file.Read(buf)
	if err != nil && err != io.EOF {
		c.JSON(500, gin.H{"error": "读取文件失败"})
		return
	}
	// Seek back to start
	if _, err := file.Seek(0, 0); err != nil {
		c.JSON(500, gin.H{"error": "处理文件失败"})
		return
	}

	mimeType := strings.Split(c.Request.Header.Get("Content-Type"), ";")[0]
	// Also detect from file content
	detectedType := detectMIMEType(buf[:n])
	if detectedType == "" {
		detectedType = mimeType
	}

	ext, ok := allowedMIMETypes[detectedType]
	if !ok {
		// Try the header's content type as fallback
		ext, ok = allowedMIMETypes[mimeType]
		if !ok {
			c.JSON(400, gin.H{"error": "不支持的文件格式，仅支持 jpg, png, gif"})
			return
		}
	}

	// Generate UUID filename
	fileUUID := uuid.New().String()
	filename := fileUUID + ext

	// Ensure upload directory exists
	if err := os.MkdirAll(uploadDir, 0755); err != nil {
		c.JSON(500, gin.H{"error": "创建上传目录失败"})
		return
	}

	// Save original file
	originalPath := filepath.Join(uploadDir, filename)
	dst, err := os.Create(originalPath)
	if err != nil {
		c.JSON(500, gin.H{"error": "保存文件失败"})
		return
	}
	defer dst.Close()

	if _, err := io.Copy(dst, file); err != nil {
		os.Remove(originalPath)
		c.JSON(500, gin.H{"error": "保存文件失败"})
		return
	}
	dst.Close()

	// Generate compressed thumbnail
	uploadType := c.PostForm("type")
	if err := generateThumbnail(originalPath, fileUUID, uploadType); err != nil {
		// Log the error but don't fail the upload — thumbnail is optional
		fmt.Printf("Warning: failed to generate thumbnail for %s: %v\n", filename, err)
	}

	url := "/uploads/" + filename
	c.JSON(200, UploadResponse{URL: url})
}

// ServeUploadHandler serves uploaded files.
// GET /uploads/:filename
// If query param thumb=1, returns the compressed version.
func ServeUploadHandler(c *gin.Context) {
	filename := c.Param("filename")
	if filename == "" {
		c.Status(404)
		return
	}

	// Prevent directory traversal
	if strings.Contains(filename, "..") || strings.Contains(filename, "/") || strings.Contains(filename, "\\") {
		c.Status(404)
		return
	}

	filePath := filepath.Join(uploadDir, filename)

	// Check if thumb=1 query parameter is requested
	if c.Query("thumb") == "1" {
		// Try to serve the thumbnail version
		ext := filepath.Ext(filename)
		baseName := filename[:len(filename)-len(ext)]
		thumbPath := filepath.Join(uploadDir, baseName+thumbSuffix+".jpg")

		if _, err := os.Stat(thumbPath); err == nil {
			c.File(thumbPath)
			return
		}
		// If thumbnail doesn't exist, fall through to serve original
	}

	// Serve original file
	if _, err := os.Stat(filePath); os.IsNotExist(err) {
		c.Status(404)
		return
	}
	c.File(filePath)
}

// detectMIMEType detects the MIME type from file content bytes.
func detectMIMEType(data []byte) string {
	if len(data) < 3 {
		return ""
	}

	// JPEG: starts with FF D8 FF
	if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
		return "image/jpeg"
	}
	// PNG: starts with 89 50 4E 47
	if len(data) >= 4 && data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47 {
		return "image/png"
	}
	// GIF: starts with "GIF"
	if len(data) >= 3 && data[0] == 'G' && data[1] == 'I' && data[2] == 'F' {
		return "image/gif"
	}
	return ""
}

// generateThumbnail creates a compressed version of the uploaded image.
// Icon type: max 200x200, quality 80
// Background type: max 1920x1080, quality 85
// Default: max 800x600, quality 85
// Thumbnails are saved as JPG format.
func generateThumbnail(originalPath, fileUUID, uploadType string) error {
	// Open the original image
	img, err := imaging.Open(originalPath)
	if err != nil {
		return fmt.Errorf("failed to open image for thumbnail: %w", err)
	}

	var maxWidth, maxHeight int
	var quality int

	switch uploadType {
	case "icon":
		maxWidth = 200
		maxHeight = 200
		quality = 80
	case "background":
		maxWidth = 1920
		maxHeight = 1080
		quality = 85
	default:
		maxWidth = 800
		maxHeight = 600
		quality = 85
	}

	// Resize if needed, maintaining aspect ratio
	thumb := fitImage(img, maxWidth, maxHeight)

	// Save thumbnail as JPEG
	thumbPath := filepath.Join(uploadDir, fileUUID+thumbSuffix+".jpg")
	thumbFile, err := os.Create(thumbPath)
	if err != nil {
		return fmt.Errorf("failed to create thumbnail file: %w", err)
	}
	defer thumbFile.Close()

	if err := jpeg.Encode(thumbFile, thumb, &jpeg.Options{Quality: quality}); err != nil {
		os.Remove(thumbPath)
		return fmt.Errorf("failed to encode thumbnail: %w", err)
	}

	return nil
}

// fitImage resizes the image to fit within maxWidth x maxHeight while
// maintaining aspect ratio. If the image is already smaller, it is returned as-is.
func fitImage(img image.Image, maxWidth, maxHeight int) image.Image {
	bounds := img.Bounds()
	w := bounds.Dx()
	h := bounds.Dy()

	if w <= maxWidth && h <= maxHeight {
		return img
	}

	// Calculate the scaling factor to fit within bounds
	scaleW := float64(maxWidth) / float64(w)
	scaleH := float64(maxHeight) / float64(h)
	scale := scaleW
	if scaleH < scaleW {
		scale = scaleH
	}

	newW := int(float64(w) * scale)
	newH := int(float64(h) * scale)

	if newW < 1 {
		newW = 1
	}
	if newH < 1 {
		newH = 1
	}

	return imaging.Resize(img, newW, newH, imaging.Lanczos)
}