#!/usr/bin/env python3
"""check_ffmpeg_build.py -- hold a Windows FFmpeg bin/ directory to ffmpeg_allowlist.json.

    python scripts/ffmpeg/check_ffmpeg_build.py <bin-dir> [--consumer <exe|dll> ...] [--json out.json]

PASS only if ALL of these hold (each is its own row in the output):

  files       the 7 DLLs carry the allowlisted major versions (avcodec-62.dll ...), ffmpeg.exe, ffprobe.exe
  license     <lib>_license() of every library == "LGPL version 2.1 or later"   (called via ctypes)
  version     <lib>_version() major == allowlist; ffmpeg.exe reports the allowlisted release
  configure   no --enable-gpl / --enable-version3 / --enable-nonfree in any library's
              <lib>_configuration() nor anywhere in the bytes of any DLL/EXE; no --enable-lib*
              other than allowed_external_libs
  encoders    `ffmpeg -encoders`  == allowlist "encoders"   (exact set)
  muxers      `ffmpeg -muxers`    == allowlist "muxers"     (exact set)
  protocols   `ffmpeg -protocols` == allowlist "protocols"  (exact set, input and output)
  devices     `ffmpeg -devices`   == allowlist indevs / outdevs (exact sets)
  hwaccels / decoders / demuxers / filters   required subsets present
  guess       av_guess_format(ext) picks the allowlisted muxer (clip_trim remux keeps the extension)
  imports     every DLL/EXE imports only its sibling FFmpeg DLLs and Windows system DLLs
              (no libwinpthread/libgcc/zlib1/libdav1d/msvcrt... next to it)
  consumer:*  (--consumer) every av*/sw* symbol a consumer binary imports is exported by this build,
              i.e. the build can replace the DLLs under it without a relink

Exit: 0 PASS, 1 FAIL, 2 could not run.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import struct
import subprocess
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
LIBS = ["avcodec", "avformat", "avutil", "avdevice", "avfilter", "swscale", "swresample"]

# Windows system DLLs an FFmpeg build may import. Anything else next to ffmpeg is a bundled
# third-party runtime (libwinpthread-1, libgcc_s, zlib1, libdav1d, msvcrt, ...) and fails.
SYSTEM_DLL = re.compile(
    r"^(api-ms-win-[a-z0-9-]+|kernel32|kernelbase|ntdll|user32|gdi32|advapi32|ole32|oleaut32|"
    r"shell32|shlwapi|psapi|bcrypt|secur32|crypt32|ws2_32|mfplat|mf|mfreadwrite|evr|d3d11|dxgi|"
    r"d3d12|dxva2|vfw32|avicap32|winmm|version|rpcrt4|setupapi|cfgmgr32|userenv|imm32|dwrite|"
    r"d2d1|usp10|comdlg32|strmiids|uuid|iphlpapi|ncrypt|avrt)\.dll$", re.I)


# ---------------------------------------------------------------- PE parsing (imports / exports)
class PE:
    def __init__(self, path: Path):
        self.path = path
        self.data = path.read_bytes()
        d = self.data
        if d[:2] != b"MZ":
            raise ValueError(f"{path}: not a PE file")
        pe = struct.unpack_from("<I", d, 0x3C)[0]
        if d[pe:pe + 4] != b"PE\0\0":
            raise ValueError(f"{path}: bad PE signature")
        nsec, = struct.unpack_from("<H", d, pe + 6)
        optsz, = struct.unpack_from("<H", d, pe + 20)
        opt = pe + 24
        magic, = struct.unpack_from("<H", d, opt)
        self.pe64 = magic == 0x20B
        dd = opt + (112 if self.pe64 else 96)
        ndd, = struct.unpack_from("<I", d, dd - 4)
        self.dirs = [struct.unpack_from("<II", d, dd + 8 * i) for i in range(ndd)]
        self.sections = []
        s = opt + optsz
        for i in range(nsec):
            vsize, va, rsize, rptr = struct.unpack_from("<IIII", d, s + 40 * i + 8)
            self.sections.append((va, max(vsize, rsize), rptr))

    def off(self, rva: int) -> int:
        for va, size, rptr in self.sections:
            if va <= rva < va + size:
                return rva - va + rptr
        raise ValueError(f"{self.path}: rva {rva:#x} outside sections")

    def cstr(self, rva: int) -> str:
        o = self.off(rva)
        return self.data[o:self.data.index(b"\0", o)].decode("ascii", "replace")

    def _thunks(self, rva: int):
        step, fmt, flag = (8, "<Q", 1 << 63) if self.pe64 else (4, "<I", 1 << 31)
        names = []
        o = self.off(rva)
        while True:
            v, = struct.unpack_from(fmt, self.data, o)
            if v == 0:
                return names
            names.append(f"#ord{v & 0xFFFF}" if v & flag else self.cstr((v & 0x7FFFFFFF) + 2))
            o += step

    def imports(self) -> dict[str, list[str]]:
        out: dict[str, list[str]] = {}
        rva, size = self.dirs[1] if len(self.dirs) > 1 else (0, 0)
        if rva:
            o = self.off(rva)
            while True:
                oft, _, _, name, ft = struct.unpack_from("<IIIII", self.data, o)
                if name == 0:
                    break
                out.setdefault(self.cstr(name), []).extend(self._thunks(oft or ft))
                o += 20
        rva, size = self.dirs[13] if len(self.dirs) > 13 else (0, 0)  # delay-load
        if rva:
            o = self.off(rva)
            while True:
                attrs, name, _, _, int_rva = struct.unpack_from("<IIIII", self.data, o)[:5]
                if name == 0:
                    break
                out.setdefault(self.cstr(name), []).extend(self._thunks(int_rva))
                o += 32
        return out

    def exports(self) -> set[str]:
        rva, size = self.dirs[0]
        if not rva:
            return set()
        o = self.off(rva)
        nnames, = struct.unpack_from("<I", self.data, o + 24)
        names_rva, = struct.unpack_from("<I", self.data, o + 32)
        no = self.off(names_rva)
        return {self.cstr(struct.unpack_from("<I", self.data, no + 4 * i)[0]) for i in range(nnames)}


# ---------------------------------------------------------------- ffmpeg.exe listings
def run_ff(bindir: Path, *args: str) -> str:
    exe = bindir / "ffmpeg.exe"
    r = subprocess.run([str(exe), "-hide_banner", *args], cwd=bindir, capture_output=True,
                       text=True, encoding="utf-8", errors="replace", timeout=60)
    return r.stdout + r.stderr


def parse_codecs(text: str) -> set[str]:
    out, body = set(), False
    for line in text.splitlines():
        if line.strip().startswith("---"):
            body = True
            continue
        if body and line.strip():
            parts = line.split()
            if len(parts) >= 2:
                out.add(parts[1])
    return out


def parse_formats(text: str) -> dict[str, set[str]]:
    """-muxers / -demuxers / -devices: ' <flags> name  long name'. The flag field is as wide as the
    legend's first token: 'D..' (3) for -muxers/-demuxers, 'D.' (2) for -devices."""
    out = {"D": set(), "E": set(), "d": set()}
    body, width = False, 3
    for line in text.splitlines():
        m = re.match(r"^ (D\.+) = ", line)
        if m:
            width = len(m.group(1))
        if line.strip().startswith("--"):
            body = True
            continue
        if not body or len(line) < width + 2:
            continue
        flags, rest = line[1:1 + width], line[1 + width:].split()
        if not rest:
            continue
        for name in rest[0].split(","):
            for f in "DEd":
                if f in flags:
                    out[f].add(name)
    return out


