#!/usr/bin/env python3
"""Encoder and decoder selection, probed on the real machine.

The presence of an encoder in `ffmpeg -encoders` means nothing. This build lists
`av1_nvenc` on a Pascal card that physically cannot encode AV1, and `h264_nvenc`
stays listed when the driver is too old for the API ffmpeg was built against. Both
fail only after the full render has been queued.

So every candidate is exercised with a real one-frame encode before it is offered,
and the result is cached with the driver version as part of the key — a driver
update must invalidate it.

Decode acceleration is probed separately from encode. `-hwaccel cuda` helps on long
timelines, but it is not free: any filter graph that needs frames in system memory
forces a download per frame, which on a 1070 Ti is slower than decoding on the CPU.
So it is only recommended for the cases that measure faster.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from skill_config import CACHE, emit, utf8_stdout  # noqa: E402

PROBE_CACHE = CACHE / "render_core_probe.json"
CACHE_TTL = 7 * 24 * 3600


# --------------------------------------------------------------------- profiles

@dataclass(frozen=True)
class Profile:
    id: str
    use: str
    # Constant-quality target per family. CRF and NVENC's CQ are not the same
    # scale: matching perceptual quality needs a slightly lower CQ number.
    x264_crf: int
    nvenc_cq: int
    x264_preset: str
    nvenc_preset: str


PROFILES: dict[str, Profile] = {
    "master": Profile("master", "Финальный мастер. Максимум качества, время не важно.",
                      16, 17, "slow", "p7"),
    "delivery": Profile("delivery", "Готовый ролик для платформы. Баланс.",
                        18, 19, "medium", "p6"),
    "draft": Profile("draft", "Черновик для проверки монтажа и читаемости.",
                     22, 25, "veryfast", "p4"),
    "proxy": Profile("proxy", "Прокси для монтажа длинного видео. Скорость важнее всего.",
                     26, 30, "ultrafast", "p1"),
}
DEFAULT_PROFILE = "delivery"


def driver_version() -> str | None:
    if shutil.which("nvidia-smi") is None:
        return None
    try:
        completed = subprocess.run(
            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=20, errors="replace",
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    if completed.returncode:
        return None
    return (completed.stdout or "").strip().splitlines()[0].strip() or None


def gpu_name() -> str | None:
    if shutil.which("nvidia-smi") is None:
        return None
    try:
        completed = subprocess.run(
            ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
            capture_output=True, text=True, timeout=20, errors="replace",
        )
    except (OSError, subprocess.TimeoutExpired):
        return None
    return (completed.stdout or "").strip().splitlines()[0].strip() if not completed.returncode else None


# ---------------------------------------------------------------------- probing

ENCODER_CANDIDATES: tuple[tuple[str, str, list[str]], ...] = (
    ("h264_nvenc", "h264", ["-preset", "p4", "-rc", "vbr", "-cq", "26", "-b:v", "0"]),
    ("hevc_nvenc", "hevc", ["-preset", "p4", "-rc", "vbr", "-cq", "28", "-b:v", "0"]),
    ("av1_nvenc", "av1", ["-preset", "p4", "-rc", "vbr", "-cq", "30", "-b:v", "0"]),
    ("h264_qsv", "h264", ["-global_quality", "26"]),
    ("hevc_qsv", "hevc", ["-global_quality", "28"]),
    ("h264_amf", "h264", ["-quality", "balanced"]),
    ("hevc_amf", "hevc", ["-quality", "balanced"]),
    ("libx264", "h264", ["-preset", "ultrafast", "-crf", "26"]),
    ("libx265", "hevc", ["-preset", "ultrafast", "-crf", "28"]),
)


def probe_encoder(name: str, arguments: list[str], width: int = 1280, height: int = 720) -> dict:
    """Encode a handful of real frames. Anything short of that proves nothing.

    A 16x16 probe passes on drivers that then fail at 1080p, so the probe runs at
    a realistic size and for several frames, which is also what surfaces the
    "no capable devices" class of error.
    """
    started = time.perf_counter()
    completed = subprocess.run(
        [
            "ffmpeg", "-hide_banner", "-v", "error", "-nostdin",
            "-f", "lavfi", "-i", f"testsrc2=size={width}x{height}:rate=30:duration=0.4",
            "-c:v", name, *arguments, "-pix_fmt", "yuv420p",
            "-f", "null", "-",
        ],
        capture_output=True, text=True, errors="replace",
    )
    elapsed = time.perf_counter() - started
    return {
        "encoder": name,
        "ok": completed.returncode == 0,
        "seconds": round(elapsed, 3),
        # The whole stderr is kept: the line that explains an NVENC refusal
        # ("Driver does not support the required nvenc API version") appears near
        # the start, and a tail-truncated capture loses exactly the useful part.
        "error": (completed.stderr or "").strip() if completed.returncode else None,
    }


def probe_hwaccel(kind: str = "cuda") -> dict:
    """Probe decode acceleration and measure whether it is actually faster.

    Two distinct paths, and mixing them up is why this check is easy to get wrong:

      `-hwaccel cuda` alone decodes on the GPU and hands frames back in system
      memory, so ordinary CPU filters just work.

      `-hwaccel cuda -hwaccel_output_format cuda` keeps frames on the GPU, which
      needs `hwdownload` before any CPU filter — and inserting `hwdownload` on the
      first path fails outright, which is what made an earlier version of this
      probe report CUDA as broken on a machine where it works.

    Both are probed, and the timing against a plain CPU decode is reported,
    because on a mid-range card the gain can be a few percent — not worth the
    extra graph complexity.
    """
    if shutil.which("ffmpeg") is None:
        return {"hwaccel": kind, "ok": False, "error": "ffmpeg not found"}
    sample = CACHE / "hwaccel_probe.mp4"
    sample.parent.mkdir(parents=True, exist_ok=True)
    built = subprocess.run(
        ["ffmpeg", "-hide_banner", "-v", "error", "-nostdin", "-y",
         "-f", "lavfi", "-i", "testsrc2=size=1920x1080:rate=30:duration=3",
         "-c:v", "libx264", "-preset", "veryfast", "-g", "60", str(sample)],
        capture_output=True, text=True, errors="replace",
    )
    if built.returncode:
        return {"hwaccel": kind, "ok": False, "error": "could not create the probe clip"}

    def timed(arguments: list[str], chain: str) -> tuple[bool, float, str]:
        started = time.perf_counter()
        run = subprocess.run(
            ["ffmpeg", "-hide_banner", "-v", "error", "-nostdin", *arguments,
             "-i", str(sample), "-vf", chain, "-f", "null", "-"],
            capture_output=True, text=True, errors="replace",
        )
        return run.returncode == 0, time.perf_counter() - started, (run.stderr or "").strip()[-300:]

    cpu_ok, cpu_seconds, _ = timed([], "scale=640:360")
    auto_ok, auto_seconds, auto_error = timed(["-hwaccel", kind], "scale=640:360")
    keep_ok, keep_seconds, keep_error = timed(
        ["-hwaccel", kind, "-hwaccel_output_format", kind],
        "scale_cuda=640:360,hwdownload,format=nv12" if kind == "cuda" else "hwdownload,scale=640:360",
    )
    sample.unlink(missing_ok=True)
    speedup = round(cpu_seconds / auto_seconds, 2) if auto_ok and auto_seconds > 0 else None
    return {
        "hwaccel": kind,
        "ok": auto_ok,
        "cpu_path_ok": cpu_ok,
        "system_memory_path": {"ok": auto_ok, "seconds": round(auto_seconds, 2),
                               "error": auto_error if not auto_ok else None},
        "gpu_memory_path": {"ok": keep_ok, "seconds": round(keep_seconds, 2),
                            "error": keep_error if not keep_ok else None},
        "cpu_seconds": round(cpu_seconds, 2),
        "speedup_vs_cpu": speedup,
        "worth_using": bool(speedup and speedup >= 1.25),
        "note": (
            "Включай только если speedup_vs_cpu заметный. На средней карте "
            "выигрыш на декоде часто в пределах погрешности, а любой CPU-фильтр "
            "всё равно требует выгрузки кадров из памяти GPU."
        ),
    }


NVENC_DRIVER_HINT = re.compile(
    r"minimum required Nvidia driver for nvenc is ([\d.]+)", re.IGNORECASE
)
NVENC_API_HINT = re.compile(
    r"required nvenc API version\.\s*Required:\s*([\d.]+)\s*Found:\s*([\d.]+)", re.IGNORECASE
)


def diagnose(encoder: str, stderr: str, driver: str | None) -> str:
    """Turn an encoder failure into something the operator can act on."""
    text = stderr or ""
    api = NVENC_API_HINT.search(text)
    needed = NVENC_DRIVER_HINT.search(text)
    if api:
        required, found = api.group(1), api.group(2)
        remedy = (
            f"эта сборка FFmpeg собрана под NVENC API {required}, а драйвер "
            f"{driver or 'установленный'} даёт {found}."
        )
        if needed:
            remedy += (
                f" Нужен драйвер NVIDIA {needed.group(1)} или новее — либо сборка "
                f"FFmpeg под старый NVENC SDK."
            )
        return f"{encoder}: {remedy} Рендер идёт на CPU, это не ошибка скрипта."
    if "No capable devices found" in text:
        return f"{encoder}: карта не умеет этот кодек в железе."
    if "OpenEncodeSessionEx failed" in text and "out of memory" in text:
        return f"{encoder}: превышен лимит одновременных NVENC-сессий, закрой другие рендеры."
    if "not built into this ffmpeg" in text:
        return f"{encoder}: отсутствует в этой сборке FFmpeg."
    return f"{encoder}: {text.strip()[-200:]}"


def capabilities(refresh: bool = False) -> dict:
    """Probe (or reuse) the machine's real encoding capabilities."""
    driver = driver_version()
    key = {
        "ffmpeg": ffmpeg_version(),
        "driver": driver,
        "gpu": gpu_name(),
    }
    if not refresh and PROBE_CACHE.exists():
        try:
            cached = json.loads(PROBE_CACHE.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            cached = None
        if cached and cached.get("key") == key and \
                time.time() - cached.get("stamp", 0) < CACHE_TTL:
            return cached

    available = set()
    if shutil.which("ffmpeg"):
        listing = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"],
                                 capture_output=True, text=True, errors="replace")
        for line in (listing.stdout or "").splitlines():
            parts = line.split()
            if len(parts) >= 2:
                available.add(parts[1])

    results = []
    for name, codec, arguments in ENCODER_CANDIDATES:
        if name not in available:
            results.append({"encoder": name, "codec": codec, "ok": False,
                            "error": "not built into this ffmpeg"})
            continue
        report = probe_encoder(name, arguments)
        report["codec"] = codec
        results.append(report)

    working = [item["encoder"] for item in results if item["ok"]]
    payload = {
        "key": key,
        "stamp": time.time(),
        "gpu": key["gpu"],
        "driver": driver,
        "encoders": [
            {**item, "error": (item["error"] or "")[-600:] or None} for item in results
        ],
        "working": working,
        "hwaccel": probe_hwaccel("cuda") if "cuda" in hwaccels() else
                   {"hwaccel": "cuda", "ok": False, "error": "not built into this ffmpeg"},
        "rejected": [
            diagnose(item["encoder"], item.get("error") or "", driver)
            for item in results if not item["ok"]
        ],
        "cpu_threads": os.cpu_count(),
    }
    PROBE_CACHE.parent.mkdir(parents=True, exist_ok=True)
    PROBE_CACHE.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
                           encoding="utf-8")
    return payload


