// Package mediaimage 对 gateway 收到的完整图片字节做不解码像素的结构校验。
// GIF/WebP 只扫描容器和帧边界，避免攻击者用巨大画布或高帧数触发解压内存放大。
package mediaimage

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"image/jpeg"
	"image/png"
)

const (
	MaxFrames     = int32(120)
	MaxDurationMS = int64(15_000)
	MaxDimension  = int32(4096)
	MaxPixels     = int64(16_777_216)
)

// Info 是从真实文件内容解析出的规范化事实；调用方必须原样保存这些字节，不能转码后复用本结果。
type Info struct {
	ContentType string
	Format      string
	SizeBytes   int64
	Width       int32
	Height      int32
	Animated    bool
	FrameCount  int32
	DurationMS  int64
	SHA256      string
}

// Inspect 只以完整文件字节为图片格式事实来源，并拒绝 APNG 及畸形动画容器。
// multipart 文件名和 Content-Type 都由客户端提供，Android 相册缓存尤其可能保留错误后缀，
// 因而调用方必须使用这里返回的规范后缀与 Content-Type 保存对象，不能再信任上传元数据。
func Inspect(content []byte) (Info, string, error) {
	var info Info
	var extension string
	switch {
	case len(content) >= 3 && content[0] == 0xff && content[1] == 0xd8 && content[2] == 0xff:
		cfg, err := jpeg.DecodeConfig(bytes.NewReader(content))
		if err != nil {
			return info, "", fmt.Errorf("invalid jpeg")
		}
		info.Format, info.ContentType, info.Width, info.Height, info.FrameCount = "jpeg", "image/jpeg", int32(cfg.Width), int32(cfg.Height), 1
		extension = "jpg"
	case len(content) >= 8 && bytes.Equal(content[:8], []byte{137, 80, 78, 71, 13, 10, 26, 10}):
		if pngHasAnimationChunk(content) {
			return info, "", fmt.Errorf("animated png is not supported")
		}
		cfg, err := png.DecodeConfig(bytes.NewReader(content))
		if err != nil {
			return info, "", fmt.Errorf("invalid png")
		}
		info.Format, info.ContentType, info.Width, info.Height, info.FrameCount = "png", "image/png", int32(cfg.Width), int32(cfg.Height), 1
		extension = "png"
	case len(content) >= 6 && (string(content[:6]) == "GIF87a" || string(content[:6]) == "GIF89a"):
		width, height, frames, durationMS, err := inspectGIF(content)
		if err != nil {
			return info, "", err
		}
		info.Format, info.ContentType, info.Width, info.Height = "gif", "image/gif", width, height
		info.FrameCount, info.Animated, info.DurationMS = frames, frames > 1, durationMS
		if !info.Animated {
			info.DurationMS = 0
		}
		extension = "gif"
	case len(content) >= 12 && string(content[:4]) == "RIFF" && string(content[8:12]) == "WEBP":
		width, height, frames, durationMS, animated, err := inspectWebP(content)
		if err != nil {
			return info, "", err
		}
		info.Format, info.ContentType, info.Width, info.Height = "webp", "image/webp", width, height
		info.FrameCount, info.DurationMS, info.Animated = frames, durationMS, animated
		extension = "webp"
	default:
		return info, "", fmt.Errorf("only jpeg, png, gif and webp are supported")
	}
	if err := validateGeometry(info); err != nil {
		return Info{}, "", err
	}
	sum := sha256.Sum256(content)
	info.SizeBytes, info.SHA256 = int64(len(content)), hex.EncodeToString(sum[:])
	return info, extension, nil
}

func validateGeometry(info Info) error {
	if info.Width <= 0 || info.Height <= 0 || info.Width > MaxDimension || info.Height > MaxDimension ||
		int64(info.Width)*int64(info.Height) > MaxPixels || info.FrameCount <= 0 || info.FrameCount > MaxFrames {
		return fmt.Errorf("image dimensions or frame count exceed limit")
	}
	if info.Animated {
		if info.FrameCount < 2 || info.DurationMS <= 0 || info.DurationMS > MaxDurationMS {
			return fmt.Errorf("animation exceeds frame or duration limit")
		}
	} else if info.FrameCount != 1 || info.DurationMS != 0 {
		return fmt.Errorf("static image properties are invalid")
	}
	return nil
}

