from __future__ import annotations

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


ROOT = Path(__file__).resolve().parent
ARTIFACT = ROOT / "artifacts" / "underflow_appends.pt"
EXPECTED_SHA256 = "bd34cabfd35feb96ced8a63489c5899aa9b17c000ce552406519998d5d3855f1"


LOAD_ONCE = r"""
import os
import sys
import torch

artifact = sys.argv[1]
print(f"python={sys.version.split()[0]}", flush=True)
print(f"torch={torch.__version__}", flush=True)
print(f"artifact={os.path.abspath(artifact)}", flush=True)
torch.jit.load(artifact, map_location="cpu")
print("loaded", flush=True)
"""


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--python",
        default=sys.executable,
        help="Python executable with torch==1.4.0 installed.",
    )
    args = parser.parse_args()

    digest = sha256(ARTIFACT)
    if digest != EXPECTED_SHA256:
        raise SystemExit(f"unexpected artifact sha256: {digest}")

    proc = subprocess.run(
        [args.python, "-c", LOAD_ONCE, str(ARTIFACT)],
        cwd=ROOT,
        text=True,
        capture_output=True,
        timeout=20,
    )
    result = {
        "artifact": str(ARTIFACT),
        "artifact_sha256": digest,
        "python": args.python,
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "native_crash": proc.returncode < 0 or proc.returncode in {134, 139},
    }
    print(json.dumps(result, indent=2))

    if not result["native_crash"]:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())