import argparse
import json
import re
import subprocess
from pathlib import Path
from typing import Dict, List, Optional


def _active_worker_dirs(result_dir: Path) -> Dict[str, int]:
    command = [
        "pgrep",
        "-af",
        f"{result_dir}/worker_",
    ]
    try:
        output = subprocess.check_output(command, text=True)
    except subprocess.CalledProcessError:
        return {}

    active: Dict[str, int] = {}
    for line in output.splitlines():
        pid_match = re.match(r"\s*(\d+)\s+", line)
        worker_match = re.search(r"(worker_\d+)", line)
        if not pid_match or not worker_match:
            continue
        active[worker_match.group(1)] = int(pid_match.group(1))
    return active


def _selected_episode_indices(worker_dir: Path) -> Optional[List[int]]:
    templates_path = worker_dir.joinpath("templates.json")
    if not templates_path.exists():
        worker_match = re.fullmatch(r"worker_(\d+)", worker_dir.name)
        if worker_match:
            return [int(worker_match.group(1))]
        return None
    with templates_path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    selected = payload.get("selected_episode_indices")
    if isinstance(selected, list):
        return [int(index) for index in selected]
    episode_offset = payload.get("episode_offset")
    if episode_offset is not None:
        return [int(episode_offset)]
    return None


def _worker_status(worker_dir: Path, active_workers: Dict[str, int]) -> Dict[str, object]:
    name = worker_dir.name
    summary_path = worker_dir.joinpath("summary.json")
    worker_log = worker_dir.joinpath("worker.log")
    selected_episode_indices = _selected_episode_indices(worker_dir)
    log_text = worker_log.read_text(encoding="utf-8", errors="ignore") if worker_log.exists() else ""

    if summary_path.exists():
        status = "completed"
    elif name in active_workers:
        status = "running"
    elif re.search(r"Traceback|RuntimeError|signal 11", log_text):
        status = "crashed"
    elif worker_log.exists():
        status = "stalled"
    else:
        status = "empty"

    return {
        "worker": name,
        "status": status,
        "pid": active_workers.get(name),
        "selected_episode_indices": selected_episode_indices,
        "summary_path": str(summary_path) if summary_path.exists() else None,
        "worker_log": str(worker_log) if worker_log.exists() else None,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--result-dir", required=True)
    args = parser.parse_args()

    result_dir = Path(args.result_dir)
    worker_dirs = sorted(path for path in result_dir.glob("worker_*") if path.is_dir())
    active_workers = _active_worker_dirs(result_dir)
    worker_rows = [_worker_status(worker_dir, active_workers) for worker_dir in worker_dirs]

    rerun_episode_indices: List[int] = []
    for row in worker_rows:
        if row["status"] != "crashed":
            continue
        selected = row.get("selected_episode_indices") or []
        rerun_episode_indices.extend(int(index) for index in selected)

    counts: Dict[str, int] = {}
    for row in worker_rows:
        status = str(row["status"])
        counts[status] = counts.get(status, 0) + 1

    payload = {
        "result_dir": str(result_dir),
        "counts": counts,
        "rerun_episode_indices": sorted(set(rerun_episode_indices)),
        "workers": worker_rows,
    }
    print(json.dumps(payload, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())