"""Remove burnt-in captions by inpainting the glyphs, not the box around them.

ffmpeg's delogo can only patch a rectangle, interpolating from its border. Over
a caption sitting on solar panels that destroys every pixel inside the box and
leaves vertical smears. Inpainting instead touches only the glyph pixels and
reconstructs them from surrounding texture, so the panels survive.

Getting the glyph mask is the trick. Within the stretch where a caption is
held, the text is pinned to the same pixels while the camera moves the scene
underneath. A per-pixel MINIMUM across that stretch therefore keeps the bright
text and drags the moving background toward black — threshold that and the
glyphs fall out.

Usage:
    python scripts/remove_captions.py --preview   # masks + 4 sample frames
    python scripts/remove_captions.py             # full run
"""

from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

import cv2
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
IN_DIR = ROOT / "public" / "services-frame"
OUT_DIR = ROOT / "public" / "frames-services"
PREVIEW_DIR = ROOT / "public" / "_caption_preview"

SRC = "ezgif-frame-{:03d}.jpg"
TOTAL = 240

# Each caption: the frames it is patched over, the core window used to build
# its mask, and the band it may occupy as (x0, y0, x1, y1).
#
# The core matters more than it looks. These captions animate in word by word,
# so a window that starts too early has frames where the later words are absent
# — the minimum then washes those letters out and the mask ends up with holes
# exactly where text still is. Each core sits inside the settled stretch.
CAPTIONS = [
    {
        "name": "power-from-the-sun",
        "apply": (1, 82),
        "core": (26, 48),
        "band": (420, 410, 1660, 600),
    },
    {
        "name": "seamless-installation",
        "apply": (78, 155),
        "core": (114, 142),
        "band": (70, 125, 1420, 290),
    },
    {
        "name": "clean-energy-flow",
        "apply": (148, 240),
        "core": (205, 238),
        "band": (70, 125, 1420, 290),
    },
]

# The "Veo" wordmark is constant for the whole clip, so it gets its own mask
# built the same way but over every frame.
VEO_BAND = (1830, 1005, 1915, 1070)

# Fraction of the min-composite's dynamic range above which a pixel counts as
# text. High, because the scene's glowing cables are bright too.
# Tuned by sweeping both against a frame where the caption sits on open sky.
# A tight mask leaves a legible ghost, because inpainting then samples from the
# glyph's own soft outer glow and fills text colour back in. Going wide costs a
# little more reconstructed area and removes the text completely.
THRESHOLD = 0.22
DILATE = 41         # covers the glyph plus the halo bleeding off its edges
INPAINT_RADIUS = 15


def read_gray(path: Path) -> np.ndarray:
    img = cv2.imread(str(path), cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise SystemExit("Cannot read {}".format(path))
    return img


def build_mask(a: int, b: int, band: tuple[int, int, int, int]) -> np.ndarray:
    """Glyph mask from the per-pixel minimum over frames a..b, inside `band`."""
    acc = None
    for n in range(a, b + 1):
        g = read_gray(IN_DIR / SRC.format(n))
        acc = g if acc is None else np.minimum(acc, g)

    x0, y0, x1, y1 = band
    region = acc[y0:y1, x0:x1]
    lo, hi = int(region.min()), int(region.max())
    cut = lo + (hi - lo) * THRESHOLD

    mask = np.zeros(acc.shape, np.uint8)
    mask[y0:y1, x0:x1] = (region > cut).astype(np.uint8) * 255
    if DILATE:
        k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (DILATE, DILATE))
        mask = cv2.dilate(mask, k)
    return mask


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--preview", action="store_true")
    ap.add_argument("--quality", type=int, default=92, help="JPEG quality")
    args = ap.parse_args()

    if not IN_DIR.is_dir():
        print("Input folder not found: {}".format(IN_DIR), file=sys.stderr)
        return 1

    print("Building masks…")
    masks = []
    for c in CAPTIONS:
        m = build_mask(*c["core"], c["band"])
        pct = 100.0 * (m > 0).sum() / m.size
        print("  {:24} frames {:>3}-{:<3}  {:.2f}% of frame masked".format(
            c["name"], c["apply"][0], c["apply"][1], pct))
        masks.append(m)
    veo = build_mask(1, TOTAL, VEO_BAND)
    print("  {:24} all frames        {:.3f}% of frame masked".format("veo-watermark",
                                                                    100.0 * (veo > 0).sum() / veo.size))

    if args.preview:
        PREVIEW_DIR.mkdir(parents=True, exist_ok=True)
        for c, m in zip(CAPTIONS, masks):
            cv2.imwrite(str(PREVIEW_DIR / "mask-{}.png".format(c["name"])), m)
        cv2.imwrite(str(PREVIEW_DIR / "mask-veo.png"), veo)
        for n in (20, 115, 200, 236):
            img = cv2.imread(str(IN_DIR / SRC.format(n)))
            combined = veo.copy()
            for c, m in zip(CAPTIONS, masks):
                if c["apply"][0] <= n <= c["apply"][1]:
                    combined = np.maximum(combined, m)
            out = cv2.inpaint(img, combined, INPAINT_RADIUS, cv2.INPAINT_TELEA)
            cv2.imwrite(str(PREVIEW_DIR / "sample-{:03d}.jpg".format(n)), out,
                        [cv2.IMWRITE_JPEG_QUALITY, args.quality])
        print("Preview written to public/_caption_preview/")
        return 0

    if OUT_DIR.exists():
        shutil.rmtree(OUT_DIR)
    OUT_DIR.mkdir(parents=True)

    print("Inpainting {} frames…".format(TOTAL))
    for n in range(1, TOTAL + 1):
        img = cv2.imread(str(IN_DIR / SRC.format(n)))
        combined = veo.copy()
        for c, m in zip(CAPTIONS, masks):
            if c["apply"][0] <= n <= c["apply"][1]:
                combined = np.maximum(combined, m)
        out = cv2.inpaint(img, combined, INPAINT_RADIUS, cv2.INPAINT_TELEA)
        cv2.imwrite(str(OUT_DIR / "frame_{:04d}.jpg".format(n)), out,
                    [cv2.IMWRITE_JPEG_QUALITY, args.quality])
        if n % 40 == 0:
            print("  {}/{}".format(n, TOTAL))

    written = sorted(OUT_DIR.glob("frame_*.jpg"))
    mb = sum(f.stat().st_size for f in written) / (1024 * 1024)
    print("Wrote   : {} frames, {:.1f} MB".format(len(written), mb))
    if len(written) != TOTAL:
        print("Expected {} frames".format(TOTAL), file=sys.stderr)
        return 1
    print("FRAME_COUNT = {}".format(len(written)))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())