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

import argparse
import json
import os
import sys


def _ensure_import_path() -> None:
    # Allow running as a file path (not just `python3 -m argus.cli.run`).
    # Ensure `/workspace/mvp/py` is on sys.path so `import argus.*` works.
    here = os.path.dirname(os.path.abspath(__file__))
    py_root = os.path.dirname(os.path.dirname(here))  # .../py
    if py_root not in sys.path:
        sys.path.insert(0, py_root)


def main() -> int:
    _ensure_import_path()

    from argus.ray.models import JobSpec, RayConfig
    from argus.ray.ray_job_tool import RayJobTool
    from argus.ray.yaml_io import load_yaml

    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True, help="MVP config yaml (ray/service)")
    parser.add_argument("--taskspec", help="Training TaskSpec yaml (required for submit)")
    parser.add_argument("--jobspec", dest="taskspec", help=argparse.SUPPRESS)  # backward compatible alias
    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.taskspec:
            raise SystemExit("--taskspec is required for submit")
        spec = JobSpec.from_dict(load_yaml(args.taskspec))
        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())
