"""Shared archive extraction utilities with ZIP bomb and size protection.

All subtitle providers that handle archive downloads should use these helpers
instead of inline zipfile/rarfile calls, to ensure consistent security checks
across the entire download pipeline.
"""

import contextlib
import hashlib
import io
import logging
import os
import zipfile

from werkzeug.utils import secure_filename

logger = logging.getLogger(__name__)

_MAX_ARCHIVE_BYTES = 20 * 1024 * 1024  # 20 MB — reject before extraction
_MAX_EXTRACTED_BYTES = 50 * 1024 * 1024  # 50 MB — total extracted
_MAX_COMPRESSION_RATIO = 100  # 100:1 — ZIP bomb ratio
_SUBTITLE_EXTENSIONS = frozenset({".ass", ".srt", ".ssa", ".vtt"})


def _safe_basename(path: str) -> str:
    """Return a filename-safe basename (P2).

    ``os.path.basename`` is platform-dependent: on Linux it does not strip
    backslashes, so a ZIP entry named ``..\\..\\cmd.exe.srt`` would slip
    through on non-Windows hosts. ``secure_filename`` normalises to ASCII,
    strips traversal components, and blocks control chars / nul / slashes
    on every platform.
    """
    raw = os.path.basename(path.replace("\\", "/"))
    safe = secure_filename(raw)
    stem, ext = os.path.splitext(raw)
    if ext and not safe.lower().endswith(ext.lower()) and secure_filename(ext.lstrip(".")):
        # Every letter of the name was non-ASCII ("简体.srt" -> "srt"): keep the
        # extension, or the subtitle check drops the member, and make the name
        # unique, or two members collapse into one.
        digest = hashlib.sha1(
            raw.encode("utf-8", "surrogateescape"), usedforsecurity=False
        ).hexdigest()[:8]
        safe_stem = secure_filename(stem) or "subtitle"
        return f"{safe_stem}-{digest}{ext.lower()}"
    return safe


def _plain_basename(path: str) -> str:
    """Basename that keeps non-Latin letters, still safe as a label.

    ``secure_filename`` reduces "简体.srt" to "srt" and "Ταινία.srt" to "srt":
    the extension check then drops the file, and two members can collapse to
    one name. Callers that only match or display names (never write them as a
    path) use this instead: directories, control characters and leading dots
    are removed, letters of every script stay.
    """
    raw = os.path.basename(path.replace("\\", "/"))
    cleaned = "".join(ch for ch in raw if ch.isprintable()).strip().lstrip(".")
    return cleaned


def extract_subtitles_from_zip(
    data: bytes,
    subtitle_exts: frozenset[str] = _SUBTITLE_EXTENSIONS,
    raw_names: bool = False,
) -> list[tuple[str, bytes]]:
    """In-memory ZIP extraction with ZIP bomb and size protection.

    Args:
        data: Raw ZIP archive bytes.
        subtitle_exts: Allowed subtitle file extensions (with leading dot).
        raw_names: Return names via ``_plain_basename`` (non-Latin letters kept)
            instead of ``secure_filename``. Only for callers that never use the
            name as a filesystem path.

    Returns:
        List of (basename, content) tuples for matching subtitle files.

    Raises:
        ValueError: If archive exceeds size limits or compression ratio.
    """
    if len(data) > _MAX_ARCHIVE_BYTES:
        raise ValueError(
            f"Archive too large: {len(data) // (1024 * 1024)} MB > "
            f"{_MAX_ARCHIVE_BYTES // (1024 * 1024)} MB limit"
        )

    try:
        with zipfile.ZipFile(io.BytesIO(data)) as zf:
            total_compressed = sum(info.compress_size for info in zf.infolist())
            total_uncompressed = sum(info.file_size for info in zf.infolist())

            if total_uncompressed > _MAX_EXTRACTED_BYTES:
                raise ValueError(
                    f"Archive extraction too large: {total_uncompressed // (1024 * 1024)} MB > "
                    f"{_MAX_EXTRACTED_BYTES // (1024 * 1024)} MB limit"
                )

            # ZIP bomb ratio check (skip if nothing compressed — empty/stored zips)
            if total_compressed > 0:
                ratio = total_uncompressed / total_compressed
                if ratio > _MAX_COMPRESSION_RATIO:
                    raise ValueError(
                        f"ZIP bomb detected: compression ratio {ratio:.0f}:1 exceeds "
                        f"{_MAX_COMPRESSION_RATIO}:1 limit"
                    )

            results = []
            for info in zf.infolist():
                # Skip directory entries
                if info.filename.endswith("/"):
                    continue
                # Strip any path components + sanitize for disk write (ZIP slip + P2)
                basename = (_plain_basename if raw_names else _safe_basename)(info.filename)
                if not basename:
                    continue
                ext = os.path.splitext(basename)[1].lower()
                if ext not in subtitle_exts:
                    continue
                results.append((basename, zf.read(info.filename)))

            return results

    except zipfile.BadZipFile:
        logger.warning("archive_utils: bad ZIP archive, skipping")
        return []


