// Copyright 2024 E99p1ant. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package storage

import (
	"bytes"
	"context"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"io"
	"mime/multipart"
	"os"
	"os/exec"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/chai2010/webp"
	"github.com/disintegration/imaging"
	"github.com/pkg/errors"
	"github.com/thanhpk/randstr"
	"github.com/wuhan005/gadget"

	"github.com/tamakyi/TamaBox/internal/conf"
	"github.com/tamakyi/TamaBox/internal/imagecheck"
)

// MaxUploadBytes 是单次图片上传的内存读取上限（8MB）。
// 调用方虽在 HTTP 层对 fileHeader.Size 做了校验，但该值可被客户端伪造，
// 因此在读取前用 io.LimitReader 兜底，避免超大文件将内存撑爆（DoS）。
const MaxUploadBytes = 8 * 1024 * 1024

// aspect 约束一次上传的宽高比（wantW:wantH，容差 DefaultAspectTolerance）。
// 头像固定 1:1、页面背景固定 16:9：客户端裁剪后即为精确比例，
// 服务端复校验用于拦截绕过前端直接 POST 的越界图片。
type aspect struct {
	W, H float64
}

func (a *aspect) check(data []byte) error {
	if a == nil {
		return nil
	}
	return imagecheck.CheckAspectRatio(bytes.NewReader(data), a.W, a.H, imagecheck.DefaultAspectTolerance)
}

// UploadAvatarToS3 上传头像：除通用校验/转换外强制 1:1 宽高比。
func UploadAvatarToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, &aspect{W: 1, H: 1})
}

// UploadBackgroundToS3 上传页面背景：除通用校验/转换外强制 16:9 宽高比。
func UploadBackgroundToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, &aspect{W: 16, H: 9})
}

// UploadPictureToS3 上传通用图片到 Cloudflare R2，自动转换为 WebP（动画 GIF 可选转换）
func UploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader) (string, error) {
	return uploadPictureToS3(file, fileHeader, nil)
}

// uploadPictureToS3 是上传管线实现：读入 → 魔数 → （可选）比例校验 → 像素总量校验
// → 解码 → WebP 转换 → 落盘 R2。avatar/background 的固定比例校验复用同一管线，
// 仅通过 aspect 参数区分，避免为一种用途复制一套上传代码。
func uploadPictureToS3(file multipart.File, fileHeader *multipart.FileHeader, asp *aspect) (string, error) {
	// 1. 限制读取大小，防止超大上传撑爆内存（DoS）。
	//    先用 LimitReader 截断到 MaxUploadBytes+1 字节；若实际读到的字节数超过上限则直接拒绝。
	limited := io.LimitReader(file, MaxUploadBytes+1)
	fileData, err := io.ReadAll(limited)
	if err != nil {
		return "", errors.Wrap(err, "read file data")
	}
	if int64(len(fileData)) > MaxUploadBytes {
		return "", errors.New("uploaded file is too large")
	}

	// 2. 校验文件魔数，仅允许 JPEG / PNG / GIF / WebP
	if !isSupportedImage(fileData) {
		return "", errors.New("unsupported file type: not a valid image")
	}

	// 2.5 用途宽高比校验（头像 1:1 / 背景 16:9），绕过前端裁剪的直传在此被拦截
	if err := asp.check(fileData); err != nil {
		return "", err
	}

	fileReader := bytes.NewReader(fileData)

	// 3. 解码前先校验像素总量（防解压炸弹），再解码图片获取格式；解码失败则直接返回错误。
	if err := imagecheck.CheckPixelLimit(bytes.NewReader(fileData)); err != nil {
		return "", err
	}

	_, format, err := image.Decode(fileReader)
	if err != nil {
		return "", errors.Wrap(err, "image decode failed, file may be corrupted")
	}

	var uploadData []byte
	var contentType string

	// 重置 reader，以便后续可能的 resize 操作
	if _, err := fileReader.Seek(0, io.SeekStart); err != nil {
		return "", errors.Wrap(err, "reset reader")
	}

	// 4. 按格式转换或直接使用原图
	switch format {
	case "gif":
		webpData, err := convertGifToAnimatedWebP(fileData)
		if err != nil {
			uploadData = fileData // 转换失败回退原始 GIF
			contentType = "image/gif"
		} else {
			uploadData = webpData
			contentType = "image/webp"
		}
	default:
		// 静态图片（JPEG / PNG / WebP 等）转静态 WebP
		webpData, err := ConvertStaticImageToWebP(fileReader)
		if err != nil {
			uploadData = fileData // 转换失败回退原图
			contentType = mimeTypeFromFormat(format)
		} else {
			uploadData = webpData
			contentType = "image/webp"
		}
	}

	// 5. 生成存储路径（无后缀名）
	now := time.Now()
	year := now.Year()
	month := int(now.Month())
	day := now.Day()
	key := fmt.Sprintf("%s%d/%02d/%02d/%s", PictureKeyPrefix, year, month, day, randstr.Hex(15))

	// 6. 获取共享 S3 客户端
	client, err := GetS3Client()
	if err != nil {
		return "", errors.Wrap(err, "get S3 client")
	}

	// 7. 上传到 R2
	if err := gadget.Retry(5, func() error {
		_, err := client.PutObject(context.Background(), &s3.PutObjectInput{
			Bucket:        aws.String(conf.Upload.ImageBucket),
			Key:           aws.String(key),
			Body:          bytes.NewReader(uploadData),
			ContentLength: aws.Int64(int64(len(uploadData))),
			ContentType:   aws.String(contentType),
			// key 为 picture/日期/随机hex，内容不可变 → 可安全启用 immutable 强缓存，
			// 避免每次页面加载都回源校验（曾导致头像/二维码 logo 每次重新下载）。
			CacheControl: aws.String("public, max-age=31536000, immutable"),
		})
		return err
	}); err != nil {
		return "", errors.Wrap(err, "put object after retry")
	}

	// 8. 返回 CDN URL
	return fmt.Sprintf("https://%s/%s", conf.Upload.ImageBucketCDNHost, key), nil
}

