drive/naming.py

"""Filename and folder-name rules: sanitising, extension detection, collisions."""
from __future__ import annotations

import os
import re
import unicodedata
from pathlib import Path

from errors import PermanentError

VIDEO_EXTENSIONS = frozenset(
    {".mp4", ".m4v", ".mov", ".mkv", ".webm", ".avi", ".flv", ".wmv", ".ts", ".m2ts", ".mts", ".3gp", ".mpg", ".mpeg", ".ogv"}
)
# Extensions that share a container format; a title saying ".mov" for an MP4 file is fine.
_FAMILIES = (
    frozenset({".mp4", ".m4v", ".mov", ".3gp"}),
    frozenset({".mkv", ".webm"}),
    frozenset({".ts", ".m2ts", ".mts"}),
    frozenset({".mpg", ".mpeg"}),
)
MIME_TO_EXTENSION = {
    "video/mp4": ".mp4",
    "video/x-m4v": ".m4v",
    "video/quicktime": ".mov",
    "video/x-matroska": ".mkv",
    "video/webm": ".webm",
    "video/x-msvideo": ".avi",
    "video/x-flv": ".flv",
    "video/x-ms-wmv": ".wmv",
    "video/mp2t": ".ts",
    "video/3gpp": ".3gp",
    "video/mpeg": ".mpg",
    "video/ogg": ".ogv",
}
EXTENSION_TO_MIME = {ext: mime for mime, ext in MIME_TO_EXTENSION.items()} | {
    ".mpeg": "video/mpeg",
    ".m2ts": "video/mp2t",
    ".mts": "video/mp2t",
}
MAX_NAME_LENGTH = 200
_ALWAYS_REPLACE = re.compile(r"[/\\]")
_STRICT_REPLACE = re.compile(r'[<>:"|?*]')


def normalize_key(text: str) -> str:
    """Comparison key for folder names: NFC, collapsed whitespace, case-insensitive."""
    return " ".join(unicodedata.normalize("NFC", text).split()).casefold()


def clean_component(text: str | None, *, strict: bool = False) -> str:
    """Make a caption value safe as a Drive file/folder name while keeping Unicode.

    Control characters and path separators are replaced; letters, bullets, dots,
    parentheses and other punctuation are kept. ``strict`` additionally replaces
    characters Windows cannot store (useful with Google Drive for desktop).
    """
    if not text:
        return ""
    text = unicodedata.normalize("NFC", text)
    text = "".join(" " if unicodedata.category(ch) == "Cc" else ch for ch in text)
    text = _ALWAYS_REPLACE.sub("_", text)
    if strict:
        text = _STRICT_REPLACE.sub("_", text).rstrip(". ")
    text = " ".join(text.split())
    return truncate_name(text)


def split_video_extension(name: str) -> tuple[str, str]:
    """Split off a *video* extension only, so 'Lecture 1.5' keeps its '.5'."""
    stem, ext = os.path.splitext(name)
    if ext.lower() in VIDEO_EXTENSIONS and stem:
        return stem, ext
    return name, ""


def truncate_name(name: str, limit: int = MAX_NAME_LENGTH) -> str:
    if len(name) <= limit:
        return name
    stem, ext = split_video_extension(name)
    return stem[: limit - len(ext)].rstrip() + ext


def same_family(a: str, b: str) -> bool:
    a, b = a.lower(), b.lower()
    if a == b:
        return True
    return any(a in family and b in family for family in _FAMILIES)


def extension_from_mime(mime_type: str | None) -> str | None:
    return MIME_TO_EXTENSION.get((mime_type or "").lower())


def mime_for_name(name: str, fallback: str | None = None) -> str:
    _, ext = split_video_extension(name)
    return EXTENSION_TO_MIME.get(ext.lower()) or fallback or "application/octet-stream"


def detect_video_extension(path: Path) -> str | None:
    """Identify the container from magic bytes; ``None`` if unknown."""
    with open(path, "rb") as fh:
        head = fh.read(4096)
    if len(head) >= 12 and head[4:8] == b"ftyp":
        brand = head[8:12]
        if brand == b"qt  ":
            return ".mov"
        if brand.startswith(b"3g"):
            return ".3gp"
        if brand.startswith(b"M4V"):
            return ".m4v"
        return ".mp4"
    if head.startswith(b"\x1a\x45\xdf\xa3"):
        return ".webm" if b"webm" in head[:64] else ".mkv"
    if head[:4] == b"RIFF" and head[8:12] == b"AVI ":
        return ".avi"
    if head[:3] == b"FLV":
        return ".flv"
    if head[:8] == b"\x30\x26\xb2\x75\x8e\x66\xcf\x11":
        return ".wmv"
    if head[:4] in (b"\x00\x00\x01\xba", b"\x00\x00\x01\xb3"):
        return ".mpg"
    if head[:4] == b"OggS":
        return ".ogv"
    if len(head) > 376 and head[0] == 0x47 and head[188] == 0x47 and head[376] == 0x47:
        return ".ts"
    return None


def build_filename(
    title: str | None,
    *,
    message_id: int,
    original_filename: str | None,
    mime_type: str | None,
    detected_ext: str | None,
    missing_title_policy: str = "fallback",
    mismatch_policy: str = "keep",
    strict: bool = False,
) -> tuple[str, list[str]]:
    """Return the Drive filename for a video plus any warnings worth recording.

    The Title field is used verbatim (sanitised). The extension is only added when
    the title has none; a title whose extension disagrees with the real container
    is kept as-is unless EXTENSION_MISMATCH_POLICY=append.
    """
    warnings: list[str] = []
    name = clean_component(title, strict=strict)
    if not name:
        if missing_title_policy == "fail":
            raise PermanentError(
                f"message {message_id} has no usable Title", public="Caption has no usable Title (MISSING_TITLE_POLICY=fail)"
            )
        name = clean_component(original_filename, strict=strict)
        if name:
            warnings.append("Title missing; used the Telegram file name")
        else:
            name = f"telegram_{message_id}"
            warnings.append("Title missing; used a generated name")

    _, ext = split_video_extension(name)
    _, original_ext = split_video_extension(original_filename or "")
    actual = detected_ext or extension_from_mime(mime_type) or (original_ext.lower() or None)
    if not ext:
        if actual:
            name = truncate_name(name, MAX_NAME_LENGTH - len(actual)) + actual
        else:
            warnings.append("could not determine the file extension")
    elif detected_ext and not same_family(ext, detected_ext):
        warnings.append(f"title extension {ext} does not match detected container {detected_ext}")
        if mismatch_policy == "append":
            name = truncate_name(name, MAX_NAME_LENGTH - len(detected_ext)) + detected_ext
    return name, warnings


def add_suffix(name: str, suffix: str) -> str:
    """'Class 01.mp4' + '(2)' -> 'Class 01 (2).mp4'."""
    stem, ext = split_video_extension(name)
    if not ext:
        stem, ext = os.path.splitext(name)
        if len(ext) > 6 or not stem:
            stem, ext = name, ""
    return truncate_name(f"{stem} {suffix}{ext}")