#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys


def _ensure_import_path() -> None:
    # Allow `python3 /workspace/.../py/run.py` to import `mvp_v11.*`
    here = os.path.dirname(os.path.abspath(__file__))
    if here not in sys.path:
        sys.path.insert(0, here)


def main() -> int:
    _ensure_import_path()

    from mvp_v11.models import JobSpec, RayConfig
    from mvp_v11.ray_job_tool import RayJobTool
    from mvp_v11.yaml_io import load_yaml

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="Ray base config yaml")
    parser.add_argument("--jobspec", help="Training jobspec yaml (required for submit)")
    parser.add_argument("--action", required=True, choices=["submit", "status", "stop", "logs", "list"])
    parser.add_argument("--submission-id", help="For status/stop/logs")
    parser.add_argument("--no-wait", action="store_true", help="Submit and return immediately")
    parser.add_argument("--tail", type=int, default=0, help="Tail N lines for logs")
    args = parser.parse_args()

    cfg = RayConfig.from_dict(load_yaml(args.config))
    tool = RayJobTool(cfg)

    if args.action == "submit":
        if not args.jobspec:
            raise SystemExit("--jobspec is required for submit")
        spec = JobSpec.from_dict(load_yaml(args.jobspec))
        submitted = tool.submit(spec, no_wait=args.no_wait)
        print(submitted)
        return 0

    if args.action in ("status", "stop", "logs"):
        sid = args.submission_id or ""
        if not sid:
            raise SystemExit("--submission-id is required for status/stop/logs")
        if args.action == "status":
            print(tool.status(sid))
            return 0
        if args.action == "stop":
            print(tool.stop(sid))
            return 0
        logs = tool.logs(sid)
        if args.tail and args.tail > 0:
            lines = logs.splitlines()
            logs = "\n".join(lines[-args.tail :]) + ("\n" if lines else "")
        print(logs, end="")
        return 0

    if args.action == "list":
        print(json.dumps(tool.list(), indent=2))
        return 0

    raise SystemExit(f"unknown action: {args.action}")


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