"""内容净化与上传安全 — 文件类型白名单、大小限制、zip/image bomb 检测、SSRF 防护、Markdown 净化。

安全边界：
- 文件类型白名单（pdf/docx/xlsx/pptx/txt/md/csv/html）
- 文件大小上限（50MB 默认）
- ZIP bomb 检测（.docx/.xlsx/.pptx 本质是 ZIP）
- Image bomb 检测（像素数 > 100MP）
- SSRF 防护（拒绝 loopback/RFC1918/link-local/ULA）
- Markdown 危险 HTML 标签净化（script/iframe/object/embed）
- XXE 防护（forbid_dtd XML 解析）
"""

from __future__ import annotations

import ipaddress
import re
import struct
import zipfile
from pathlib import Path
from typing import Any

# 文件类型白名单：扩展名 → MIME 类型
ALLOWED_FILE_TYPES: dict[str, str] = {
    ".pdf": "application/pdf",
    ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    ".xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    ".pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
    ".txt": "text/plain",
    ".md": "text/markdown",
    ".csv": "text/csv",
    ".html": "text/html",
}

# 文件大小上限：50MB
MAX_FILE_SIZE = 50 * 1024 * 1024

# ZIP bomb 检测阈值
MAX_ZIP_RATIO = 100  # 压缩比上限 100:1
MAX_ZIP_DECOMPRESSED_SIZE = 500 * 1024 * 1024  # 解压后上限 500MB

# Image bomb 检测阈值
MAX_IMAGE_PIXELS = 100_000_000  # 100MP

# 需要净化的文本格式
_TEXT_FORMATS = frozenset({"md", "html", "txt", "csv"})

# 危险 HTML 标签（含内容整体移除）
_DANGEROUS_TAG_RE = re.compile(
    r"<\s*(script|iframe|object|embed|style|form|input|textarea|button|link|meta)"
    r"[^>]*>.*?<\s*/\s*\1\s*>",
    re.IGNORECASE | re.DOTALL,
)
# 危险自闭合/孤立标签
_DANGEROUS_SOLO_RE = re.compile(
    r"<\s*(script|iframe|object|embed|style|form|input|textarea|button|link|meta)[^>]*/?>",
    re.IGNORECASE,
)
# 危险闭合标签
_DANGEROUS_CLOSE_RE = re.compile(
    r"<\s*/\s*(script|iframe|object|embed|style|form|input|textarea|button|link|meta)\s*>",
    re.IGNORECASE,
)
# javascript: 协议
_JS_URL_RE = re.compile(r"javascript\s*:", re.IGNORECASE)
# 事件处理器属性 (onclick=, onload=, onerror=, ...)
_EVENT_HANDLER_RE = re.compile(
    r"\s+on\w+\s*=\s*(?:\"[^\"]*\"|'[^']*'|[^\s>]+)",
    re.IGNORECASE,
)

# SSRF 防护：拒绝的网络段
_DENIED_NETWORKS = [
    ipaddress.ip_network("127.0.0.0/8"),  # IPv4 loopback
    ipaddress.ip_network("10.0.0.0/8"),  # RFC1918
    ipaddress.ip_network("172.16.0.0/12"),  # RFC1918
    ipaddress.ip_network("192.168.0.0/16"),  # RFC1918
    ipaddress.ip_network("169.254.0.0/16"),  # link-local
    ipaddress.ip_network("0.0.0.0/8"),  # "this network"
    ipaddress.ip_network("::1/128"),  # IPv6 loopback
    ipaddress.ip_network("fc00::/7"),  # IPv6 ULA
    ipaddress.ip_network("fe80::/10"),  # IPv6 link-local
    ipaddress.ip_network("::/128"),  # IPv6 unspecified
]


def validate_file_type(filename: str, file_type: str | None = None) -> str:
    """校验文件类型是否在白名单内。

    Args:
        filename: 文件名（用于提取扩展名）
        file_type: 可选 MIME 类型（当前仅用于日志，不做强制校验）

    Returns:
        归一化文件类型（扩展名，不含点），如 "pdf"、"md"

    Raises:
        ValueError: 文件类型不在白名单
    """
    ext = Path(filename).suffix.lower()
    if ext not in ALLOWED_FILE_TYPES:
        allowed = ", ".join(sorted(ALLOWED_FILE_TYPES.keys()))
        raise ValueError(f"File type '{ext}' is not allowed. Allowed types: {allowed}")
    return ext.lstrip(".")


