from __future__ import annotations

import argparse
import os

from one_run import __version__
from one_run.constants import AUTOMATION_ENV_VAR
from one_run.commands.init_cmd import init_listener, init_manifest
from one_run.commands.ls_cmd import list_runs
from one_run.commands.ps_cmd import list_processes
from one_run.commands.cancel_cmd import cancel_run
from one_run.commands.cron_cmd import install_cron, remove_cron, show_cron
from one_run.commands.listen_cmd import listen_events
from one_run.commands.rm_cmd import remove_run, remove_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__}")
    parser.add_argument(
        "--automation",
        action="store_true",
        help=f"prefer machine-readable output defaults (also enabled by {AUTOMATION_ENV_VAR}=1)",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    init_p = sub.add_parser("init", help="interactive config generators")
    init_sub = init_p.add_subparsers(dest="init_target", required=True)
    init_manifest_p = init_sub.add_parser("manifest", help="generate a manifest yaml interactively")
    init_manifest_p.add_argument("-o", "--output", default=None, help="output path (default: ./manifests/<experiment>.manifest.yaml)")
    init_manifest_p.add_argument("--force", action="store_true", help="overwrite output file if it exists")
    init_listener_p = init_sub.add_parser("listener", help="generate a listener yaml interactively")
    init_listener_p.add_argument("-o", "--output", default=None, help="output path (default: ./listeners/<name>.listener.yaml)")
    init_listener_p.add_argument("--force", action="store_true", help="overwrite output file if it exists")

    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")
    run_p.add_argument("--wait-for", action="append", help="dependency run_id to wait for before start (repeatable)")
    run_p.add_argument("--wait-timeout", help="max dependency wait duration, e.g. 30m")
    run_p.add_argument("--wait-interval", type=float, help="dependency poll interval seconds")
    run_p.add_argument("--json", action="store_true", help="print machine-readable result json")

    submit_p = sub.add_parser("submit", help="submit a slurm run")
    submit_p.add_argument("config", help="path to manifest yaml")
    submit_p.add_argument("--wait-for", action="append", help="dependency run_id to wait for before submit (repeatable)")
    submit_p.add_argument("--wait-timeout", help="max dependency wait duration, e.g. 30m")
    submit_p.add_argument("--wait-interval", type=float, help="dependency poll interval seconds")
    submit_p.add_argument("--json", action="store_true", help="print machine-readable result json")

    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")
    status_p.add_argument("--wait", action="store_true", help="wait until run reaches terminal state")
    status_p.add_argument("--timeout", help="max wait duration, e.g. 30m")
    status_p.add_argument("--interval", type=float, default=2.0, help="poll interval in seconds (default: 2)")

    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")
    cancel_p.add_argument("--json", action="store_true", help="print machine-readable result json")

    rm_p = sub.add_parser("rm", help="remove runs (single run_id or filtered cleanup)")
    rm_p.add_argument("run_id", nargs="?", help="run identifier")
    rm_p.add_argument("--force", action="store_true", help="force kill when stopping active run")
    rm_p.add_argument("--state", choices=["pending", "running", "success", "failed"], help="filter by run state (bulk mode)")
    rm_p.add_argument("--older-than", help="delete runs older than duration, e.g. 7d (bulk mode)")
    rm_p.add_argument("--dry-run", action="store_true", help="show what would be deleted (bulk mode)")
    rm_p.add_argument("--json", action="store_true", help="print machine-readable result json")

    validate_p = sub.add_parser("validate", help="validate manifest only")
    validate_p.add_argument("config", help="path to manifest yaml")
    validate_p.add_argument("--json", action="store_true", help="print machine-readable result json")

    cron_p = sub.add_parser("cron", help="manage crontab entries from manifest schedule")
    cron_sub = cron_p.add_subparsers(dest="cron_target", required=True)
    cron_install_p = cron_sub.add_parser("install", help="install or update crontab entry for manifest")
    cron_install_p.add_argument("config", help="path to manifest yaml")
    cron_remove_p = cron_sub.add_parser("remove", help="remove crontab entry for manifest")
    cron_remove_p.add_argument("config", help="path to manifest yaml")
    cron_show_p = cron_sub.add_parser("show", help="show crontab entry for manifest")
    cron_show_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)")

    listen_p = sub.add_parser("listen", help="listen to lifecycle events (all runs or selected run_ids)", add_help=False)
    listen_p.add_argument("--help", action="help", help="show this help message and exit")
    listen_p.add_argument("run_ids", nargs="*", help="optional run identifiers")
    listen_p.add_argument("-m", "--manifest", help="manifest with listeners config")
    listen_p.add_argument("--output", choices=["table", "jsonl"], default="table", help="output format (default: table)")
    listen_p.add_argument("--jsonl", action="store_true", help="alias for --output jsonl")
    listen_p.add_argument("--event", action="append", help="filter by event name (repeatable)")
    listen_p.add_argument("--state", action="append", help="filter by event state (repeatable)")
    listen_p.add_argument("--tag", action="append", help="filter by run tag (repeatable)")
    listen_p.add_argument("--handler", help="python callback in module:function format")
    listen_mode_group = listen_p.add_mutually_exclusive_group()
    listen_mode_group.add_argument("--once", action="store_true", help="read currently available events once and exit")
    listen_mode_group.add_argument("-w", "--watch", action="store_true", help="watch event stream continuously")
    listen_p.add_argument("-h", "--history", action="store_true", help="include historical events that happened before listener start")
    listen_p.add_argument("--since", help="history window, e.g. 10m or 2h (requires --history)")
    listen_p.add_argument("--interval", type=float, default=None, help="poll interval in seconds (default: manifest interval_sec or 1.0)")

    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",
    )
    ls_p.add_argument("--tag", action="append", help="filter by tag (can be used multiple times)")
    ls_p.add_argument("--group-by-tag", action="store_true", help="group output by tag")
    ls_p.add_argument("--json", action="store_true", help="print machine-readable json")

    ps_p = sub.add_parser("ps", help="list active runs")
    ps_p.add_argument(
        "--state",
        choices=["pending", "running"],
        help="filter by active state",
    )
    ps_p.add_argument("--tag", action="append", help="filter by tag (can be used multiple times)")
    ps_p.add_argument("--json", action="store_true", help="print machine-readable json")
    ps_p.add_argument("-w", "--watch", action="store_true", help="refresh process table continuously")
    ps_p.add_argument("--interval", type=float, default=2.0, help="watch refresh interval in seconds (default: 2)")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    try:
        import argcomplete  # type: ignore

        argcomplete.autocomplete(parser)
    except Exception:
        pass
    args = parser.parse_args(argv)
    automation_mode = args.automation or os.getenv(AUTOMATION_ENV_VAR) == "1"

    try:
        if args.command == "init":
            if args.init_target == "manifest":
                return init_manifest(output=args.output, force=args.force)
            if args.init_target == "listener":
                return init_listener(output=args.output, force=args.force)
        if args.command == "run":
            return run_local(
                args.config,
                watch=args.watch,
                json_output=(args.json or automation_mode),
                wait_for=args.wait_for,
                wait_timeout=args.wait_timeout,
                wait_interval=args.wait_interval,
            )
        if args.command == "submit":
            return submit_slurm(
                args.config,
                json_output=(args.json or automation_mode),
                wait_for=args.wait_for,
                wait_timeout=args.wait_timeout,
                wait_interval=args.wait_interval,
            )
        if args.command == "status":
            return show_status(
                args.run_id,
                json_output=(args.json or automation_mode),
                long=args.long,
                wait=args.wait,
                timeout=args.timeout,
                interval_sec=args.interval,
            )
        if args.command == "cancel":
            return cancel_run(args.run_id, force=args.force, json_output=(args.json or automation_mode))
        if args.command == "rm":
            if args.run_id:
                if args.state or args.older_than or args.dry_run:
                    raise OneRunError("rm <run_id> cannot be combined with --state/--older-than/--dry-run")
                return remove_run(args.run_id, force=args.force, json_output=(args.json or automation_mode))
            return remove_runs(
                state=args.state,
                older_than=args.older_than,
                dry_run=args.dry_run,
                json_output=(args.json or automation_mode),
            )
        if args.command == "validate":
            return validate_manifest(args.config, json_output=(args.json or automation_mode))
        if args.command == "cron":
            if args.cron_target == "install":
                return install_cron(args.config)
            if args.cron_target == "remove":
                return remove_cron(args.config)
            if args.cron_target == "show":
                return show_cron(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 == "listen":
            watch_mode = args.watch or not args.once
            return listen_events(
                run_ids=args.run_ids,
                watch=watch_mode,
                jsonl=args.jsonl,
                output=args.output,
                interval_sec=args.interval,
                events=args.event,
                states=args.state,
                tags=args.tag,
                handler=args.handler,
                manifest_path=args.manifest,
                once=args.once,
                history=args.history,
                since=args.since,
            )
        if args.command == "ls":
            return list_runs(
                state_filter=args.state,
                long=args.long,
                tags=args.tag,
                group_by_tag=args.group_by_tag,
                json_output=(args.json or automation_mode),
            )
        if args.command == "ps":
            return list_processes(
                state_filter=args.state,
                tags=args.tag,
                json_output=(args.json or automation_mode),
                watch=args.watch,
                interval_sec=args.interval,
            )
    except OneRunError as exc:
        parser.exit(status=1, message=f"error: {exc}\n")

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