"""Streaming multipart/form-data parser for the stdlib HTTP server.

The upload body is consumed in chunks from the request stream and never
assembled into memory as a whole. File parts are streamed straight to a temp
file on disk with an incremental size cap; text fields are capped at a small
field size. File bytes are preserved exactly (no trailing-byte trimming), the
first file part is the only one accepted, and both the extension and the file
signature are validated by the caller.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import re
from typing import BinaryIO

MAX_FIELD_BYTES = 16 * 1024
MAX_HEADER_BYTES = 32 * 1024
CHUNK_SIZE = 64 * 1024
# .xlsx files are ZIP containers (PK\x03\x04 local header or an empty-archive
# PK\x05\x06); .xls files are OLE2 compound documents.
XLSX_MAGIC = (b"PK\x03\x04", b"PK\x05\x06")
XLS_MAGIC = (b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",)


class MultipartError(ValueError):
    pass


@dataclass(frozen=True)
class UploadedFile:
    filename: str
    path: Path
    size: int


def valid_file_signature(path: Path, suffix: str) -> bool:
    """True when the file's leading bytes match its declared extension."""
    with path.open("rb") as handle:
        head = handle.read(8)
    if suffix == ".xlsx":
        return any(head.startswith(magic) for magic in XLSX_MAGIC)
    if suffix == ".xls":
        return any(head.startswith(magic) for magic in XLS_MAGIC)
    return False