def validate_file_size(file_size: int) -> None:
    """校验文件大小是否在上限内。

    Args:
        file_size: 文件大小（字节）

    Raises:
        ValueError: 文件超过 MAX_FILE_SIZE
    """
    if file_size > MAX_FILE_SIZE:
        raise ValueError(
            f"File size {file_size} bytes exceeds limit {MAX_FILE_SIZE} bytes "
            f"({MAX_FILE_SIZE // (1024 * 1024)}MB)"
        )
    if file_size <= 0:
        raise ValueError(f"File size must be positive, got {file_size}")


def check_zip_bomb(file_path: str) -> None:
    """检测 ZIP bomb — 适用于 .docx/.xlsx/.pptx（本质是 ZIP 格式）。

    拒绝条件：
    - 单文件压缩比 > 100:1
    - 总解压大小 > 500MB

    Args:
        file_path: ZIP 格式文件路径

    Raises:
        ValueError: 检测到 ZIP bomb
        zipfile.BadZipFile: 文件不是有效 ZIP
    """
    with zipfile.ZipFile(file_path) as zf:
        total_compressed = 0
        total_uncompressed = 0

        for info in zf.infolist():
            total_compressed += info.compress_size
            total_uncompressed += info.file_size

            # 单文件压缩比检查
            if info.compress_size > 0:
                ratio = info.file_size / info.compress_size
                if ratio > MAX_ZIP_RATIO:
                    raise ValueError(
                        f"Zip bomb detected: file '{info.filename}' has compression "
                        f"ratio {ratio:.1f}:1 (max {MAX_ZIP_RATIO}:1)"
                    )

        # 总解压大小检查
        if total_uncompressed > MAX_ZIP_DECOMPRESSED_SIZE:
            raise ValueError(
                f"Zip bomb detected: total uncompressed size "
                f"{total_uncompressed} bytes exceeds limit "
                f"{MAX_ZIP_DECOMPRESSED_SIZE} bytes"
            )

        # 总体压缩比检查
        if total_compressed > 0:
            overall_ratio = total_uncompressed / total_compressed
            if overall_ratio > MAX_ZIP_RATIO:
                raise ValueError(
                    f"Zip bomb detected: overall compression ratio "
                    f"{overall_ratio:.1f}:1 (max {MAX_ZIP_RATIO}:1)"
                )


def check_image_bomb(image_path: str) -> None:
    """检测 image bomb — 像素数超过 100MP 则拒绝。

    支持 PNG、JPEG、GIF 格式（通过文件头解析，不依赖 PIL）。
    无法识别的格式视为安全（返回不抛异常）。

    Args:
        image_path: 图片文件路径

    Raises:
        ValueError: 像素数超过 MAX_IMAGE_PIXELS
    """
    try:
        with open(image_path, "rb") as f:
            header = f.read(32)
    except OSError:
        return  # 无法读取，视为安全

    width, height = _read_image_dimensions(header, image_path)
    if width <= 0 or height <= 0:
        return  # 无法识别尺寸，视为安全

    pixels = width * height
    if pixels > MAX_IMAGE_PIXELS:
        raise ValueError(
            f"Image bomb detected: {width}x{height} = {pixels} pixels "
            f"exceeds limit {MAX_IMAGE_PIXELS} pixels (100MP)"
        )


def _read_image_dimensions(header: bytes, file_path: str) -> tuple[int, int]:
    """从文件头读取图片尺寸（PNG/GIF/JPEG）。"""
    # PNG: \x89PNG\r\n\x1a\n + IHDR (width@16, height@20, big-endian)
    if len(header) >= 24 and header[:8] == b"\x89PNG\r\n\x1a\n":
        width = struct.unpack(">I", header[16:20])[0]
        height = struct.unpack(">I", header[20:24])[0]
        return width, height

    # GIF: GIF87a/GIF89a (width@6, height@8, little-endian)
    if len(header) >= 10 and header[:6] in (b"GIF87a", b"GIF89a"):
        width = struct.unpack("<H", header[6:8])[0]
        height = struct.unpack("<H", header[8:10])[0]
        return width, height

    # JPEG: \xff\xd8 — 需要扫描 SOF0/SOF2 标记
    if len(header) >= 2 and header[:2] == b"\xff\xd8":
        return _read_jpeg_dimensions(file_path)

    return 0, 0


