"""图片校验与元信息解析。不依赖 Pillow，只读文件头，避免解码带来的攻击面和额外依赖。"""
from __future__ import annotations

import re
import struct

from fastapi import HTTPException


MAX_IMAGE_BYTES = 8 * 1024 * 1024
MAX_THUMBNAIL_BYTES = 512 * 1024
MAX_IMAGES_PER_RECORD = 6
MAX_PIXELS = 60_000_000
MAX_EDGE = 20_000
# 只收 JPEG 和 PNG：客户端一律用 canvas 转成 JPEG（小尺寸 PNG 原样透传），
# WebP 没有任何收益，而它的 RIFF 头几乎不校验内容，很容易被拿来夹带任意二进制。
ALLOWED_TYPES = ("image/jpeg", "image/png")
EXTENSIONS = {"image/jpeg": ".jpg", "image/png": ".png"}


def _png(data: bytes) -> tuple[int, int] | None:
    if data[:8] != b"\x89PNG\r\n\x1a\n" or data[12:16] != b"IHDR":
        return None
    return struct.unpack(">II", data[16:24])


JPEG_SOF_MARKERS = {0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF}


def _jpeg(data: bytes) -> tuple[int, int] | None:
    if data[:3] != b"\xff\xd8\xff":
        return None
    index, size = 2, len(data)
    while index + 3 < size:
        if data[index] != 0xFF:
            return None
        marker = data[index + 1]
        if marker == 0xFF:
            index += 1
            continue
        if marker == 0x01 or 0xD0 <= marker <= 0xD9:
            index += 2
            continue
        length = struct.unpack(">H", data[index + 2:index + 4])[0]
        if length < 2:
            return None
        if marker in JPEG_SOF_MARKERS:
            if index + 9 > size:
                return None
            height, width = struct.unpack(">HH", data[index + 5:index + 9])
            return width, height
        if marker == 0xDA:
            return None
        index += 2 + length
    return None


PARSERS = (("image/png", _png), ("image/jpeg", _jpeg))


def trim_trailing_payload(data: bytes, media_type: str) -> bytes:
    """截掉图像结束标记之后的内容。

    否则可以在一张合法图片后面附上任意字节，把本站变成任意文件的托管点；
    结束标记之后的数据本来就不属于图像，丢掉不会影响任何正常图片。
    """
    if media_type == "image/jpeg":
        end = data.rfind(b"\xff\xd9")
        return data[:end + 2] if end > 0 else data
    if media_type == "image/png":
        end = data.rfind(b"IEND")
        return data[:end + 8] if end > 0 else data   # IEND 后面还有 4 字节 CRC
    return data


def sniff(data: bytes) -> tuple[str, int, int] | None:
    """按文件头识别类型并取出尺寸；无法解析即视为非法图片。"""
    for media_type, parser in PARSERS:
        try:
            size = parser(data)
        except (struct.error, IndexError):
            size = None
        if size and size[0] > 0 and size[1] > 0:
            return media_type, size[0], size[1]
    return None


UNSAFE_NAME_CHARS = re.compile(r"[^\w.\-一-鿿（）() ]")


def safe_original_name(raw: str | None, content_type: str) -> str:
    """文件名完全由客户端 multipart 头控制，可能带引号、尖括号等用于注入的字符，入库前统一清洗。"""
    name = (raw or "").strip().replace("\\", "/").rsplit("/", 1)[-1]
    name = UNSAFE_NAME_CHARS.sub("_", name).strip("._ ")[:120]
    return name or f"image{EXTENSIONS.get(content_type, '.jpg')}"


def inspect_image(data: bytes, limit: int = MAX_IMAGE_BYTES, label: str = "图片") -> tuple[str, int, int]:
    """校验图片字节并返回 (content_type, width, height)，任何不合规都抛出中文错误。"""
    if not data:
        raise HTTPException(400, f"{label}内容为空")
    if len(data) > limit:
        raise HTTPException(413, f"{label}不能超过 {limit // (1024 * 1024) or 1} MB")
    result = sniff(data)
    if not result:
        raise HTTPException(415, f"{label}格式不支持，仅接受 JPEG 或 PNG")
    media_type, width, height = result
    if width > MAX_EDGE or height > MAX_EDGE or width * height > MAX_PIXELS:
        raise HTTPException(413, f"{label}分辨率过大，请压缩后再上传")
    return media_type, width, height