from __future__ import annotations import argparseimport subprocessimport sysimport tarfileimport timefrom pathlib import Path from . import run_index, sourcefrom .common import (    default_namespace,    deployed_function,    print_json,    print_table,    resolve_run_location,    run_command,    safe_extract_archive,) RECENT_RUN_COLUMNS = [    ("created_at", "created"),    ("run_id", "run"),    ("experiment_name", "experiment"),    ("namespace", "namespace"),    ("agent_model", "model"),    ("judge_preset", "judge"),    ("build_id", "build"),    ("suite_id", "suite"),]REMOTE_RUN_COLUMNS = [    ("updated_at", "updated"),    ("namespace", "namespace"),    ("experiment_name", "experiment"),    ("run_id", "run"),    ("status", "status"),    ("agent_model", "model"),    ("judge_preset", "judge"),    ("build_id", "build"),    ("suite_id", "suite"),]BUILD_COLUMNS = [    ("built_at_utc", "built"),    ("build_id", "build"),    ("ready", "ready"),    ("base_sha", "base"),    ("patch_sha256", "patch"),]SUITE_COLUMNS = [    ("part", "part"),    ("experiment_name", "experiment"),    ("run_id", "run"),    ("status", "status"),    ("updated_at", "updated"),    ("trial_count", "trials"),]  def volume_get(    args: argparse.Namespace, remote_path: str, local_path: str) -> subprocess.CompletedProcess[str]:    return run_command(        ["modal", "volume", "get", args.volume, remote_path, local_path],        capture=local_path == "-",    )  def read_volume_text(args: argparse.Namespace, remote_path: str) -> str:    result = volume_get(args, remote_path, "-")    return result.stdout  def build_ready_on_volume(args: argparse.Namespace, build_id: str) -> bool:    try:        read_volume_text(args, f"builds/{build_id}/READY")        return True    except (subprocess.CalledProcessError, OSError):        return False  def ensure_volume_exists(args: argparse.Namespace) -> int:    print(f"\nEnsuring Modal volume '{args.volume}' exists...")    try:        result = subprocess.run(            ["modal", "volume", "create", args.volume],            text=True,            capture_output=True,            check=True,        )        if result.stdout:            print(result.stdout, end="")    except subprocess.CalledProcessError as error:        output = (error.stdout or "") + (error.stderr or "")        if "already exists" in output.lower():            print(f"Volume '{args.volume}' already exists; continuing.")        else:            if output:                print(output, end="", file=sys.stderr)            return error.returncode    return 0  def print_rows(    rows: list[dict], columns: list[tuple[str, str]], *, as_json: bool = False) -> None:    if as_json:        print_json(rows)    else:        print_table(rows, columns)  def run_remote_prefix(args: argparse.Namespace, run_id: str) -> str:    namespace, experiment_name = resolve_run_location(args, run_id)    return f"runs/{namespace}/{experiment_name}/{run_id}"  def read_run_file(args: argparse.Namespace, file_name: str) -> str:    run_id = resolve_run_id(args)    return read_volume_text(args, f"{run_remote_prefix(args, run_id)}/{file_name}")  def suite_rows(args: argparse.Namespace) -> list[dict]:    namespace = default_namespace(args)    return deployed_function(args, "suite_status").remote(namespace, args.suite_id)  def suite_member_args(args: argparse.Namespace, row: dict) -> argparse.Namespace:    run_args = argparse.Namespace(**vars(args))    run_args.experiment_name = row["experiment_name"]    run_args.run_id = row["run_id"]    return run_args  def resolve_run_id(args: argparse.Namespace) -> str:    """Return the run id to act on, defaulting to the most recent local run.     Commands like `status`/`logs` accept no run id at all, in which case we use    the most recently launched run recorded in the local index so a bare    `zed-eval status` answers "how's my latest run doing?".    """    run_id = getattr(args, "run_id", None)    if run_id:        return run_id    entry = run_index.most_recent()    if not entry:        raise ValueError(            "no run id given and the local run index is empty. Launch a run "            "first, or pass a run id explicitly."        )    args.run_id = entry["run_id"]    print(        f"(no run id given; using most recent: {entry['run_id']} / "        f"{entry.get('experiment_name')})",        file=sys.stderr,    )    return entry["run_id"]  def command_runs(args: argparse.Namespace) -> int:    entries = run_index.recent(args.limit)    if getattr(args, "json", False):        print_json(entries)        return 0    if not entries:        print(            "No runs recorded locally yet. Launch a run, or use "            "`zed-eval list --details` to query runs on the volume."        )        return 0    print_table(entries, RECENT_RUN_COLUMNS)    return 0  def command_list(args: argparse.Namespace) -> int:    namespace = None if args.all_namespaces else default_namespace(args)    experiment_name = getattr(args, "experiment_name", None)    if experiment_name:        experiment_name = source.sanitize_namespace(experiment_name)    if args.details or args.json:        rows = deployed_function(args, "list_runs").remote(            namespace, experiment_name, args.limit        )        print_rows(rows, REMOTE_RUN_COLUMNS, as_json=args.json)        return 0    remote_path = (        f"runs/{namespace}/{experiment_name}/"        if experiment_name        else f"runs/{namespace}/"    )    run_command(["modal", "volume", "ls", args.volume, remote_path])    return 0  def command_builds(args: argparse.Namespace) -> int:    if args.details or args.json:        rows = deployed_function(args, "list_builds").remote(args.limit)        print_rows(rows, BUILD_COLUMNS, as_json=args.json)        return 0    run_command(["modal", "volume", "ls", args.volume, "builds/"])    return 0  def command_status(args: argparse.Namespace) -> int:    # One-shot status ping: the controller's state.json only carries a coarse    # status (pending -> running -> completed/failed), so there is no per-trial    # progress worth following. Print it once and return.    print(read_run_file(args, "state.json"), end="")    return 0  def command_logs(args: argparse.Namespace) -> int:    print(read_run_file(args, "controller.log"), end="")    return 0  def command_fetch(args: argparse.Namespace) -> int:    jobs_dir = Path(args.jobs_dir).expanduser()    jobs_dir.mkdir(parents=True, exist_ok=True)    temporary_archive = jobs_dir / f".{args.run_id}.tar.gz"    remote_path = f"{run_remote_prefix(args, args.run_id)}/harbor-job.tar.gz"    print(f"Fetching {args.volume}:/{remote_path} -> {temporary_archive}")    volume_get(args, remote_path, str(temporary_archive))    with tarfile.open(temporary_archive, "r:gz") as archive:        safe_extract_archive(archive, jobs_dir)    temporary_archive.unlink(missing_ok=True)    print(f"Extracted Harbor job under {jobs_dir / args.run_id}")    return 0  def command_suite_status(args: argparse.Namespace) -> int:    print_rows(suite_rows(args), SUITE_COLUMNS, as_json=args.json)    return 0  def command_suite_logs(args: argparse.Namespace) -> int:    rows = suite_rows(args)    while True:        for row in rows:            print(                f"\n=== {row['part']} / {row['experiment_name']} / {row['run_id']} ==="            )            run_args = suite_member_args(args, row)            run_args.follow = False            command_logs(run_args)        if not args.follow:            return 0        time.sleep(args.interval)  def command_suite_fetch(args: argparse.Namespace) -> int:    result = 0    for row in suite_rows(args):        result = command_fetch(suite_member_args(args, row)) or result    return result