from __future__ import annotations

import argparse

from one_run import __version__
from one_run.commands.ls_cmd import list_runs
from one_run.commands.cancel_cmd import cancel_run
from one_run.commands.gc_cmd import gc_runs
from one_run.commands.run_cmd import run_local
from one_run.commands.status_cmd import show_status
from one_run.commands.submit_cmd import submit_slurm
from one_run.commands.tail_cmd import tail_run_log
from one_run.commands.validate_cmd import validate_manifest
from one_run.errors import OneRunError


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="one-run", description="ml run wrapper for reprodusible experiments and documentation")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    sub = parser.add_subparsers(dest="command", required=True)

    run_p = sub.add_parser("run", help="execute a local run")
    run_p.add_argument("config", help="path to manifest yaml")
    run_p.add_argument("-w", "--watch", action="store_true", help="run in background and return immediately")

    submit_p = sub.add_parser("submit", help="submit a slurm run")
    submit_p.add_argument("config", help="path to manifest yaml")

    status_p = sub.add_parser("status", help="show run status")
    status_p.add_argument("run_id", help="run identifier")
    status_p.add_argument("--json", action="store_true", help="print status payload as json")
    status_p.add_argument("-l", "--long", action="store_true", help="print extended status fields")

    cancel_p = sub.add_parser("cancel", help="cancel a running run")
    cancel_p.add_argument("run_id", help="run identifier")
    cancel_p.add_argument("--force", action="store_true", help="force kill local run / kill signal for slurm")

    validate_p = sub.add_parser("validate", help="validate manifest only")
    validate_p.add_argument("config", help="path to manifest yaml")

    tail_p = sub.add_parser("tail", help="show run logs")
    tail_p.add_argument("run_id", help="run identifier")
    tail_p.add_argument("-n", "--lines", type=int, default=30, help="number of lines from the end")
    tail_p.add_argument("-f", "--follow", action="store_true", help="follow log output")
    tail_p.add_argument("--stderr", action="store_true", help="tail stderr instead of stdout")
    tail_p.add_argument("--since", help="show logs only for runs updated within duration (e.g. 10m, 2h)")

    ls_p = sub.add_parser("ls", help="list runs")
    ls_p.add_argument("--long", "-l", action="store_true", help="show detailed columns")
    ls_p.add_argument(
        "--state",
        choices=["pending", "running", "success", "failed"],
        help="filter by run state",
    )

    gc_p = sub.add_parser("gc", help="cleanup run directories")
    gc_p.add_argument("--state", choices=["pending", "running", "success", "failed"], help="filter by run state")
    gc_p.add_argument("--older-than", help="delete runs older than duration, e.g. 7d")
    gc_p.add_argument("--dry-run", action="store_true", help="show what would be deleted")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    try:
        if args.command == "run":
            return run_local(args.config, watch=args.watch)
        if args.command == "submit":
            return submit_slurm(args.config)
        if args.command == "status":
            return show_status(args.run_id, json_output=args.json, long=args.long)
        if args.command == "cancel":
            return cancel_run(args.run_id, force=args.force)
        if args.command == "validate":
            return validate_manifest(args.config)
        if args.command == "tail":
            return tail_run_log(args.run_id, stderr=args.stderr, follow=args.follow, lines=args.lines, since=args.since)
        if args.command == "ls":
            return list_runs(state_filter=args.state, long=args.long)
        if args.command == "gc":
            return gc_runs(state=args.state, older_than=args.older_than, dry_run=args.dry_run)
    except OneRunError as exc:
        parser.exit(status=1, message=f"error: {exc}\n")

    parser.exit(status=2, message="unknown command\n")
    return 2
