#!/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 / "cpu-smoke-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

    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = ""
    env["MULTISCREEN_BACKEND"] = "torch"
    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",
        "float32",
        "--device",
        "cpu",
        "--seed",
        "1337",
        "--output",
        str(OUTPUT),
    ]

    print(f"Checkpoint: {ROOT}")
    print("Device: cpu")
    print("Backend: torch")
    print("Triton imported: False")
    print()

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

    if result.returncode != 0:
        print(
            f"CPU-only 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("CPU-only smoke: PASS")
    print("Backend: torch")
    print("Triton imported: False")
    print(f"Results: {OUTPUT}")
    return 0


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