func pngHasAnimationChunk(content []byte) bool {
	for offset := 8; offset+12 <= len(content); {
		size := int(binary.BigEndian.Uint32(content[offset : offset+4]))
		if size < 0 || offset+12+size > len(content) {
			return false
		}
		if string(content[offset+4:offset+8]) == "acTL" {
			return true
		}
		offset += 12 + size
	}
	return false
}

func inspectGIF(content []byte) (int32, int32, int32, int64, error) {
	if len(content) < 13 {
		return 0, 0, 0, 0, fmt.Errorf("invalid gif")
	}
	width := int32(binary.LittleEndian.Uint16(content[6:8]))
	height := int32(binary.LittleEndian.Uint16(content[8:10]))
	if width <= 0 || height <= 0 || width > MaxDimension || height > MaxDimension || int64(width)*int64(height) > MaxPixels {
		return 0, 0, 0, 0, fmt.Errorf("gif dimensions exceed limit")
	}
	offset := 13
	if content[10]&0x80 != 0 {
		colorTableBytes := 3 * (1 << (int(content[10]&0x07) + 1))
		if offset+colorTableBytes > len(content) {
			return 0, 0, 0, 0, fmt.Errorf("invalid gif color table")
		}
		offset += colorTableBytes
	}
	var frames int32
	var durationMS, pendingDelayMS int64
	for offset < len(content) {
		marker := content[offset]
		offset++
		switch marker {
		case 0x3b:
			if frames <= 0 {
				return 0, 0, 0, 0, fmt.Errorf("gif has no image frame")
			}
			if frames == 1 {
				durationMS = 0
			}
			return width, height, frames, durationMS, nil
		case 0x21:
			if offset >= len(content) {
				return 0, 0, 0, 0, fmt.Errorf("invalid gif extension")
			}
			label := content[offset]
			offset++
			if label == 0xf9 {
				if offset+6 > len(content) || content[offset] != 4 || content[offset+5] != 0 {
					return 0, 0, 0, 0, fmt.Errorf("invalid gif graphic control extension")
				}
				pendingDelayMS = int64(binary.LittleEndian.Uint16(content[offset+2:offset+4])) * 10
				offset += 6
				continue
			}
			var err error
			offset, _, err = skipGIFSubBlocks(content, offset)
			if err != nil {
				return 0, 0, 0, 0, err
			}
		case 0x2c:
			if offset+9 > len(content) {
				return 0, 0, 0, 0, fmt.Errorf("invalid gif image descriptor")
			}
			left := int(binary.LittleEndian.Uint16(content[offset : offset+2]))
			top := int(binary.LittleEndian.Uint16(content[offset+2 : offset+4]))
			frameWidth := int(binary.LittleEndian.Uint16(content[offset+4 : offset+6]))
			frameHeight := int(binary.LittleEndian.Uint16(content[offset+6 : offset+8]))
			packed := content[offset+8]
			if frameWidth <= 0 || frameHeight <= 0 || left+frameWidth > int(width) || top+frameHeight > int(height) {
				return 0, 0, 0, 0, fmt.Errorf("invalid gif frame dimensions")
			}
			offset += 9
			if packed&0x80 != 0 {
				colorTableBytes := 3 * (1 << (int(packed&0x07) + 1))
				if offset+colorTableBytes > len(content) {
					return 0, 0, 0, 0, fmt.Errorf("invalid gif local color table")
				}
				offset += colorTableBytes
			}
			if offset >= len(content) || content[offset] < 2 || content[offset] > 8 {
				return 0, 0, 0, 0, fmt.Errorf("invalid gif lzw code size")
			}
			offset++
			var hasData bool
			var err error
			offset, hasData, err = skipGIFSubBlocks(content, offset)
			if err != nil || !hasData {
				return 0, 0, 0, 0, fmt.Errorf("invalid gif image data")
			}
			frames++
			durationMS += pendingDelayMS
			pendingDelayMS = 0
			if frames > MaxFrames || durationMS > MaxDurationMS {
				return 0, 0, 0, 0, fmt.Errorf("gif animation exceeds frame or duration limit")
			}
		default:
			return 0, 0, 0, 0, fmt.Errorf("invalid gif block marker")
		}
	}
	return 0, 0, 0, 0, fmt.Errorf("gif trailer is missing")
}

