"""文件上传安全校验。

提供 validate_upload() 对上传文件做：
1. 文件名清洗（去除路径、危险字符、控制长度）
2. 扩展名白名单校验
3. MIME 与扩展名一致性校验
4. 文件魔数（magic bytes）校验，防止伪装
5. 文件大小校验
"""
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Optional

from fastapi import HTTPException, UploadFile


# 允许的扩展名 → 期望的 MIME 前缀（; 前部分）
ALLOWED_TYPES: dict[str, str] = {
    "txt": "text/",
    "md": "text/",
    "csv": "text/",
    "json": "application/json",
    "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
    "pdf": "application/pdf",
    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "xls": "application/vnd.ms-excel",
    "png": "image/png",
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "gif": "image/gif",
    "webp": "image/webp",
    "bmp": "image/bmp",
}

# 魔数签名（前若干字节）
MAGIC_SIGNATURES: dict[str, list[bytes]] = {
    "pdf": [b"%PDF-"],
    "png": [b"\x89PNG\r\n\x1a\n"],
    "jpg": [b"\xff\xd8\xff"],
    "jpeg": [b"\xff\xd8\xff"],
    "gif": [b"GIF87a", b"GIF89a"],
    "bmp": [b"BM"],
    "webp": [b"RIFF"],
    # Office 文件均为 ZIP 容器（PK\x03\x04）
    "docx": [b"PK\x03\x04"],
    "pptx": [b"PK\x03\x04"],
    "xlsx": [b"PK\x03\x04"],
    "xls": [b"\xd0\xcf\x11\xe0"],  # OLE 复合文档
}

MAX_FILENAME_LEN = 200


@dataclass
class UploadCheck:
    filename: str
    suffix: str
    media_type: str
    size: int


_FILENAME_BAD = re.compile("[\x00-\x1f<>:\"/\\|?*]")


def sanitize_filename(name: str) -> str:
    """清洗文件名：去路径、去控制/危险字符、限长。"""
    # 仅保留文件名部分
    name = re.split(r"[\\/]", name)[-1].strip()
    name = _FILENAME_BAD.sub("_", name)
    name = name.strip(". ")
    if not name:
        name = "upload"
    if len(name) > MAX_FILENAME_LEN:
        stem, dot, ext = name.rpartition(".")
        if dot:
            name = stem[: MAX_FILENAME_LEN - len(ext) - 1] + "." + ext
        else:
            name = name[:MAX_FILENAME_LEN]
    return name


def _check_magic(suffix: str, head: bytes) -> bool:
    sigs = MAGIC_SIGNATURES.get(suffix)
    if not sigs:
        return True  # 无签名要求（文本类）默认放行
    return any(head.startswith(sig) for sig in sigs)


def validate_upload(
    file: UploadFile,
    *,
    content: bytes,
    max_size: Optional[int] = None,
) -> UploadCheck:
    """校验上传文件，失败抛出 HTTPException。

    参数 content 为已读取的字节内容（用于魔数校验）。
    """
    if max_size is not None and len(content) > max_size:
        raise HTTPException(status_code=413, detail=f"文件不能超过 {max_size // (1024 * 1024)}MB")
    if not content:
        raise HTTPException(status_code=400, detail="上传文件为空")

    raw_name = file.filename or "material"
    filename = sanitize_filename(raw_name)
    suffix = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    media_type = (file.content_type or "").split(";")[0].strip().lower()

    if suffix not in ALLOWED_TYPES:
        raise HTTPException(
            status_code=415,
            detail="暂不支持该文件类型，请上传 txt、md、csv、json、docx、pptx、pdf、xlsx 或图片",
        )

    # MIME 一致性：图片/office/pdf 必须匹配；文本类型放宽
    expected_mime = ALLOWED_TYPES[suffix]
    if expected_mime != "text/" and media_type and not media_type.startswith(expected_mime.split("/")[0]):
        # 宽松：只校验主类型一致（image/* application/* 等）
        if media_type.split("/")[0] != expected_mime.split("/")[0]:
            raise HTTPException(status_code=415, detail="文件类型与扩展名不一致")

    # 魔数校验
    head = content[:16]
    if not _check_magic(suffix, head):
        raise HTTPException(status_code=415, detail="文件内容与声明类型不符")

    return UploadCheck(filename=filename, suffix=suffix, media_type=media_type or expected_mime, size=len(content))