def safe_read_zip_member(
    zf: zipfile.ZipFile,
    name: str,
    max_bytes: int = _MAX_EXTRACTED_BYTES,
) -> bytes:
    """Read a single named entry from an open ZipFile with bomb protection.

    Use when callers know which entries they want (e.g. ``config.json``)
    and need to avoid loading a multi-GB decompressed payload into memory.
    The caller has already validated the outer archive size; this guards
    each individual entry against a single oversized member.

    Args:
        zf: Open ``zipfile.ZipFile`` instance.
        name: Entry name to read.
        max_bytes: Cap on the entry's uncompressed size (default 50 MB,
            shared with ``extract_subtitles_from_zip``).

    Returns:
        Decompressed entry bytes.

    Raises:
        KeyError: If ``name`` is not in the archive.
        ValueError: If the uncompressed size or compression ratio
            exceeds the configured limits.
    """
    info = zf.getinfo(name)  # raises KeyError if missing

    # The central-directory header fields are attacker-controlled: a crafted
    # archive can declare file_size=0 while storing a large DEFLATE payload,
    # making both header checks pass and then expanding to GBs inside zf.read().
    # Treat the headers only as a cheap fast-reject, and enforce the real cap by
    # streaming the decompressed bytes and stopping one byte past the limit.
    if info.file_size > max_bytes:
        raise ValueError(
            f"ZIP entry '{name}' too large: {info.file_size // (1024 * 1024)} MB > "
            f"{max_bytes // (1024 * 1024)} MB limit"
        )

    if info.compress_size > 0:
        ratio = info.file_size / info.compress_size
        if ratio > _MAX_COMPRESSION_RATIO:
            raise ValueError(
                f"ZIP bomb detected on '{name}': ratio {ratio:.0f}:1 exceeds "
                f"{_MAX_COMPRESSION_RATIO}:1 limit"
            )

    with zf.open(name) as member:
        data = member.read(max_bytes + 1)
    if len(data) > max_bytes:
        raise ValueError(
            f"ZIP entry '{name}' exceeds {max_bytes // (1024 * 1024)} MB cap during decompression"
        )
    return data


def safe_extract_zip_member_to(
    zf: zipfile.ZipFile,
    name: str,
    dest_path: str,
    max_bytes: int = _MAX_EXTRACTED_BYTES,
) -> int:
    """Stream one entry to ``dest_path`` with the same guards as the reader.

    The in-memory reader is fine for a config file, but a database inside a
    full backup can be gigabytes: materialising it (on top of the archive
    itself) is a memory-exhaustion risk. Returns the bytes written; the partial
    file is removed when a limit is hit.
    """
    info = zf.getinfo(name)  # raises KeyError if missing
    if info.file_size > max_bytes:
        raise ValueError(
            f"ZIP entry '{name}' too large: {info.file_size // (1024 * 1024)} MB > "
            f"{max_bytes // (1024 * 1024)} MB limit"
        )
    if info.compress_size > 0:
        ratio = info.file_size / info.compress_size
        if ratio > _MAX_COMPRESSION_RATIO:
            raise ValueError(
                f"ZIP bomb detected on '{name}': ratio {ratio:.0f}:1 exceeds "
                f"{_MAX_COMPRESSION_RATIO}:1 limit"
            )

    written = 0
    try:
        with zf.open(name) as member, open(dest_path, "wb") as out:
            while chunk := member.read(1024 * 1024):
                written += len(chunk)
                if written > max_bytes:
                    raise ValueError(
                        f"ZIP entry '{name}' exceeds {max_bytes // (1024 * 1024)} MB cap "
                        "during decompression"
                    )
                out.write(chunk)
    except BaseException:
        with contextlib.suppress(OSError):
            os.remove(dest_path)
        raise
    return written


def extract_subtitles_from_rar(
    data: bytes,
    subtitle_exts: frozenset[str] = _SUBTITLE_EXTENSIONS,
    raw_names: bool = False,
) -> list[tuple[str, bytes]]:
    """In-memory RAR extraction with size protection.

    Args:
        data: Raw RAR archive bytes.
        subtitle_exts: Allowed subtitle file extensions (with leading dot).

    Returns:
        List of (basename, content) tuples for matching subtitle files.

    Raises:
        ImportError: If the rarfile package is not installed.
        ValueError: If archive exceeds the size limit.
    """
    if len(data) > _MAX_ARCHIVE_BYTES:
        raise ValueError(
            f"Archive too large: {len(data) // (1024 * 1024)} MB > "
            f"{_MAX_ARCHIVE_BYTES // (1024 * 1024)} MB limit"
        )

    import rarfile  # optional dep — raises ImportError if absent

    results = []
    try:
        with rarfile.RarFile(io.BytesIO(data)) as rf:
            for info in rf.infolist():
                basename = (_plain_basename if raw_names else _safe_basename)(info.filename)
                if not basename:
                    continue
                ext = os.path.splitext(basename)[1].lower()
                if ext not in subtitle_exts:
                    continue
                results.append((basename, rf.read(info.filename)))
    except Exception as e:
        logger.warning("archive_utils: RAR extraction failed: %s", e)

    return results