func skipGIFSubBlocks(content []byte, offset int) (int, bool, error) {
	hasData := false
	for {
		if offset >= len(content) {
			return 0, false, fmt.Errorf("invalid gif sub-block")
		}
		blockSize := int(content[offset])
		offset++
		if blockSize == 0 {
			return offset, hasData, nil
		}
		if offset+blockSize > len(content) {
			return 0, false, fmt.Errorf("invalid gif sub-block")
		}
		hasData = true
		offset += blockSize
	}
}

func inspectWebP(content []byte) (int32, int32, int32, int64, bool, error) {
	if len(content) < 20 {
		return 0, 0, 0, 0, false, fmt.Errorf("invalid webp header")
	}
	if uint64(binary.LittleEndian.Uint32(content[4:8]))+8 != uint64(len(content)) {
		return 0, 0, 0, 0, false, fmt.Errorf("invalid webp riff size")
	}
	var width, height int32
	var frames int32
	var durationMS int64
	var animatedFlag, hasVP8X, hasANIM, hasStaticImage bool
	for offset := 12; offset < len(content); {
		if offset+8 > len(content) {
			return 0, 0, 0, 0, false, fmt.Errorf("invalid webp trailing chunk")
		}
		chunkType := string(content[offset : offset+4])
		size := uint64(binary.LittleEndian.Uint32(content[offset+4 : offset+8]))
		start, end := uint64(offset+8), uint64(offset+8)+size
		paddedEnd := end + size%2
		if end < start || paddedEnd > uint64(len(content)) {
			return 0, 0, 0, 0, false, fmt.Errorf("invalid webp chunk")
		}
		if size%2 == 1 && content[end] != 0 {
			return 0, 0, 0, 0, false, fmt.Errorf("invalid webp chunk padding")
		}
		chunk := content[int(start):int(end)]
		switch chunkType {
		case "VP8X":
			if hasVP8X || offset != 12 || len(chunk) != 10 || chunk[0]&0xc1 != 0 || chunk[1] != 0 || chunk[2] != 0 || chunk[3] != 0 {
				return 0, 0, 0, 0, false, fmt.Errorf("invalid webp header")
			}
			hasVP8X, animatedFlag = true, chunk[0]&0x02 != 0
			width, height = int32(readUint24LE(chunk[4:7]))+1, int32(readUint24LE(chunk[7:10]))+1
		case "VP8 ", "VP8L":
			if hasStaticImage || (!hasVP8X && offset != 12) {
				return 0, 0, 0, 0, false, fmt.Errorf("invalid webp image payload order")
			}
			imageWidth, imageHeight, err := inspectWebPImagePayload(chunkType, chunk)
			if err != nil {
				return 0, 0, 0, 0, false, err
			}
			hasStaticImage = true
			if width == 0 {
				width, height = imageWidth, imageHeight
			} else if imageWidth > width || imageHeight > height {
				return 0, 0, 0, 0, false, fmt.Errorf("webp image exceeds canvas")
			}
		case "ANIM":
			if !hasVP8X || !animatedFlag || hasANIM || len(chunk) != 6 {
				return 0, 0, 0, 0, false, fmt.Errorf("invalid webp animation header")
			}
			hasANIM = true
		case "ANMF":
			if !hasVP8X || !animatedFlag || !hasANIM || width <= 0 || height <= 0 || len(chunk) < 16 || chunk[15]&0xfc != 0 {
				return 0, 0, 0, 0, false, fmt.Errorf("invalid webp animation frame")
			}
			frameX, frameY := int64(readUint24LE(chunk[0:3]))*2, int64(readUint24LE(chunk[3:6]))*2
			frameWidth, frameHeight := int32(readUint24LE(chunk[6:9]))+1, int32(readUint24LE(chunk[9:12]))+1
			if frameWidth <= 0 || frameHeight <= 0 || frameX+int64(frameWidth) > int64(width) || frameY+int64(frameHeight) > int64(height) {
				return 0, 0, 0, 0, false, fmt.Errorf("invalid webp animation frame geometry")
			}
			if err := inspectWebPFramePayload(chunk[16:], frameWidth, frameHeight); err != nil {
				return 0, 0, 0, 0, false, err
			}
			frames++
			durationMS += int64(readUint24LE(chunk[12:15]))
			if frames > MaxFrames || durationMS > MaxDurationMS {
				return 0, 0, 0, 0, false, fmt.Errorf("webp animation exceeds frame or duration limit")
			}
		}
		offset = int(paddedEnd)
	}
	if width <= 0 || height <= 0 {
		return 0, 0, 0, 0, false, fmt.Errorf("invalid webp dimensions")
	}
	if animatedFlag {
		if !hasANIM || hasStaticImage || frames < 2 || durationMS <= 0 {
			return 0, 0, 0, 0, false, fmt.Errorf("invalid animated webp")
		}
		return width, height, frames, durationMS, true, nil
	}
	if hasANIM || frames != 0 || !hasStaticImage {
		return 0, 0, 0, 0, false, fmt.Errorf("invalid static webp")
	}
	return width, height, 1, 0, false, nil
}

