#!/usr/bin/env python3

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent
RUNTIME = ROOT / "runtime"
TEST_SCRIPT = RUNTIME / "lab/scripts/test_checkpoint_capabilities.py"
OUTPUT = ROOT / "gpu-triton-results.json"


def main() -> int:
    if not TEST_SCRIPT.is_file():
        print(f"ERROR: missing capability runner: {TEST_SCRIPT}", file=sys.stderr)
        return 1

    if not (ROOT / "lit_model.pth").is_file():
        print("ERROR: missing lit_model.pth", file=sys.stderr)
        return 1

    if not (ROOT / "model_config.yaml").is_file():
        print("ERROR: missing model_config.yaml", file=sys.stderr)
        return 1

    try:
        import torch
    except Exception as exc:
        print(f"ERROR: PyTorch import failed: {exc}", file=sys.stderr)
        return 1

    if not torch.cuda.is_available():
        print("ERROR: CUDA is not available.", file=sys.stderr)
        return 1

    try:
        import triton
    except Exception as exc:
        print(f"ERROR: Triton import failed: {exc}", file=sys.stderr)
        return 1

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0")
    env["MULTISCREEN_BACKEND"] = "triton"
    env["TOKENIZERS_PARALLELISM"] = "false"
    env["PYTHONUNBUFFERED"] = "1"

    old_pythonpath = env.get("PYTHONPATH", "")
    env["PYTHONPATH"] = (
        str(RUNTIME)
        if not old_pythonpath
        else str(RUNTIME) + os.pathsep + old_pythonpath
    )

    command = [
        sys.executable,
        "-u",
        str(TEST_SCRIPT),
        "--checkpoint-dir",
        str(ROOT),
        "--tokenizer-dir",
        str(ROOT),
        "--dtype",
        "bfloat16",
        "--device",
        "cuda",
        "--seed",
        "1337",
        "--output",
        str(OUTPUT),
    ]

    print(f"Checkpoint: {ROOT}")
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print("Device: cuda")
    print("Backend: triton")
    print("Triton imported: True")
    print(f"Triton version: {triton.__version__}")
    print()

    result = subprocess.run(command, cwd=ROOT, env=env)

    if result.returncode != 0:
        print(
            f"GPU Triton smoke: FAIL ({result.returncode})",
            file=sys.stderr,
        )
        return result.returncode

    if not OUTPUT.is_file() or OUTPUT.stat().st_size == 0:
        print(f"ERROR: result was not created: {OUTPUT}", file=sys.stderr)
        return 1

    print()
    print("GPU Triton smoke: PASS")
    print("Backend: triton")
    print("Triton imported: True")
    print(f"Results: {OUTPUT}")
    return 0


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