class _MultipartStream:
    def __init__(
        self,
        rfile: BinaryIO,
        content_length: int,
        content_type: str,
        *,
        max_body_bytes: int,
    ) -> None:
        lowered = content_type.lower()
        if not lowered.startswith("multipart/form-data"):
            raise MultipartError("上传请求必须是 multipart/form-data。")
        match = re.search(r"boundary=(?:\"([^\"]+)\"|([^;]+))", content_type)
        if not match:
            raise MultipartError("上传请求缺少文件边界。")
        boundary = (match.group(1) or match.group(2)).strip().encode("utf-8")
        if not boundary or len(boundary) > 200:
            raise MultipartError("上传请求的文件边界无效。")
        if content_length <= 0:
            raise MultipartError("上传请求为空。")
        if content_length > max_body_bytes:
            raise MultipartError(f"上传请求超过 {max_body_bytes // (1024 * 1024)} MB 限制。")
        self.rfile = rfile
        self.remaining = content_length
        self.total_read = 0
        self.buffer = bytearray()
        self.boundary = boundary
        self._fill()

    # ------------------------------------------------------------------
    # Low-level stream helpers
    # ------------------------------------------------------------------

    def _fill(self) -> None:
        if self.total_read >= self.remaining:
            return
        want = min(CHUNK_SIZE, self.remaining - self.total_read)
        data = self.rfile.read(want)
        if not data:
            # Client closed early; clamp remaining so every helper sees EOF.
            self.remaining = self.total_read
            return
        self.total_read += len(data)
        self.buffer.extend(data)

    def _take(self, count: int) -> bytes:
        while len(self.buffer) < count:
            if self.total_read >= self.remaining:
                raise MultipartError("multipart 请求体不完整。")
            self._fill()
        out = bytes(self.buffer[:count])
        del self.buffer[:count]
        return out

    def _read_line(self) -> bytes:
        while True:
            index = self.buffer.find(b"\n")
            if index >= 0:
                line = bytes(self.buffer[: index + 1])
                del self.buffer[: index + 1]
                return line
            if self.total_read >= self.remaining:
                if not self.buffer:
                    raise MultipartError("multipart 请求体不完整。")
                line = bytes(self.buffer)
                self.buffer.clear()
                return line
            self._fill()

    def _skip_preamble(self) -> None:
        marker = b"--" + self.boundary
        while True:
            index = self.buffer.find(marker)
            if index >= 0:
                del self.buffer[: index + len(marker)]
                return
            keep = len(marker) - 1
            if len(self.buffer) > keep:
                del self.buffer[: len(self.buffer) - keep]
            if self.total_read >= self.remaining:
                raise MultipartError("上传请求中没有找到文件。")
            self._fill()

    def _iter_content(self, boundary: bytes):
        """Yield content bytes up to (excluding) the ``\\r\\n--boundary`` marker."""
        marker = b"\r\n--" + boundary
        keep = len(marker) - 1
        while True:
            index = self.buffer.find(marker)
            if index >= 0:
                content = bytes(self.buffer[:index])
                del self.buffer[:index]
                if content:
                    yield content
                return
            if len(self.buffer) > keep:
                safe = len(self.buffer) - keep
                content = bytes(self.buffer[:safe])
                del self.buffer[:safe]
                if content:
                    yield content
            if self.total_read >= self.remaining:
                raise MultipartError("multipart 请求体缺少结束边界。")
            self._fill()

    def _after_content_boundary(self) -> str:
        """Consume the content-ending marker; return 'part' or 'end'."""
        marker = b"\r\n--" + self.boundary
        del self.buffer[: len(marker)]
        indicator = self._take(2)
        if indicator == b"--":
            if self.buffer.startswith(b"\r\n"):
                del self.buffer[:2]
            elif self.buffer.startswith(b"\n"):
                del self.buffer[:1]
            return "end"
        if indicator != b"\r\n":
            raise MultipartError("上传请求格式无效。")
        return "part"

    # ------------------------------------------------------------------
    # Part parsing
    # ------------------------------------------------------------------

    def _read_headers(self) -> bytes:
        total = 0
        lines: list[bytes] = []
        while True:
            line = self._read_line()
            total += len(line)
            if total > MAX_HEADER_BYTES:
                raise MultipartError("multipart 头部过长。")
            if line in (b"\r\n", b"\n"):
                return b"".join(lines)
            lines.append(line)

    @staticmethod
    def _parse_disposition(header: bytes) -> tuple[str | None, str | None]:
        name = None
        filename = None
        name_match = re.search(br'name="([^"]*)"', header)
        if name_match:
            name = name_match.group(1).decode("utf-8", errors="replace")
        filename_match = re.search(br'filename="([^"]*)"', header)
        if filename_match:
            filename = filename_match.group(1).decode("utf-8", errors="replace")
        return name, filename

    def parse(self, work_dir: Path) -> tuple[dict[str, str], UploadedFile | None]:
        import tempfile

        self._skip_preamble()
        next_bytes = self._take(2)
        if next_bytes == b"--":
            # Empty multipart body: opening boundary immediately closes.
            return {}, None
        if next_bytes != b"\r\n":
            raise MultipartError("上传请求格式无效。")

        fields: dict[str, str] = {}
        uploaded: UploadedFile | None = None
        while True:
            header = self._read_headers()
            name, filename = self._parse_disposition(header)
            if name is None:
                raise MultipartError("上传请求缺少字段名。")

            if filename is not None:
                if uploaded is not None:
                    raise MultipartError("一次只能上传一个文件。")
                with tempfile.NamedTemporaryFile(
                    dir=str(work_dir), prefix="upload-", delete=False
                ) as sink:
                    temp_path = Path(sink.name)
                    size = 0
                    try:
                        for chunk in self._iter_content(self.boundary):
                            size += len(chunk)
                            sink.write(chunk)
                    except Exception:
                        temp_path.unlink(missing_ok=True)
                        raise
                uploaded = UploadedFile(
                    filename=filename, path=temp_path, size=size
                )
                if self._after_content_boundary() == "end":
                    break
            else:
                chunks: list[bytes] = []
                total = 0
                for chunk in self._iter_content(self.boundary):
                    total += len(chunk)
                    if total > MAX_FIELD_BYTES:
                        raise MultipartError("表单字段超过大小限制。")
                    chunks.append(chunk)
                fields[name] = b"".join(chunks).decode(
                    "utf-8", errors="replace"
                ).strip()
                if self._after_content_boundary() == "end":
                    break
        return fields, uploaded


def parse_upload(
    rfile: BinaryIO,
    content_length: int,
    content_type: str,
    work_dir: str | Path,
    *,
    max_file_bytes: int,
) -> tuple[dict[str, str], UploadedFile | None]:
    """Stream a multipart upload; returns ``(fields, uploaded_file)``.

    Raises :class:`MultipartError` for malformed or oversized bodies. The
    caller owns the uploaded temp file and must remove it when done.
    """
    work = Path(work_dir)
    work.mkdir(parents=True, exist_ok=True)
    stream = _MultipartStream(
        rfile, content_length, content_type, max_body_bytes=max_file_bytes + MAX_FIELD_BYTES
    )
    fields, uploaded = stream.parse(work)
    if uploaded is None:
        raise MultipartError("上传请求中没有找到文件。")
    if uploaded.size > max_file_bytes:
        raise MultipartError(f"文件超过 {max_file_bytes // (1024 * 1024)} MB 限制。")
    return fields, uploaded
