"""Correct Decord's known wheel label, then check dependencies and CPU decoding.

No model loading, downloads, credentials or CUDA execution. See upstream #356:
https://github.com/dmlc/decord/issues/356
"""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import platform
import subprocess
import sys
import tempfile
from importlib.metadata import Distribution, distribution, version
from pathlib import Path

ORIGINAL_WHEEL = (
    b"Wheel-Version: 1.0\nGenerator: bdist_wheel (0.36.2)\n"
    b"Root-Is-Purelib: false\nTag: cp36-cp36m-manylinux2010_x86_64\n\n"
)
CORRECTED_WHEEL = ORIGINAL_WHEEL.replace(b"cp36-cp36m", b"py3-none")
LIBDECORD_SHA256 = "98b260c5812106648ba299279916fbe98439893e346d4efdcf5cde66ba8973da"
WHEEL_PATH = "decord-0.6.0.dist-info/WHEEL"
RECORD_PATH = "decord-0.6.0.dist-info/RECORD"


def record_digest(data: bytes) -> str:
    return "sha256=" + base64.urlsafe_b64encode(hashlib.sha256(data).digest()).decode().rstrip("=")


def normalize_decord_metadata(dist: Distribution) -> None:
    """Edit only the exact published 0.6.0 metadata, preserving every binary."""
    if dist.version != "0.6.0":
        raise RuntimeError("Decord metadata correction requires exactly version 0.6.0")
    wheel = Path(dist.locate_file(WHEEL_PATH))
    record = Path(dist.locate_file(RECORD_PATH))
    original = wheel.read_bytes()
    if original not in (ORIGINAL_WHEEL, CORRECTED_WHEEL):
        raise RuntimeError("Unexpected Decord WHEEL metadata; refusing to retag it")
    native = Path(dist.locate_file("decord/libdecord.so"))
    if hashlib.sha256(native.read_bytes()).hexdigest() != LIBDECORD_SHA256:
        raise RuntimeError("Unexpected Decord native library; refusing to retag it")
    rows = list(csv.reader(io.StringIO(record.read_text())))
    targets = [row for row in rows if row and row[0] == WHEEL_PATH]
    if len(targets) != 1 or targets[0][1:] != [record_digest(original), str(len(original))]:
        raise RuntimeError("Decord WHEEL/RECORD mismatch; refusing to modify metadata")
    if original == CORRECTED_WHEEL:
        return
    targets[0][1:] = [record_digest(CORRECTED_WHEEL), str(len(CORRECTED_WHEEL))]
    updated = io.StringIO(newline="")
    csv.writer(updated, lineterminator="\n").writerows(rows)
    wheel.write_bytes(CORRECTED_WHEEL)
    record.write_text(updated.getvalue())
    print("Corrected Decord 0.6.0 wheel tag to its published Python 3 tag; binaries unchanged.")


def verify_cpu_decode() -> None:
    import decord
    import numpy as np

    with tempfile.TemporaryDirectory(prefix="omnai-decode-") as directory:
        clip = Path(directory) / "synthetic.mp4"
        subprocess.run([
            "ffmpeg", "-nostdin", "-v", "error", "-f", "lavfi", "-i",
            "testsrc2=size=64x64:rate=4:duration=1", "-an", "-c:v", "mpeg4",
            "-pix_fmt", "yuv420p", "-threads", "1", str(clip),
        ], check=True, timeout=30)
        reader = decord.VideoReader(str(clip), ctx=decord.cpu(0), num_threads=1)
        frames = reader.get_batch([0, 3]).asnumpy()
        if len(reader) != 4 or frames.shape != (2, 64, 64, 3) or frames.dtype != np.uint8:
            raise RuntimeError("Decord failed the synthetic CPU video-decode check")
        if not np.any(frames[0] != frames[1]):
            raise RuntimeError("Decord returned identical frames for the changing test pattern")
    print("Decord CPU decode passed: four source frames, two RGB samples. No GPU/model used.")


def main() -> None:
    if platform.system() != "Linux" or platform.machine() != "x86_64" or sys.version_info[:2] != (3, 10):
        raise RuntimeError("This build check targets the pinned Python 3.10 Linux x86_64 image")
    normalize_decord_metadata(distribution("decord"))
    subprocess.run([sys.executable, "-m", "pip", "check"], check=True, timeout=120)
    if version("torch") != "2.8.0+cu128" or version("torchaudio") != "2.8.0+cu128":
        raise RuntimeError("The installed Torch/Torchaudio wheels do not match the CUDA pins")
    verify_cpu_decode()


if __name__ == "__main__":
    main()