def _read_jpeg_dimensions(file_path: str) -> tuple[int, int]:
    """扫描 JPEG 文件查找 SOF0/SOF2 标记以获取尺寸。"""
    try:
        with open(file_path, "rb") as f:
            f.read(2)  # 跳过 SOI 标记
            while True:
                marker = f.read(2)
                if len(marker) < 2:
                    return 0, 0
                if marker[0] != 0xFF:
                    continue
                # SOF0 (0xC0) 或 SOF2 (0xC2) 包含尺寸信息
                if marker[1] in (0xC0, 0xC2):
                    f.read(3)  # 跳过 length(2) + precision(1)
                    height = struct.unpack(">H", f.read(2))[0]
                    width = struct.unpack(">H", f.read(2))[0]
                    return width, height
                else:
                    # 跳过此 marker 段
                    length_bytes = f.read(2)
                    if len(length_bytes) < 2:
                        return 0, 0
                    length = struct.unpack(">H", length_bytes)[0]
                    f.read(length - 2)
    except (OSError, struct.error):
        return 0, 0


def is_safe_ip(ip_str: str) -> bool:
    """SSRF 防护：检查 IP 是否安全（非 loopback/RFC1918/link-local/ULA）。

    Args:
        ip_str: IP 地址字符串

    Returns:
        True 如果 IP 是安全的公网地址，False 如果在拒绝的网络段内
    """
    try:
        ip = ipaddress.ip_address(ip_str)
    except ValueError:
        return False  # 无效 IP，拒绝

    for network in _DENIED_NETWORKS:
        if ip in network:
            return False
    return True


def sanitize_markdown(content: str) -> str:
    """净化 Markdown 内容 — 移除危险 HTML 标签和属性。

    移除内容：
    - 危险标签（script/iframe/object/embed/style/form/input/textarea/button/link/meta）及其内容
    - javascript: 协议
    - 事件处理器属性（onclick/onload/onerror 等）

    Args:
        content: 原始 Markdown/HTML 文本

    Returns:
        净化后的文本
    """
    # 移除危险标签及其内容（如 <script>...</script>）
    result = _DANGEROUS_TAG_RE.sub("", content)
    # 移除孤立的危险标签（如 <script>、<iframe ...>）
    result = _DANGEROUS_SOLO_RE.sub("", result)
    # 移除孤立的闭合标签（如 </script>）
    result = _DANGEROUS_CLOSE_RE.sub("", result)
    # 移除 javascript: 协议
    result = _JS_URL_RE.sub("", result)
    # 移除事件处理器属性
    result = _EVENT_HANDLER_RE.sub("", result)
    return result


def sanitize_content(content: str, file_type: str) -> str:
    """内容净化主入口 — 对文本格式应用 Markdown 净化。

    Args:
        content: 文本内容
        file_type: 文件类型（扩展名，不含点），如 "md"、"html"、"pdf"

    Returns:
        净化后的内容（二进制格式原样返回）
    """
    if file_type in _TEXT_FORMATS:
        return sanitize_markdown(content)
    return content


def parse_xml_safe(content: bytes) -> Any:
    """安全解析 XML — 禁止 DTD/实体以防止 XXE 攻击。

    Args:
        content: XML 字节内容

    Returns:
        ElementTree Element

    Raises:
        xml.etree.ElementTree.ParseError: 解析失败或包含 DTD
    """
    import xml.etree.ElementTree as ET

    parser = ET.XMLParser(
        forbid_dtd=True,
        forbid_entities=True,
        forbid_external=True,
    )
    return ET.fromstring(content, parser=parser)


__all__ = [
    "ALLOWED_FILE_TYPES",
    "MAX_FILE_SIZE",
    "MAX_IMAGE_PIXELS",
    "MAX_ZIP_DECOMPRESSED_SIZE",
    "MAX_ZIP_RATIO",
    "check_image_bomb",
    "check_zip_bomb",
    "is_safe_ip",
    "parse_xml_safe",
    "sanitize_content",
    "sanitize_markdown",
    "validate_file_size",
    "validate_file_type",
]