def ffmpeg_version() -> str | None:
    if shutil.which("ffmpeg") is None:
        return None
    completed = subprocess.run(["ffmpeg", "-hide_banner", "-version"],
                               capture_output=True, text=True, errors="replace")
    if completed.returncode:
        return None
    return (completed.stdout or "").splitlines()[0].strip()


def hwaccels() -> list[str]:
    if shutil.which("ffmpeg") is None:
        return []
    completed = subprocess.run(["ffmpeg", "-hide_banner", "-hwaccels"],
                               capture_output=True, text=True, errors="replace")
    lines = (completed.stdout or "").splitlines()[1:]
    return [line.strip() for line in lines if line.strip()]


# ------------------------------------------------------------------ encoder args

def choose_encoder(codec: str = "h264", prefer_gpu: bool = True,
                   caps: dict | None = None) -> tuple[str, str]:
    """Return `(encoder, why)` for the requested codec family."""
    caps = caps or capabilities()
    working = caps["working"]
    order = {
        "h264": ["h264_nvenc", "h264_qsv", "h264_amf", "libx264"],
        "hevc": ["hevc_nvenc", "hevc_qsv", "hevc_amf", "libx265"],
        "av1": ["av1_nvenc", "libx264"],
    }[codec]
    if not prefer_gpu:
        order = [name for name in order if name.startswith("lib")] + order
    for name in order:
        if name in working:
            reason = (
                "GPU-энкодер прошёл реальный пробинг" if not name.startswith("lib")
                else "программный энкодер: GPU недоступен или отклонён пробингом"
            )
            return name, reason
    raise SystemExit(
        f"No working encoder for {codec}. Probed: "
        + "; ".join(caps.get("rejected", [])[:4])
    )