func inspectWebPImagePayload(chunkType string, chunk []byte) (int32, int32, error) {
	switch chunkType {
	case "VP8 ":
		if len(chunk) <= 10 || chunk[0]&0x01 != 0 || !bytes.Equal(chunk[3:6], []byte{0x9d, 0x01, 0x2a}) {
			return 0, 0, fmt.Errorf("invalid webp vp8 payload")
		}
		width, height := int32(binary.LittleEndian.Uint16(chunk[6:8])&0x3fff), int32(binary.LittleEndian.Uint16(chunk[8:10])&0x3fff)
		if width <= 0 || height <= 0 {
			return 0, 0, fmt.Errorf("invalid webp vp8 dimensions")
		}
		return width, height, nil
	case "VP8L":
		if len(chunk) <= 5 || chunk[0] != 0x2f {
			return 0, 0, fmt.Errorf("invalid webp vp8l payload")
		}
		bits := binary.LittleEndian.Uint32(chunk[1:5])
		if bits>>29 != 0 {
			return 0, 0, fmt.Errorf("invalid webp vp8l version")
		}
		return int32(bits&0x3fff) + 1, int32((bits>>14)&0x3fff) + 1, nil
	default:
		return 0, 0, fmt.Errorf("unsupported webp image payload")
	}
}

func inspectWebPFramePayload(content []byte, expectedWidth int32, expectedHeight int32) error {
	if len(content) < 8 {
		return fmt.Errorf("webp animation frame has no image payload")
	}
	foundImage := false
	for offset := 0; offset < len(content); {
		if offset+8 > len(content) {
			return fmt.Errorf("invalid webp animation frame chunk")
		}
		chunkType := string(content[offset : offset+4])
		size := uint64(binary.LittleEndian.Uint32(content[offset+4 : offset+8]))
		start, end := uint64(offset+8), uint64(offset+8)+size
		paddedEnd := end + size%2
		if end < start || paddedEnd > uint64(len(content)) || (size%2 == 1 && content[end] != 0) {
			return fmt.Errorf("invalid webp animation frame chunk")
		}
		chunk := content[int(start):int(end)]
		switch chunkType {
		case "ALPH":
			if foundImage || len(chunk) == 0 {
				return fmt.Errorf("invalid webp animation alpha payload")
			}
		case "VP8 ", "VP8L":
			if foundImage {
				return fmt.Errorf("webp animation frame has multiple image payloads")
			}
			width, height, err := inspectWebPImagePayload(chunkType, chunk)
			if err != nil || width != expectedWidth || height != expectedHeight {
				return fmt.Errorf("webp animation image dimensions mismatch")
			}
			foundImage = true
		default:
			return fmt.Errorf("unsupported webp animation frame chunk")
		}
		offset = int(paddedEnd)
	}
	if !foundImage {
		return fmt.Errorf("webp animation frame has no image payload")
	}
	return nil
}

func readUint24LE(value []byte) uint32 {
	return uint32(value[0]) | uint32(value[1])<<8 | uint32(value[2])<<16
}
