"""File validation and sanitization for the API plugin.

NOTE: Runs inside the Airflow api-server process. Do NOT import from src.*.
Mirrors logic from src/services/file_validator.py and
src/extraction_v2/password_detection.py.
"""

import os
from pathlib import Path

# Supported file extensions
VALID_EXTENSIONS = {
    ".xlsx", ".xls", ".xlsm",
    ".docx", ".doc", ".docm",
    ".pdf",
    ".pptx", ".ppt", ".pptm",
    ".md",
}

# Magic byte signatures
XLSX_MAGIC = b"PK"
OLE2_MAGIC = b"\xd0\xcf\x11\xe0"
PDF_MAGIC = b"%PDF-"

# OOXML extensions that should be ZIP archives (not OLE2) when unencrypted
_OOXML_EXTENSIONS = {".docx", ".docm", ".xlsx", ".xlsm", ".pptx", ".pptm"}


class PasswordProtectedError(Exception):
    """Raised when a document is password-protected."""

    pass


def validate_extension(filename: str) -> str:
    """Check file extension against whitelist.

    Returns:
        The lowercase extension.

    Raises:
        ValueError: If extension is not supported.
    """
    ext = Path(filename).suffix.lower()
    if ext not in VALID_EXTENSIONS:
        raise ValueError(f"Unsupported extension: {ext}")
    return ext


def validate_magic_bytes(content: bytes) -> None:
    """Check file magic bytes match expected document formats.

    Raises:
        ValueError: If signature doesn't match any known format.
    """
    if content[:2] == XLSX_MAGIC or content[:4] == OLE2_MAGIC or content[:5] == PDF_MAGIC:
        return
    raise ValueError("File signature does not match expected format")


def validate_markdown_utf8(content: bytes) -> None:
    """Check markdown content is valid UTF-8.

    Raises:
        ValueError: If bytes are not decodable as UTF-8.
    """
    try:
        content.decode("utf-8")
    except UnicodeDecodeError:
        raise ValueError("File is not valid UTF-8 markdown")


def check_password_protected(file_path: str) -> None:
    """Check if an Office document is password-protected.

    Raises:
        PasswordProtectedError: If the document is encrypted.
    """
    ext = Path(file_path).suffix.lower()

    if ext in _OOXML_EXTENSIONS:
        with open(file_path, "rb") as f:
            header = f.read(8)
        if header[:4] == OLE2_MAGIC:
            raise PasswordProtectedError("Document is password-protected")
    elif ext in {".doc", ".xls", ".ppt"}:
        try:
            import olefile
        except ImportError:
            return
        if not olefile.isOleFile(file_path):
            return
        try:
            ole = olefile.OleFileIO(file_path)
            if ole.exists("EncryptionInfo") or ole.exists("encryption"):
                ole.close()
                raise PasswordProtectedError("Document is password-protected")
            ole.close()
        except PasswordProtectedError:
            raise
        except Exception:
            pass


def sanitize_filename(filename: str) -> str:
    """Strip path traversal while preserving CJK/Unicode characters."""
    safe = os.path.basename(filename)
    safe = safe.replace("\x00", "")
    safe = safe.replace("..", "")
    safe = safe.strip(". ")
    return safe