def video_arguments(encoder: str, profile: str = DEFAULT_PROFILE,
                    ten_bit: bool = False) -> list[str]:
    """Quality-targeted arguments for one encoder and profile."""
    spec = PROFILES.get(profile) or PROFILES[DEFAULT_PROFILE]
    if encoder.endswith("_nvenc"):
        arguments = [
            "-c:v", encoder,
            "-preset", spec.nvenc_preset,
            "-tune", "hq",
            "-rc", "vbr",
            "-cq", str(spec.nvenc_cq),
            "-b:v", "0",
            # Pascal ignores multipass; on Turing and later it is a real quality
            # win at almost no cost, and passing it is harmless either way.
            "-multipass", "qres",
            "-rc-lookahead", "20",
            "-spatial-aq", "1",
            "-aq-strength", "8",
            "-bf", "3",
            "-g", "120",
        ]
        if encoder.startswith("h264"):
            arguments += ["-profile:v", "high", "-coder", "cabac"]
        else:
            arguments += ["-profile:v", "main10" if ten_bit else "main"]
        return arguments
    if encoder.endswith("_qsv"):
        return ["-c:v", encoder, "-global_quality", str(spec.nvenc_cq),
                "-look_ahead", "1", "-g", "120"]
    if encoder.endswith("_amf"):
        return ["-c:v", encoder, "-quality", "quality", "-rc", "cqp",
                "-qp_i", str(spec.nvenc_cq), "-qp_p", str(spec.nvenc_cq + 2), "-g", "120"]
    if encoder == "libx265":
        return ["-c:v", "libx265", "-preset", spec.x264_preset,
                "-crf", str(spec.x264_crf), "-tag:v", "hvc1", "-g", "120"]
    # With NVENC unavailable, software encoding speed is what matters. x264
    # already threads well; capping frame threads keeps latency sane on long
    # graphs without giving up throughput.
    threads = max(1, min(16, (os.cpu_count() or 4)))
    return [
        "-c:v", "libx264",
        "-preset", spec.x264_preset,
        "-crf", str(spec.x264_crf),
        "-profile:v", "high",
        "-threads", str(threads),
        "-g", "120",
    ]


