epub2md/epub2md/media.py

"""Media registry: stable, collision-free output names; each source extracted once."""

from __future__ import annotations

import mimetypes
import posixpath

from .models import MediaFile, Package
from .utils import safe_zip_path, slugify

KIND_DIRS = {"images": "images", "audio": "audio", "video": "video", "other": "other"}


def media_kind(media_type: str, path: str) -> str:
    mt = media_type or (mimetypes.guess_type(path)[0] or "")
    if mt.startswith("image/"):
        return "images"
    if mt.startswith("audio/"):
        return "audio"
    if mt.startswith("video/"):
        return "video"
    return "other"


class MediaRegistry:
    def __init__(self, package: Package, strategy: str, slug_language: str) -> None:
        self.package = package
        self.strategy = strategy  # relative | flat | original-path
        self.slug_language = slug_language
        self.files: dict[str, MediaFile] = {}
        self._used: set[str] = set()
        self._counter = 0
        self.opf_dir = posixpath.dirname(package.opf_path)

    def _media_type(self, path: str) -> str:
        item = self.package.item_by_path(path)
        if item and item.media_type:
            return item.media_type
        return mimetypes.guess_type(path)[0] or "application/octet-stream"

    def _unique(self, candidate: str) -> str:
        if candidate not in self._used:
            self._used.add(candidate)
            return candidate
        stem, ext = posixpath.splitext(candidate)
        n = 2
        while f"{stem}-{n}{ext}" in self._used:
            n += 1
        result = f"{stem}-{n}{ext}"
        self._used.add(result)
        return result

    def _output_name(self, source: str, kind: str, hint: str = "") -> str:
        base = posixpath.basename(source)
        stem, ext = posixpath.splitext(base)
        ext = ext.lower()
        if self.strategy == "original-path":
            rel = source[len(self.opf_dir) + 1 :] if self.opf_dir and source.startswith(self.opf_dir + "/") else source
            rel = safe_zip_path(rel) or base
            return self._unique(f"media/{rel}")
        if self.strategy == "flat":
            return self._unique(f"media/{base}")
        self._counter += 1
        slug = slugify(hint or stem, self.slug_language, max_len=40, fallback="media")
        return self._unique(f"media/{KIND_DIRS[kind]}/{self._counter:04d}-{slug}{ext}")

    def register(self, source: str, role: str = "", hint: str = "") -> MediaFile:
        existing = self.files.get(source)
        if existing:
            existing.references += 1
            if role and not existing.role:
                existing.role = role
            return existing
        media_type = self._media_type(source)
        kind = media_kind(media_type, source)
        mf = MediaFile(
            source=source,
            output=self._output_name(source, kind, hint),
            media_type=media_type,
            kind=kind,
            references=1,
            role=role,
        )
        self.files[source] = mf
        return mf

    def register_generated(self, name: str, ext: str, data: bytes, media_type: str, role: str) -> MediaFile:
        source = f"generated:{name}{ext}"
        kind = media_kind(media_type, name + ext)
        if self.strategy == "relative":
            self._counter += 1
            out = self._unique(f"media/{KIND_DIRS[kind]}/{self._counter:04d}-{slugify(name, self.slug_language, 40, 'media')}{ext}")
        else:
            out = self._unique(f"media/generated/{slugify(name, self.slug_language, 40, 'media')}{ext}")
        mf = MediaFile(source=source, output=out, media_type=media_type, kind=kind, size=len(data),
                       references=1, role=role, data=data)
        self.files[source] = mf
        return mf