from __future__ import annotations

import argparse
import os
import shlex
import subprocess
import sys
from pathlib import Path


def _preflight() -> None:
    print("MVP_PRECHECK_PYTHON:", sys.executable, flush=True)
    print("MVP_PRECHECK_PYTHONPATH:", os.environ.get("PYTHONPATH"), flush=True)
    print("MVP_PRECHECK_MVP_CODE_PATH:", os.environ.get("MVP_CODE_PATH"), flush=True)
    try:
        import verl  # type: ignore

        print("MVP_PRECHECK_VERL_FILE:", getattr(verl, "__file__", None), flush=True)
    except Exception as e:
        print("MVP_PRECHECK_VERL_IMPORT_ERROR:", repr(e), flush=True)

    try:
        import mvp_marker  # type: ignore

        print("MVP_PRECHECK_MARKER:", getattr(mvp_marker, "MARKER", None), flush=True)
    except Exception as e:
        print("MVP_PRECHECK_MARKER_MISSING:", repr(e), flush=True)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--job-dir", required=True)
    parser.add_argument("cmd", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    job_dir = Path(args.job_dir)
    job_dir.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("MVP_JOB_DIR", str(job_dir))
    os.chdir(job_dir)

    _preflight()

    if not args.cmd:
        print("no command provided", file=sys.stderr)
        return 2

    # argparse includes the leading "--" if the caller uses it; strip it.
    cmd = list(args.cmd)
    if cmd and cmd[0] == "--":
        cmd = cmd[1:]
    if not cmd:
        print("no command provided", file=sys.stderr)
        return 2

    # Execute training command as a subprocess so that logs are captured by Ray job logs.
    cmd_str = " ".join(shlex.quote(x) for x in cmd)
    print("MVP_DRIVER_EXEC:", cmd_str, flush=True)

    proc = subprocess.run(cmd, check=False)
    print("MVP_DRIVER_EXIT_CODE:", proc.returncode, flush=True)
    return proc.returncode


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