def audio_arguments(bitrate: str = "192k") -> list[str]:
    return ["-c:a", "aac", "-b:a", bitrate, "-ar", "48000", "-ac", "2"]


def container_arguments() -> list[str]:
    return ["-movflags", "+faststart"]


def output_arguments(codec: str = "h264", profile: str = DEFAULT_PROFILE,
                     prefer_gpu: bool = True, audio_bitrate: str = "192k",
                     ten_bit: bool = False, caps: dict | None = None) -> tuple[list[str], dict]:
    """One call for a complete, probed output configuration."""
    caps = caps or capabilities()
    encoder, why = choose_encoder(codec, prefer_gpu, caps)
    arguments = [
        *video_arguments(encoder, profile, ten_bit),
        "-pix_fmt", "yuv420p10le" if ten_bit else "yuv420p",
        *audio_arguments(audio_bitrate),
        *container_arguments(),
    ]
    return arguments, {
        "encoder": encoder,
        "reason": why,
        "profile": profile,
        "gpu": caps.get("gpu"),
        "driver": caps.get("driver"),
        "hardware": not encoder.startswith("lib"),
    }


def decode_arguments(use_hwaccel: bool, caps: dict | None = None) -> list[str]:
    """Input-side flags. Empty unless CUDA decode actually probed working."""
    if not use_hwaccel:
        return []
    caps = caps or capabilities()
    if not (caps.get("hwaccel") or {}).get("ok"):
        return []
    # Frames are downloaded on demand by the filter graph; keeping the output
    # format unset lets ffmpeg insert hwdownload where it is needed.
    return ["-hwaccel", "cuda"]