def parse_protocols(text: str) -> dict[str, set[str]]:
    out = {"Input": set(), "Output": set()}
    cur = None
    for line in text.splitlines():
        s = line.strip()
        if s in ("Input:", "Output:"):
            cur = s[:-1]
        elif cur and s:
            out[cur].add(s)
    return out


def parse_hwaccels(text: str) -> set[str]:
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    if lines and lines[0].startswith("Hardware acceleration"):
        lines = lines[1:]
    return set(lines)


def parse_filters(text: str) -> set[str]:
    out = set()
    for line in text.splitlines():
        p = line.split()
        if len(p) >= 3 and re.fullmatch(r"[T.][S.][C.]?", p[0]) and "->" in p[2]:
            out.add(p[1])
    return out


# ---------------------------------------------------------------- ctypes probe (own process)
CTYPES_PROBE = r"""
import ctypes, json, os, sys
bindir = sys.argv[1]; majors = json.loads(sys.argv[2]); exts = json.loads(sys.argv[3])
os.add_dll_directory(bindir)
out = {"libs": {}, "guess": {}}
for lib, major in majors.items():
    dll = ctypes.CDLL(os.path.join(bindir, "%s-%d.dll" % (lib, major)))
    info = {}
    for fn in ("license", "configuration"):
        f = getattr(dll, "%s_%s" % (lib, fn)); f.restype = ctypes.c_char_p
        info[fn] = f().decode("utf-8", "replace")
    v = getattr(dll, lib + "_version"); v.restype = ctypes.c_uint
    ver = v(); info["version"] = "%d.%d.%d" % (ver >> 16, (ver >> 8) & 255, ver & 255)
    out["libs"][lib] = info
fmt = ctypes.CDLL(os.path.join(bindir, "avformat-%d.dll" % majors["avformat"]))
fmt.av_guess_format.restype = ctypes.c_void_p
fmt.av_guess_format.argtypes = [ctypes.c_char_p] * 3
for ext in exts:
    p = fmt.av_guess_format(None, ("x." + ext).encode(), None)
    out["guess"][ext] = ctypes.cast(ctypes.c_void_p(p), ctypes.POINTER(ctypes.c_char_p))[0].decode() if p else None
print(json.dumps(out))
"""


