adam/remote_media.py

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import mimetypes
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from adam.remote_api import RemoteApiError


MAX_SOURCE_BYTES = 60 * 1024 * 1024
MAX_THUMBNAIL_SIZE = 1200
DEFAULT_THUMBNAIL_SIZE = 320


class OpaqueIdCodec:
    def __init__(self, secret: str) -> None:
        self.secret = (secret or "adam-remote").encode("utf-8")

    def encode(self, payload: dict[str, Any]) -> str:
        body = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
        token = base64.urlsafe_b64encode(body).decode("ascii").rstrip("=")
        signature = hmac.new(self.secret, token.encode("ascii"), hashlib.sha256).hexdigest()[:24]
        return f"{token}.{signature}"

    def decode(self, value: str) -> dict[str, Any]:
        try:
            token, signature = value.rsplit(".", 1)
        except ValueError as exc:
            raise RemoteApiError("Unknown resource id.", status=404) from exc
        expected = hmac.new(self.secret, token.encode("ascii"), hashlib.sha256).hexdigest()[:24]
        if not hmac.compare_digest(signature, expected):
            raise RemoteApiError("Unknown resource id.", status=404)
        padding = "=" * (-len(token) % 4)
        try:
            payload = json.loads(base64.urlsafe_b64decode((token + padding).encode("ascii")).decode("utf-8"))
        except (ValueError, TypeError, json.JSONDecodeError) as exc:
            raise RemoteApiError("Unknown resource id.", status=404) from exc
        if not isinstance(payload, dict):
            raise RemoteApiError("Unknown resource id.", status=404)
        return payload


@dataclass(frozen=True, slots=True)
class RemoteMediaFile:
    path: Path
    content_type: str
    cache_hit: bool


class RemoteMediaStore:
    def __init__(self, root: Path, codec: OpaqueIdCodec) -> None:
        self.root = root.resolve()
        self.codec = codec
        self.cache_root = self.root / "data" / "remote_thumbnails"
        self.cache_root.mkdir(parents=True, exist_ok=True)

    def media_id(self, *, kind: str, asset_id: str, index: int) -> str:
        return self.codec.encode({"kind": kind, "asset_id": asset_id, "index": int(index)})

    def thumbnail(self, source: Path, *, size: int = DEFAULT_THUMBNAIL_SIZE) -> RemoteMediaFile:
        source = source.expanduser().resolve()
        if not source.is_file():
            raise RemoteApiError("Media file was not found.", status=404)
        try:
            stat = source.stat()
        except OSError as exc:
            raise RemoteApiError("Media file could not be read.", status=404) from exc
        if stat.st_size > MAX_SOURCE_BYTES:
            raise RemoteApiError("Media file is too large for remote preview.", status=413)
        bounded_size = max(64, min(int(size or DEFAULT_THUMBNAIL_SIZE), MAX_THUMBNAIL_SIZE))
        key = hashlib.sha256(
            f"{source}|{stat.st_mtime_ns}|{stat.st_size}|{bounded_size}".encode("utf-8", errors="ignore")
        ).hexdigest()
        target = self.cache_root / f"{key}.jpg"
        if target.is_file():
            return RemoteMediaFile(target, "image/jpeg", True)
        try:
            from PIL import Image, ImageOps

            with Image.open(source) as image:
                image = ImageOps.exif_transpose(image)
                image.thumbnail((bounded_size, bounded_size), Image.Resampling.LANCZOS)
                if image.mode not in {"RGB", "L"}:
                    image = image.convert("RGB")
                image.save(target, format="JPEG", quality=82, optimize=True)
        except Exception as exc:
            raise RemoteApiError("Thumbnail could not be generated.", status=415) from exc
        return RemoteMediaFile(target, "image/jpeg", False)

    def original(self, source: Path) -> RemoteMediaFile:
        source = source.expanduser().resolve()
        if not source.is_file():
            raise RemoteApiError("Media file was not found.", status=404)
        try:
            if source.stat().st_size > MAX_SOURCE_BYTES:
                raise RemoteApiError("Media file is too large for remote viewing.", status=413)
        except OSError as exc:
            raise RemoteApiError("Media file could not be read.", status=404) from exc
        return RemoteMediaFile(source, mimetypes.guess_type(str(source))[0] or "application/octet-stream", True)