def benchmark(seconds: float = 3.0, width: int = 1080, height: int = 1920) -> list[dict]:
    """Measure real throughput per working encoder at the delivery profile."""
    caps = capabilities()
    results = []
    for encoder in caps["working"]:
        arguments = video_arguments(encoder, "delivery")
        started = time.perf_counter()
        completed = subprocess.run(
            ["ffmpeg", "-hide_banner", "-v", "error", "-nostdin",
             "-f", "lavfi", "-i",
             f"testsrc2=size={width}x{height}:rate=30:duration={seconds}",
             *arguments, "-pix_fmt", "yuv420p", "-f", "null", "-"],
            capture_output=True, text=True, errors="replace",
        )
        elapsed = time.perf_counter() - started
        results.append({
            "encoder": encoder,
            "ok": completed.returncode == 0,
            "wall_seconds": round(elapsed, 2),
            "realtime_factor": round(seconds / elapsed, 2) if elapsed > 0 else None,
            "error": (completed.stderr or "").strip()[-200:] if completed.returncode else None,
        })
    results.sort(key=lambda item: -(item["realtime_factor"] or 0))
    return results


def main() -> None:
    utf8_stdout()
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--refresh", action="store_true", help="Re-probe, ignoring the cache")
    parser.add_argument("--benchmark", action="store_true",
                        help="Measure throughput of every working encoder at 1080x1920")
    parser.add_argument("--codec", default="h264", choices=["h264", "hevc", "av1"])
    parser.add_argument("--profile", default=DEFAULT_PROFILE, choices=list(PROFILES))
    parser.add_argument("--cpu", action="store_true", help="Ask for the software encoder")
    args = parser.parse_args()

    caps = capabilities(refresh=args.refresh)
    arguments, chosen = output_arguments(args.codec, args.profile, not args.cpu, caps=caps)
    payload = {
        "gpu": caps.get("gpu"),
        "driver": caps.get("driver"),
        "working_encoders": caps["working"],
        "cpu_threads": caps.get("cpu_threads"),
        "rejected": caps.get("rejected", []),
        "hwaccel": caps.get("hwaccel"),
        "profiles": {key: {"use": value.use, "crf": value.x264_crf, "nvenc_cq": value.nvenc_cq}
                     for key, value in PROFILES.items()},
        "selection": chosen,
        "ffmpeg_arguments": arguments,
    }
    if args.benchmark:
        payload["benchmark"] = benchmark()
    emit(payload)


if __name__ == "__main__":
    main()