def ctypes_probe(bindir: Path, majors: dict, exts: list[str]) -> dict:
    # A separate interpreter: two builds' avutil-60.dll cannot coexist in one process (the loader
    # reuses a module by name), so a check never sees DLLs another check loaded.
    r = subprocess.run([sys.executable, "-c", CTYPES_PROBE, str(bindir), json.dumps(majors), json.dumps(exts)],
                       capture_output=True, text=True, timeout=60)
    if r.returncode != 0:
        raise RuntimeError("ctypes probe failed:\n" + r.stderr[-2000:])
    return json.loads(r.stdout)


# ---------------------------------------------------------------- main
def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("bindir", type=Path)
    ap.add_argument("--allowlist", type=Path, default=HERE / "ffmpeg_allowlist.json")
    ap.add_argument("--consumer", type=Path, action="append", default=[],
                    help="exe/dll linked against these DLLs; all of its av*/sw* imports must be exported")
    ap.add_argument("--json", type=Path, help="write the result rows here")
    a = ap.parse_args()

    allow = json.loads(a.allowlist.read_text(encoding="utf-8"))
    keys = lambda k: {x for x in allow.get(k, {}) if not x.startswith("_")}
    majors = allow["lib_majors"]
    bindir = a.bindir.resolve()
    rows: list[tuple[str, bool, str]] = []

    def row(name: str, ok: bool, detail: str = "") -> None:
        rows.append((name, ok, detail))

    def exact(name: str, got: set[str], want: set[str]) -> None:
        extra, missing = sorted(got - want), sorted(want - got)
        detail = f"{len(got)} built"
        if extra:
            detail += f"; NOT allowlisted ({len(extra)}): {' '.join(extra[:40])}{' ...' if len(extra) > 40 else ''}"
        if missing:
            detail += f"; missing: {' '.join(missing)}"
        row(name, not extra and not missing, detail)

    def subset(name: str, got: set[str], want: set[str]) -> None:
        missing = sorted(want - got)
        row(name, not missing, f"{len(want)} required, {len(got)} built" + (f"; missing: {' '.join(missing)}" if missing else ""))

    # files
    dlls = {lib: bindir / f"{lib}-{m}.dll" for lib, m in majors.items()}
    exes = [bindir / "ffmpeg.exe", bindir / "ffprobe.exe"]
    absent = [p.name for p in [*dlls.values(), *exes] if not p.is_file()]
    row("files", not absent, "missing: " + " ".join(absent) if absent else
        " ".join(p.name for p in dlls.values()) + " ffmpeg.exe ffprobe.exe")
    if absent:
        return report(rows, a.json)

    # license / version / configure via the libraries themselves
    guess = {k: v for k, v in allow.get("guess_format", {}).items() if not k.startswith("_")}
    try:
        probe = ctypes_probe(bindir, majors, sorted(guess))
    except Exception as e:  # noqa: BLE001
        print(f"could not load the DLLs: {e}", file=sys.stderr)
        return 2
    bad = {lib: i["license"] for lib, i in probe["libs"].items() if i["license"] != allow["license"]}
    row("license", not bad, "; ".join(f"{l}: {v}" for l, v in bad.items()) if bad
        else f'all 7 libraries: "{allow["license"]}"')
    badv = {lib: i["version"] for lib, i in probe["libs"].items() if int(i["version"].split(".")[0]) != majors[lib]}
    verline = next((l for l in run_ff(bindir, "-version").splitlines() if l.startswith("ffmpeg version")), "?")
    want_ver = allow["ffmpeg_version"]
    ver_ok = not badv and re.match(rf"ffmpeg version n?{re.escape(want_ver)}(\s|$)", verline) is not None
    row("version", ver_ok, f"{verline.split(' Copyright')[0]}; " +
        " ".join(f"{l}={i['version']}" for l, i in probe["libs"].items()))

    forbidden = allow["forbidden_configure_flags"]
    hits = []
    for lib, i in probe["libs"].items():
        flags = i["configuration"].split()
        hits += [f"{lib}:{f}" for f in flags if f in forbidden]
        hits += [f"{lib}:{f}" for f in flags if f.startswith("--enable-lib")
                 and f[len("--enable-"):] not in allow["allowed_external_libs"]]
    for p in [*dlls.values(), *exes]:
        blob = p.read_bytes()
        hits += [f"{p.name}:bytes:{f}" for f in forbidden if f.encode() in blob]
    extlibs = sorted({f for i in probe["libs"].values() for f in i["configuration"].split() if f.startswith("--enable-lib")})
    row("configure", not hits, ("; ".join(sorted(set(hits))[:30]) if hits else "no gpl/version3/nonfree") +
        f"; external libs: {' '.join(extlibs) or '(none)'}")

    # component sets from ffmpeg.exe
    exact("encoders", parse_codecs(run_ff(bindir, "-encoders")), keys("encoders"))
    exact("muxers", parse_formats(run_ff(bindir, "-muxers"))["E"], keys("muxers"))
    prot = parse_protocols(run_ff(bindir, "-protocols"))
    exact("protocols", prot["Input"] | prot["Output"], keys("protocols"))
    dev = parse_formats(run_ff(bindir, "-devices"))
    exact("indevs", dev["D"], keys("indevs"))
    exact("outdevs", dev["E"], keys("outdevs"))
    subset("hwaccels", parse_hwaccels(run_ff(bindir, "-hwaccels")), set(allow["required_hwaccels"]))
    subset("decoders", parse_codecs(run_ff(bindir, "-decoders")), set(allow["required_decoders"]))
    subset("demuxers", parse_formats(run_ff(bindir, "-demuxers"))["D"], set(allow["required_demuxers"]))
    subset("filters", parse_filters(run_ff(bindir, "-filters")), set(allow["required_filters"]))

    wrong = {e: (probe["guess"].get(e), m) for e, m in guess.items() if probe["guess"].get(e) != m}
    row("guess", not wrong, "; ".join(f".{e} -> {g} (want {m})" for e, (g, m) in wrong.items()) if wrong
        else f"{len(guess)} extensions map to their allowlisted muxer")

    # imports: nothing but siblings + Windows
    siblings = {p.name.lower() for p in dlls.values()}
    stray = []
    for p in [*dlls.values(), *exes]:
        for dll in PE(p).imports():
            if dll.lower() not in siblings and not SYSTEM_DLL.match(dll):
                stray.append(f"{p.name}->{dll}")
    row("imports", not stray, "; ".join(stray) if stray else "system DLLs + sibling FFmpeg DLLs only")

    # consumers: every imported av*/sw* symbol must be exported here
    exports = {p.name.lower(): PE(p).exports() for p in dlls.values()}
    for c in a.consumer:
        if not c.is_file():
            row(f"consumer:{c.name}", False, f"not found: {c}")
            continue
        missing, n, foreign = [], 0, []
        for dll, syms in PE(c).imports().items():
            key = dll.lower()
            if not re.match(r"^(av|sw)[a-z]+-\d+\.dll$", key):
                continue
            if key not in exports:
                foreign.append(dll)
                continue
            for s in syms:
                n += 1
                if s not in exports[key]:
                    missing.append(f"{dll}!{s}")
        ok = not missing and not foreign
        detail = f"{n} imported symbols resolved"
        if foreign:
            detail = f"links DLL names this build does not have: {' '.join(foreign)}"
        if missing:
            detail += f"; NOT exported ({len(missing)}): {' '.join(missing[:20])}"
        row(f"consumer:{c.name}", ok, detail)

    return report(rows, a.json)


def report(rows, json_path) -> int:
    w = max(len(r[0]) for r in rows)
    for name, ok, detail in rows:
        print(f"{'PASS' if ok else 'FAIL'}  {name:<{w}}  {detail}")
    failed = [r[0] for r in rows if not r[1]]
    print(f"RESULT: {'PASS' if not failed else 'FAIL (' + ', '.join(failed) + ')'}")
    if json_path:
        json_path.write_text(json.dumps([{"check": n, "pass": o, "detail": d} for n, o, d in rows], indent=1),
                             encoding="utf-8")
    return 0 if not failed else 1


if __name__ == "__main__":
    sys.exit(main())