watermark.py

"""Brand watermark overlay for generated images and videos.

A watermark is requested by passing a JSON spec (a string or an already-parsed
dict) to a generation endpoint via the ``watermark`` parameter. When the spec is
valid, a small badge — the brand logo with the site domain rendered beneath it —
is composited onto the generated image (or every frame of the generated video,
via a single ffmpeg overlay pass). An empty or invalid spec is a no-op, so any
caller that omits the parameter is completely unaffected.

Spec fields (JSON object; at least one of ``logo`` / ``domain`` must be present
for the spec to count as "valid"):

    logo      str    logo image as an http(s) URL, a ``data:`` URI, or raw
                     base64. Optional — a domain-only watermark is allowed.
    domain    str    text drawn under the logo, e.g. "nsfwailab.com". Optional.
    position  str    bottom-right (default), bottom-left, top-right, top-left,
                     bottom-center or top-center.
    scale     float  badge width as a fraction of the media width (default 0.18,
                     clamped to 0.05..0.6).
    opacity   float  overall badge opacity 0..1 (default 0.9).
    margin    float  badge margin from the edge as a fraction of media width
                     (default 0.03).

Every public entry point is wrapped so a watermarking failure degrades to
"asset returned un-watermarked" rather than breaking generation.
"""
from __future__ import annotations

import base64
import json
import os
import subprocess
import tempfile
import urllib.request

_VIDEO_REF_W = 1024  # fallback media width when a video's size can't be probed.

# Cache decoded logos so a per-frame / per-call render doesn't refetch or
# re-decode the same image. Keyed on the raw logo string.
_logo_cache: dict[str, object] = {}

_POSITIONS = {
    "bottom-right", "bottom-left", "top-right", "top-left",
    "bottom-center", "top-center",
}


def parse_spec(spec) -> dict | None:
    """Validate/normalise a watermark spec; return a clean dict or None.

    Accepts a JSON string or a dict. Returns None (a no-op signal) when the spec
    is empty, unparseable, or carries neither a logo nor a domain.
    """
    if not spec:
        return None
    data = spec
    if isinstance(spec, str):
        s = spec.strip()
        if not s:
            return None
        try:
            data = json.loads(s)
        except (ValueError, TypeError):
            return None
    if not isinstance(data, dict):
        return None

    logo = data.get("logo") or data.get("logo_url") or ""
    domain = (data.get("domain") or data.get("text") or "").strip()
    if not (logo or domain):
        return None

    def _f(key, default, lo, hi):
        try:
            return max(lo, min(hi, float(data.get(key, default))))
        except (TypeError, ValueError):
            return default

    position = str(data.get("position", "bottom-right")).strip().lower()
    if position not in _POSITIONS:
        position = "bottom-right"

    return {
        "logo": logo if isinstance(logo, str) else "",
        "domain": domain,
        "position": position,
        "scale": _f("scale", 0.18, 0.05, 0.6),
        "opacity": _f("opacity", 0.9, 0.1, 1.0),
        "margin": _f("margin", 0.03, 0.0, 0.2),
    }


def is_valid(spec) -> bool:
    """True when ``spec`` would produce a watermark."""
    return parse_spec(spec) is not None


def _load_logo(logo: str):
    """Decode a logo string (URL / data-URI / base64) into an RGBA PIL image."""
    if not logo:
        return None
    if logo in _logo_cache:
        return _logo_cache[logo]
    from PIL import Image
    import io

    img = None
    try:
        raw = None
        if logo.startswith("data:"):
            _, _, b64 = logo.partition(",")
            raw = base64.b64decode(b64)
        elif logo.startswith("http://") or logo.startswith("https://"):
            req = urllib.request.Request(logo, headers={"User-Agent": "watermark/1.0"})
            with urllib.request.urlopen(req, timeout=10) as resp:
                raw = resp.read()
        else:
            raw = base64.b64decode(logo)
        if raw:
            img = Image.open(io.BytesIO(raw)).convert("RGBA")
    except Exception as exc:  # noqa: BLE001 - logo is best-effort
        print(f"[watermark] logo decode failed: {exc}")
        img = None
    _logo_cache[logo] = img
    return img


_FONT_CANDIDATES = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    "DejaVuSans-Bold.ttf",
]


def _font(size: int):
    from PIL import ImageFont

    for path in _FONT_CANDIDATES:
        try:
            return ImageFont.truetype(path, size)
        except Exception:  # noqa: BLE001
            continue
    return ImageFont.load_default()


