#!/usr/bin/env python3
import os
import argparse
from typing import Optional

SENTINEL = "View project at"
SKIP_DIRS = {".git", ".hg", ".svn", "__pycache__", ".mamba", ".conda"}

def extract_run_name(log_path: str) -> Optional[str]:
    """
    Return the run name if we find a line that:
      - starts with 'wandb:' (after leading whitespace is stripped)
      - whose second-to-last word is 'run' and the last word is the run name
    and this occurs before the first line containing SENTINEL.
    Otherwise return None.
    """
    try:
        with open(log_path, "r", errors="ignore") as f:
            for raw in f:
                if SENTINEL in raw:
                    return None
                line = raw.strip()
                if not line.startswith("wandb: "):
                    continue
                toks = line.split()
                if len(toks) >= 2 and toks[-2] == "run":
                    return toks[-1]  # the run name (last token)
    except OSError:
        return None
    return None

def list_runs(root: str, followlinks: bool = False, do_rename=False):
    for dirpath, dirs, files in os.walk(root, followlinks=followlinks):
        # prune junk dirs
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        if "run.log" in files:
            log_path = os.path.join(dirpath, "run.log")
            name = extract_run_name(log_path)
            new_dir_path = dirpath
            if name:
                if name not in dirpath:
                    new_dir_path = f"{dirpath}-{name}"
            print(f"{name if name else '<unknown>'}\t{new_dir_path}")
            
            if do_rename and new_dir_path != dirpath:
                parent = os.path.dirname(dirpath)
                # resolve absolute path for safety
                abs_old = os.path.abspath(dirpath)
                abs_new = os.path.abspath(new_dir_path)

                if os.path.exists(abs_new):
                    print(f"⚠️ Target {abs_new} already exists, skipping rename.")
                else:
                    print(f"Renaming {abs_old} → {abs_new}")
                    os.rename(abs_old, abs_new)

def main():
    ap = argparse.ArgumentParser(
        description="List W&B run names by parsing lines that start with 'wandb:' and end with 'run <name>'."
    )
    ap.add_argument("root", help="Root directory to search")
    ap.add_argument("--followlinks", action="store_true", help="Follow symlinks while walking")
    args = ap.parse_args()
    list_runs(args.root, followlinks=args.followlinks)

if __name__ == "__main__":
    main()