#!/usr/bin/env python3
"""Накладывать готовые VFX-элементы: чёрный фон, зелёный фон или альфа.

Стоки вроде Pexels и Pixabay — библиотеки **съёмки**, а не элементов. По запросу
«arrow overlay» они честно отдают видео, где кто-то снял улицу. Настоящие
элементы — указатели, глитч-рамки, вспышки, частицы, дым — лежат на профильных
VFX-сайтах, и API у них нет: файл скачивается руками.

Поэтому этот скрипт не ищет. Он делает так, чтобы **любой** скачанный элемент
работал без ручной возни: определяет, как он сделан, и накладывает правильно.

Три способа, которыми раздают элементы, и что с каждым делать:

* **на чёрном фоне** — самый частый. Складывается режимом `screen`: чёрное
  становится прозрачным само по себе, потому что screen оставляет максимум из
  двух пикселей. Хромакей здесь навредит: он вырежет ещё и тёмные части самого
  эффекта.
* **на зелёном фоне** — вырезается `chromakey` по измеренному цвету фона, а не
  по константе «зелёный». Оттенок гуляет от сайта к сайту, и захардкоженный
  0x00FF00 оставляет кайму на половине паков.
* **с альфа-каналом** (ProRes 4444, WebM, MOV) — просто накладывается.

Способ определяется по самому файлу: наличие альфы читается из pix_fmt, а фон —
по цвету рамки кадра. Пользователю не нужно ничего указывать; если определение
ошиблось, есть `--mode`.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))

from media_probe import probe  # noqa: E402
from skill_config import canvas as resolve_canvas, emit, require_binary, utf8_stdout  # noqa: E402

ALPHA_FORMATS = ("yuva", "rgba", "argb", "abgr", "bgra", "ya")


def border_colour(path: Path) -> tuple[float, float, float, float]:
    """Средний цвет рамки кадра и насколько он однороден.

    Рамка, а не весь кадр: эффект живёт в середине, а фон — по краям. Судить по
    среднему всего кадра значит смешать фон с самим эффектом и получить серое
    ни о чём.
    """
    done = subprocess.run(
        ["ffmpeg", "-v", "error", "-nostdin", "-i", str(path),
         "-vf", "select='between(n\\,10\\,40)',scale=64:64,format=rgb24",
         "-frames:v", "8", "-f", "rawvideo", "-"],
        capture_output=True)
    if done.returncode or not done.stdout:
        return 0.0, 0.0, 0.0, 999.0
    data = np.frombuffer(done.stdout, dtype=np.uint8)
    frames = data.size // (64 * 64 * 3)
    if frames < 1:
        return 0.0, 0.0, 0.0, 999.0
    stack = data[:frames * 64 * 64 * 3].reshape(frames, 64, 64, 3).astype(np.float64)
    edge = np.concatenate([
        stack[:, :4, :, :].reshape(-1, 3), stack[:, -4:, :, :].reshape(-1, 3),
        stack[:, :, :4, :].reshape(-1, 3), stack[:, :, -4:, :].reshape(-1, 3),
    ])
    mean = edge.mean(axis=0)
    spread = float(edge.std(axis=0).mean())
    return float(mean[0]), float(mean[1]), float(mean[2]), spread


def detect_mode(path: Path) -> tuple[str, dict]:
    info = probe(path)
    pix = ((info.get("video") or {}).get("pix_fmt") or "").lower()
    if any(pix.startswith(prefix) for prefix in ALPHA_FORMATS):
        return "alpha", {"pix_fmt": pix, "why": "в потоке есть альфа-канал"}
    r, g, b, spread = border_colour(path)
    facts = {"border_rgb": [round(r), round(g), round(b)], "border_spread": round(spread, 1)}
    if g > 90 and g > r * 1.6 and g > b * 1.6:
        return "green", {**facts, "why": "рамка кадра зелёная — хромакей по измеренному цвету"}
    if max(r, g, b) < 34 and spread < 26:
        return "black", {**facts, "why": "рамка кадра чёрная и ровная — режим screen"}
    return "unknown", {**facts,
                       "why": "фон не чёрный, не зелёный и без альфы — это обычная "
                              "съёмка, а не элемент. Как оверлей она не сложится."}


def build_chain(mode: str, facts: dict, width: int, height: int,
                scale: float, opacity: float, similarity: float, blend: float) -> tuple[str, str]:
    """Цепочка обработки элемента и способ смешивания с базой."""
    target = max(2, int(width * scale)) & ~1
    common = [f"scale={target}:-2:flags=bicubic"]
    if mode == "green":
        r, g, b = facts.get("border_rgb", [0, 255, 0])
        key = f"0x{int(r):02X}{int(g):02X}{int(b):02X}"
        common += [
            f"chromakey={key}:{similarity}:{blend}",
            # despill: без него вокруг элемента остаётся зелёная кайма, и её
            # видно именно на светлом фоне, куда элемент обычно и кладут.
            "colorchannelmixer=gr=0.35:gb=0.35:gg=0.30",
            "format=rgba",
        ]
        if opacity < 0.999:
            common.append(f"colorchannelmixer=aa={opacity:.3f}")
        return ",".join(common), "overlay"
    if mode == "alpha":
        common.append("format=rgba")
        if opacity < 0.999:
            common.append(f"colorchannelmixer=aa={opacity:.3f}")
        return ",".join(common), "overlay"
    # black
    if opacity < 0.999:
        # У screen нет альфы, поэтому прозрачность делается яркостью: тусклее
        # элемент — слабее вклад.
        common.append(f"colorlevels=romax={opacity:.3f}:gomax={opacity:.3f}:bomax={opacity:.3f}")
    common.append("format=gbrp")
    return ",".join(common), "blend=all_mode=screen"


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--plan", required=True,
                        help="JSON {elements:[{at, path, duration?, source_in?, scale?, "
                             "opacity?, mode?, why?}]}")
    parser.add_argument("--canvas", default="reels")
    parser.add_argument("--similarity", type=float, default=0.18)
    parser.add_argument("--blend", type=float, default=0.08)
    parser.add_argument("--crf", type=int, default=17)
    parser.add_argument("--preset", default="slow")
    parser.add_argument("--inspect", help="Определить режим для каждого файла в папке и выйти")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    require_binary("ffmpeg")

    if args.inspect:
        folder = Path(args.inspect)
        found = []
        for item in sorted(folder.rglob("*")):
            if item.suffix.lower() not in (".mov", ".mp4", ".webm", ".mkv"):
                continue
            mode, facts = detect_mode(item)
            found.append({"file": str(item), "mode": mode, **facts})
        emit({"status": "ok", "inspected": len(found), "elements": found,
              "note": "unknown означает обычную съёмку: как элемент она не подходит."})
        return

    source = Path(args.input).resolve()
    if not source.is_file():
        raise SystemExit(f"Not found: {source}")
    board = resolve_canvas(args.canvas)
    info = probe(source)
    video = info.get("video") or {}
    width = int(video.get("display_width") or board.width)
    height = int(video.get("display_height") or board.height)
    fps = int(round(float(video.get("avg_frame_rate") or board.fps)))
    duration = float(video.get("duration") or info.get("duration") or 0.0)

    payload = json.loads(Path(args.plan).read_text(encoding="utf-8"))
    elements = payload.get("elements") or payload
    prepared: list[dict] = []
    skipped: list[dict] = []
    for item in sorted(elements, key=lambda e: float(e["at"])):
        path = Path(item["path"]).resolve()
        if not path.is_file():
            skipped.append({**item, "why": "файл не найден"})
            continue
        mode = item.get("mode")
        facts: dict = {}
        if not mode:
            mode, facts = detect_mode(path)
        if mode == "unknown":
            skipped.append({**item, "mode": mode, **facts})
            continue
        at = float(item["at"])
        if at < 0.2 or at > duration - 0.4:
            skipped.append({**item, "why": "вне длительности базы"})
            continue
        prepared.append({
            "at": at, "path": path, "mode": mode, "facts": facts,
            "length": float(item.get("duration", 1.6)),
            "source_in": float(item.get("source_in", 0.0)),
            "scale": float(item.get("scale", 1.0)),
            "opacity": float(item.get("opacity", 0.85)),
            "why": item.get("why"),
        })

    if not prepared:
        emit({"status": "nothing_to_do", "skipped": skipped})
        return

    command = ["ffmpeg", "-y", "-v", "warning", "-nostdin", "-i", str(source)]
    for item in prepared:
        command += ["-ss", f"{item['source_in']:.3f}", "-t", f"{item['length']:.3f}",
                    "-i", str(item["path"])]

    graph: list[str] = []
    label = "0:v"
    for index, item in enumerate(prepared, start=1):
        chain, how = build_chain(item["mode"], item["facts"], width, height,
                                 item["scale"], item["opacity"],
                                 args.similarity, args.blend)
        # `screen` через blend требует одинакового размера, поэтому элемент
        # всегда доводится до канваса; при overlay он может быть меньше.
        if how.startswith("blend"):
            chain += f",scale={width}:{height},setsar=1"
        chain += f",setpts=PTS-STARTPTS+{item['at']:.4f}/TB"
        graph.append(f"[{index}:v]{chain}[e{index}]")
        out = f"v{index}"
        gate = (f":enable='between(t,{item['at']:.4f},"
                f"{item['at'] + item['length']:.4f})'")
        if how == "overlay":
            graph.append(f"[{label}][e{index}]overlay=(W-w)/2:(H-h)/2:"
                         f"eof_action=pass{gate}[{out}]")
        else:
            graph.append(f"[{label}][e{index}]{how}:eof_action=pass{gate}[{out}]")
        label = out
    graph.append(f"[{label}]format=yuv420p[vout]")

    report = {
        "status": "planned",
        "input": str(source),
        "elements": [{"at": round(i["at"], 3), "file": i["path"].name, "mode": i["mode"],
                      "how": {"black": "screen", "green": "хромакей", "alpha": "альфа"}[i["mode"]],
                      "detected_by": i["facts"].get("why"), "why": i["why"]}
                     for i in prepared],
        "skipped": skipped,
        "count": len(prepared),
    }
    if args.dry_run:
        emit(report)
        return

    output = Path(args.output).resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    command += [
        "-filter_complex", ";".join(graph),
        "-map", "[vout]", "-map", "0:a?",
        "-c:v", "libx264", "-preset", args.preset, "-crf", str(args.crf),
        "-profile:v", "high", "-pix_fmt", "yuv420p", "-r", str(fps),
        "-c:a", "copy", "-movflags", "+faststart", str(output),
    ]
    done = subprocess.run(command, capture_output=True, text=True, errors="replace")
    if done.returncode:
        report.update({"status": "failed", "graph": ";".join(graph),
                       "tail": (done.stderr or "").strip()[-1800:]})
        emit(report)
        raise SystemExit(done.returncode)

    actual = probe(output)
    report["status"] = "ok"
    report["output"] = str(output)
    drift = (round(actual["duration"] - duration, 4) if actual.get("duration") else None)
    report["duration_drift"] = drift
    report["drift_ok"] = drift is not None and abs(drift) <= max(0.08, 2.0 / fps)
    emit(report)
    if not report["drift_ok"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()