#!/usr/bin/env python3
"""
Asserts the Python and torch the image ended up with can actually run ComfyUI.

Both engine failures so far were import-time and only showed up at runtime, in
a container log, long after the build reported success. Checking here turns
each of them into one line that fails the build instead.
"""
import os
import shutil
import subprocess
import sys
import tempfile

failures = []

# Ubuntu 22.04's python3.11 package is 3.11.0~rc1, which predates the
# CVE-2020-10735 fix that added these. torch._dynamo references
# sys.get_int_max_str_digits unconditionally at import.
if not hasattr(sys, "get_int_max_str_digits"):
    failures.append(
        f"python {sys.version.split()[0]} lacks sys.get_int_max_str_digits, "
        "so torch._dynamo cannot import. Use a python past 3.10.7 / 3.11.0rc2."
    )

try:
    import torch

    version = tuple(int(part) for part in torch.__version__.split(".")[:2])
    # comfy_kitchen annotates custom ops with PEP 585 list[int], which
    # torch.library.infer_schema only accepts from 2.7 onward.
    if version < (2, 7):
        failures.append(
            f"torch {torch.__version__} is older than 2.7, so comfy_kitchen "
            "will fail with an infer_schema ValueError."
        )
    else:
        print(f"verify-engine: torch {torch.__version__}")
except Exception as error:  # noqa: BLE001
    failures.append(f"importing torch failed: {error}")

# Triton compiles a small C launcher the first time a kernel runs. Without a
# toolchain that surfaces as "Failed to find C compiler" partway through a
# render, long after the build claimed success.
compiler = os.environ.get("CC") or "cc"
found = shutil.which(compiler) or shutil.which("gcc") or shutil.which("cc")
if not found:
    failures.append(
        f"no C compiler on PATH (looked for {compiler}, gcc, cc). Triton needs "
        "one to build kernels at runtime. Install build-essential."
    )
else:
    with tempfile.TemporaryDirectory() as work:
        source = os.path.join(work, "probe.c")
        with open(source, "w") as handle:
            handle.write("int main(void) { return 0; }\n")
        probe = subprocess.run(
            [found, source, "-o", os.path.join(work, "probe")],
            capture_output=True,
            text=True,
        )
        if probe.returncode != 0:
            failures.append(
                f"{found} is present but cannot compile: {probe.stderr.strip()[:200]}"
            )
        else:
            print(f"verify-engine: C compiler {found} works")

# The import that actually blew up on the Space.
try:
    import torch._dynamo  # noqa: F401

    print("verify-engine: torch._dynamo imports")
except Exception as error:  # noqa: BLE001
    failures.append(f"importing torch._dynamo failed: {error}")

if failures:
    print("verify-engine: FAILED", file=sys.stderr)
    for failure in failures:
        print(f"  - {failure}", file=sys.stderr)
    raise SystemExit(1)

print(f"verify-engine: ok on python {sys.version.split()[0]}")