def _render_badge(cfg: dict, media_w: int):
    """Build the watermark badge (logo + domain on a translucent plate)."""
    from PIL import Image, ImageDraw

    badge_w = max(72, int(media_w * cfg["scale"]))
    pad = max(6, badge_w // 16)

    logo = _load_logo(cfg["logo"])
    logo_img = None
    logo_h = 0
    if logo is not None and logo.width > 0:
        lw = max(1, badge_w - 2 * pad)
        lh = max(1, int(logo.height * lw / logo.width))
        logo_img = logo.resize((lw, lh), Image.LANCZOS)
        logo_h = lh

    domain = cfg["domain"]
    font = None
    text_w = text_h = 0
    bbox = (0, 0, 0, 0)
    if domain:
        measure = ImageDraw.Draw(Image.new("RGBA", (1, 1)))
        fs = max(11, int(badge_w * 0.14))
        font = _font(fs)
        bbox = measure.textbbox((0, 0), domain, font=font)
        text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
        max_tw = badge_w - 2 * pad
        if text_w > max_tw and text_w > 0:
            fs = max(8, int(fs * max_tw / text_w))
            font = _font(fs)
            bbox = measure.textbbox((0, 0), domain, font=font)
            text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]

    gap = pad if (logo_img is not None and domain) else 0
    badge_h = logo_h + gap + text_h + 2 * pad

    badge = Image.new("RGBA", (badge_w, badge_h), (0, 0, 0, 0))
    draw = ImageDraw.Draw(badge)
    radius = max(8, badge_w // 12)
    draw.rounded_rectangle([0, 0, badge_w - 1, badge_h - 1], radius=radius,
                           fill=(0, 0, 0, 110))

    y = pad
    if logo_img is not None:
        badge.alpha_composite(logo_img, ((badge_w - logo_img.width) // 2, y))
        y += logo_h + gap
    if domain:
        tx = (badge_w - text_w) // 2 - bbox[0]
        draw.text((tx, y - bbox[1]), domain, font=font, fill=(255, 255, 255, 235))

    opacity = cfg["opacity"]
    if opacity < 1.0:
        alpha = badge.split()[3].point(lambda v: int(v * opacity))
        badge.putalpha(alpha)
    return badge


def _offset(position: str, media_w: int, media_h: int, bw: int, bh: int, margin: int):
    """Top-left pixel offset for the badge given a named position."""
    if "right" in position:
        x = media_w - bw - margin
    elif "left" in position:
        x = margin
    else:  # center
        x = (media_w - bw) // 2
    y = margin if position.startswith("top") else media_h - bh - margin
    return max(0, x), max(0, y)


def apply_to_image(img, spec):
    """Return ``img`` with the watermark composited on, or unchanged on no-op."""
    cfg = parse_spec(spec)
    if cfg is None:
        return img
    try:
        base = img.convert("RGBA")
        badge = _render_badge(cfg, base.width)
        margin = int(cfg["margin"] * base.width)
        x, y = _offset(cfg["position"], base.width, base.height,
                       badge.width, badge.height, margin)
        base.alpha_composite(badge, (x, y))
        return base.convert("RGB")
    except Exception as exc:  # noqa: BLE001 - never break generation
        print(f"[watermark] image overlay failed: {exc}")
        return img


def _video_dims(path: str):
    try:
        import cv2

        cap = cv2.VideoCapture(path)
        w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        cap.release()
        if w > 0 and h > 0:
            return w, h
    except Exception:  # noqa: BLE001
        pass
    # Fallback when OpenCV is unavailable: parse the resolution out of ffmpeg's
    # probe output so the badge is still sized/placed against the real frame.
    try:
        import re

        import imageio_ffmpeg

        ff = imageio_ffmpeg.get_ffmpeg_exe()
        err = subprocess.run([ff, "-i", path], capture_output=True, text=True).stderr
        m = re.search(r"Video:.*?(\d{2,5})x(\d{2,5})", err)
        if m:
            return int(m.group(1)), int(m.group(2))
    except Exception:  # noqa: BLE001
        pass
    return None


def apply_to_video(in_path: str, spec, out_path: str | None = None) -> str:
    """Overlay the watermark on a video via one ffmpeg pass.

    Returns the path to the watermarked file, or the original ``in_path`` when
    the spec is a no-op or anything fails (so upload still proceeds).
    """
    cfg = parse_spec(spec)
    if cfg is None:
        return in_path
    badge_png = None
    try:
        import imageio_ffmpeg

        ffmpeg = imageio_ffmpeg.get_ffmpeg_exe()
        dims = _video_dims(in_path)
        media_w, media_h = dims if dims else (_VIDEO_REF_W, _VIDEO_REF_W)

        badge = _render_badge(cfg, media_w)
        margin = int(cfg["margin"] * media_w)
        x, y = _offset(cfg["position"], media_w, media_h,
                       badge.width, badge.height, margin)

        fd, badge_png = tempfile.mkstemp(suffix=".png")
        os.close(fd)
        badge.save(badge_png)

        if out_path is None:
            fd, out_path = tempfile.mkstemp(suffix=".mp4")
            os.close(fd)

        cmd = [
            ffmpeg, "-y", "-i", in_path, "-i", badge_png,
            "-filter_complex", f"overlay={x}:{y}:format=auto",
            "-c:a", "copy", "-c:v", "libx264", "-pix_fmt", "yuv420p",
            "-crf", "18", "-preset", "veryfast", "-movflags", "+faststart",
            out_path,
        ]
        subprocess.run(cmd, check=True, capture_output=True)
        return out_path
    except Exception as exc:  # noqa: BLE001 - never break generation
        detail = exc.stderr.decode("utf-8", "ignore")[-400:] if isinstance(exc, subprocess.CalledProcessError) else exc
        print(f"[watermark] video overlay failed: {detail}")
        return in_path
    finally:
        if badge_png and os.path.exists(badge_png):
            try:
                os.unlink(badge_png)
            except OSError:
                pass