// ConvertStaticImageToWebP 将静态图片（JPEG/PNG等）转换为 WebP 格式
// 限制最大宽度为 1920px，保持比例
func ConvertStaticImageToWebP(r io.Reader) ([]byte, error) {
	// 先整体读入并校验像素总量（防解压炸弹），再完整解码。
	data, err := io.ReadAll(r)
	if err != nil {
		return nil, errors.Wrap(err, "read image")
	}
	if err := imagecheck.CheckPixelLimit(bytes.NewReader(data)); err != nil {
		return nil, errors.Wrap(err, "check image pixel limit")
	}

	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, errors.Wrap(err, "decode image")
	}

	maxWidth := 1920
	if img.Bounds().Dx() > maxWidth {
		img = imaging.Resize(img, maxWidth, 0, imaging.Lanczos)
	}

	buf := new(bytes.Buffer)
	if err := webp.Encode(buf, img, &webp.Options{Quality: 80}); err != nil {
		return nil, errors.Wrap(err, "webp encode")
	}
	return buf.Bytes(), nil
}

// convertGifToAnimatedWebP 将 GIF 文件数据转换为动画 WebP（需要系统安装 cwebp 工具）
// 如果 cwebp 不可用，返回错误，调用方应回退到原始 GIF
func convertGifToAnimatedWebP(gifData []byte) ([]byte, error) {
	_, err := exec.LookPath("cwebp")
	if err != nil {
		return nil, errors.New("cwebp not found in PATH")
	}

	tmpGif, err := os.CreateTemp("", "input_*.gif")
	if err != nil {
		return nil, errors.Wrap(err, "create temp gif")
	}
	defer os.Remove(tmpGif.Name())
	if _, err := tmpGif.Write(gifData); err != nil {
		return nil, errors.Wrap(err, "write temp gif")
	}
	tmpGif.Close()

	tmpWebp, err := os.CreateTemp("", "output_*.webp")
	if err != nil {
		return nil, errors.Wrap(err, "create temp webp")
	}
	defer os.Remove(tmpWebp.Name())
	tmpWebp.Close()

	cmd := exec.Command("cwebp", "-q", "80", "-mixed", "-mt", "-quiet", tmpGif.Name(), "-o", tmpWebp.Name())
	if err := cmd.Run(); err != nil {
		return nil, errors.Wrap(err, "run cwebp")
	}

	webpData, err := os.ReadFile(tmpWebp.Name())
	if err != nil {
		return nil, errors.Wrap(err, "read webp output")
	}
	return webpData, nil
}

// mimeTypeFromFormat 根据图片格式返回 MIME 类型
func mimeTypeFromFormat(format string) string {
	switch format {
	case "jpeg":
		return "image/jpeg"
	case "png":
		return "image/png"
	case "gif":
		return "image/gif"
	case "webp":
		return "image/webp"
	default:
		return "application/octet-stream"
	}
}

// isSupportedImage 通过文件头魔数判断是否为支持的图片格式
func isSupportedImage(data []byte) bool {
	if len(data) < 4 {
		return false
	}

	// JPEG: FF D8 FF
	if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
		return true
	}

	// PNG: 89 50 4E 47
	if bytes.Equal(data[:4], []byte{0x89, 0x50, 0x4E, 0x47}) {
		return true
	}

	// GIF: GIF87a / GIF89a
	if len(data) >= 6 {
		if bytes.Equal(data[:6], []byte("GIF87a")) || bytes.Equal(data[:6], []byte("GIF89a")) {
			return true
		}
	}

	// WebP: RIFF ... WEBP
	if len(data) >= 12 && bytes.Equal(data[:4], []byte("RIFF")) && bytes.Equal(data[8:12], []byte("WEBP")) {
		return true
	}

	return false
}