#!/usr/bin/env python3
"""Reset audit harness.

Runs the latest registered version of each Farama environment and records whether it
replays bit for bit under the same seed. The orchestrator uses only the standard library
and starts a fresh subprocess, inside one of four virtual environments, for each phase.

Usage:
    python3 reset_audit.py --venvs .audit --out files [--limit N] [--workers 8]

Expected layout, one virtual environment per family of packages:
    .audit/core        gymnasium[classic-control,toy-text,box2d,mujoco]
    .audit/robotics    gymnasium-robotics
    .audit/minigrid    minigrid
    .audit/pettingzoo  pettingzoo[classic,butterfly,sisl]

Phases for each environment, each in its own subprocess with a 180 second timeout:
    probe   make the environment and record its spaces
    check   gymnasium env_checker on env.unwrapped, render check skipped, or pettingzoo api_test
    gen     sample and save 1,000 actions for each seed
    replay  replay the saved actions in two processes, PYTHONHASHSEED 1 and 2, for each seed
    perf    run serially after everything else: 3 runs of 5,000 steps, 20 resets

Nothing is estimated. A value that was not measured is recorded as a status, never a number.
"""
import argparse
import concurrent.futures as cf
import csv
import hashlib
import importlib
import json
import os
import pathlib
import pickle
import re
import statistics
import struct
import subprocess
import sys
import tempfile
import time
import traceback
import warnings

SEEDS = [0, 1, 12345]
REPLAY_STEPS = 1000
PERF_STEPS = 5000
PERF_RUNS = 3
RESET_SAMPLES = 20
TIMEOUT = 180
VENVS = ["core", "robotics", "minigrid", "pettingzoo"]
PETTINGZOO_FAMILIES = ["classic", "butterfly", "sisl"]
PACKAGES = {"core": "gymnasium", "robotics": "gymnasium-robotics", "minigrid": "minigrid", "pettingzoo": "pettingzoo"}
COMMON_ENV = {
    "SDL_VIDEODRIVER": "dummy", "SDL_AUDIODRIVER": "dummy", "PYGAME_HIDE_SUPPORT_PROMPT": "1",
    "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", "VECLIB_MAXIMUM_THREADS": "1", "MKL_NUM_THREADS": "1",
}


def clean(text, limit=240):
    """One line of error text, with no local paths."""
    line = re.sub(r"\x1b\[[0-9;]*m", "", str(text))
    line = re.sub(r"\s+", " ", line).strip()
    line = re.sub(r"/[^\s'\"]*site-packages/", "", line)
    line = re.sub(r"/Users/[^\s'\"]*/", "", line)
    return line[:limit]


# Worker side. Imports of numpy, gymnasium and pettingzoo happen only here.

def canon(value, sink):
    """Feed a canonical byte form of an observation to sink, recursing into dicts and tuples."""
    import numpy as np
    if isinstance(value, dict):
        for key in sorted(value, key=str):
            sink(b"k" + str(key).encode())
            canon(value[key], sink)
    elif isinstance(value, (tuple, list)):
        sink(b"s%d" % len(value))
        for item in value:
            canon(item, sink)
    elif isinstance(value, (str, bytes)):
        raw = value.encode() if isinstance(value, str) else value
        sink(b"t" + raw)
    elif value is None:
        sink(b"n")
    else:
        arr = np.asarray(value)
        sink(arr.dtype.str.encode())
        sink(repr(arr.shape).encode())
        sink(arr.tobytes())


def obs_digest(obs):
    h = hashlib.sha256()
    canon(obs, h.update)
    return h.hexdigest()


def link(prev, obs, reward, terminated, truncated, extra=b""):
    h = hashlib.sha256(prev)
    h.update(extra)
    canon(obs, h.update)
    h.update(struct.pack("<d", float(reward)))
    h.update(bytes([bool(terminated), bool(truncated)]))
    return h.digest()


def registry_of(venv):
    import gymnasium
    if venv == "robotics":
        import gymnasium_robotics
        gymnasium.register_envs(gymnasium_robotics)
    elif venv == "minigrid":
        import minigrid
        gymnasium.register_envs(minigrid)
    return gymnasium


def family_of(entry):
    name = entry if isinstance(entry, str) else getattr(entry, "__name__", repr(entry))
    name = name.split(":")[0]
    if name.startswith("gymnasium.envs."):
        return name.split(".")[2]
    if name.startswith("gymnasium_robotics.envs."):
        return name.split(".")[2]
    if name.startswith("minigrid."):
        return "babyai" if ".babyai" in name else "wfc" if ".wfc" in name else "minigrid"
    return "other"


def list_gym(venv):
    gym = registry_of(venv)
    prefix = {"core": "", "robotics": "gymnasium_robotics", "minigrid": "minigrid"}[venv]
    latest = {}
    for env_id, spec in gym.registry.items():
        key = (spec.namespace, spec.name)
        if key not in latest or (spec.version or 0) > (latest[key].version or 0):
            latest[key] = spec
    rows = []
    for spec in latest.values():
        entry = spec.entry_point if isinstance(spec.entry_point, str) else getattr(spec.entry_point, "__module__", "")
        if not entry.startswith(prefix or "gymnasium"):
            continue
        fam = family_of(spec.entry_point)
        if spec.namespace in ("phys2d", "tabular"):
            fam = spec.namespace
        if not isinstance(spec.entry_point, str):
            fam = "compatibility"
        rows.append({"id": spec.id, "family": fam, "max_episode_steps": spec.max_episode_steps})
    return rows


def list_pettingzoo():
    rows = []
    for fam in PETTINGZOO_FAMILIES:
        mod = importlib.import_module(f"pettingzoo.{fam}.all_modules")
        table = getattr(mod, f"{fam}_environments")
        for name, module in table.items():
            rows.append({"id": name, "family": fam, "module": module.__name__, "max_episode_steps": None})
    return rows


def make_gym(venv, env_id):
    return registry_of(venv).make(env_id, disable_env_checker=True)


def make_pz(module):
    return importlib.import_module(module).env()


def licence_of(package):
    from importlib import metadata
    dist = metadata.distribution(package)
    for f in dist.files or []:
        name = pathlib.PurePath(f).name.upper()
        if ".DIST-INFO" in str(f).upper() and (name.startswith("LICENSE") or name.startswith("LICENCE") or name.startswith("COPYING")):
            text = dist.read_text(str(f).split(".dist-info/", 1)[1]) or ""
            lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
            return {"licence": lines[0] if lines else "empty LICENSE file",
                    "licence_sha256": hashlib.sha256(text.encode()).hexdigest()}
    return {"licence": "no LICENSE file in installed package", "licence_sha256": ""}


def package_version(package):
    from importlib import metadata
    return metadata.version(package)


def phase_list(args):
    return list_pettingzoo() if args["venv"] == "pettingzoo" else list_gym(args["venv"])


def phase_probe(args):
    venv = args["venv"]
    if venv == "pettingzoo":
        env = make_pz(args["module"])
        env.reset(seed=0)
        agent = env.possible_agents[0]
        out = {"observation_space": str(env.observation_space(agent))[:400], "action_space": str(env.action_space(agent))[:400],
               "agents": len(env.possible_agents)}
    else:
        env = make_gym(venv, args["id"])
        out = {"observation_space": str(env.observation_space)[:400], "action_space": str(env.action_space)[:400]}
    package = PACKAGES[venv]
    out["package"] = package
    out["package_version"] = package_version(package)
    out.update(licence_of(package))
    env.close()
    return out


def phase_check(args):
    venv = args["venv"]
    caught = []
    with warnings.catch_warnings(record=True) as seen:
        warnings.simplefilter("always")
        try:
            if venv == "pettingzoo":
                from pettingzoo.test import api_test
                api_test(make_pz(args["module"]), num_cycles=1000, verbose_progress=False)
            else:
                from gymnasium.utils.env_checker import check_env
                env = make_gym(venv, args["id"])
                check_env(env.unwrapped, skip_render_check=True)
        except BaseException as ex:  # a failed check is a result, not a crash
            return {"check": "fail", "check_detail": clean(f"{type(ex).__name__}: {ex}")}
        caught = [clean(w.message) for w in seen]
    if caught:
        return {"check": "warnings", "check_detail": caught[0], "check_warnings": len(caught)}
    return {"check": "pass", "check_detail": ""}


def sample(space, mask=None):
    return space.sample(mask=mask) if mask is not None else space.sample()


def phase_gen(args):
    """Sample and save actions. PettingZoo actions depend on legal moves, so they are generated by playing with masks."""
    seed, path, venv = args["seed"], args["path"], args["venv"]
    actions = []
    if venv == "pettingzoo":
        env = make_pz(args["module"])
        env.reset(seed=seed)
        spaces = {a: env.action_space(a) for a in env.possible_agents}
        for s in spaces.values():
            s.seed(seed)
        for _ in range(REPLAY_STEPS):
            agent = env.agent_selection
            obs, _r, term, trunc, info = env.last()
            if term or trunc:
                action = None
            else:
                mask = obs.get("action_mask") if isinstance(obs, dict) else info.get("action_mask")
                action = sample(spaces[agent], mask)
            actions.append(action)
            env.step(action)
            if not env.agents:
                env.reset()
    else:
        env = make_gym(venv, args["id"])
        env.action_space.seed(seed)
        actions = [env.action_space.sample() for _ in range(REPLAY_STEPS)]
    with open(path, "wb") as fh:
        pickle.dump(actions, fh)
    env.close()
    return {"n": len(actions)}


def phase_replay(args):
    seed, venv = args["seed"], args["venv"]
    with open(args["path"], "rb") as fh:
        actions = pickle.load(fh)
    chain = []
    prev = hashlib.sha256(b"start").digest()
    if venv == "pettingzoo":
        env = make_pz(args["module"])
        env.reset(seed=seed)
        obs, r, term, trunc, _ = env.last()
        first = obs_digest(obs)
        prev = link(prev, obs, r, term, trunc, env.agent_selection.encode())
        chain.append(prev.hex())
        for action in actions:
            env.step(action)
            if not env.agents:
                prev = hashlib.sha256(prev + b"end").digest()
                env.reset()
            obs, r, term, trunc, _ = env.last()
            prev = link(prev, obs, r, term, trunc, env.agent_selection.encode())
            chain.append(prev.hex())
        env.reset(seed=seed)
        obs, *_ = env.last()
        again = obs_digest(obs)
    else:
        env = make_gym(venv, args["id"])
        obs, _ = env.reset(seed=seed)
        first = obs_digest(obs)
        prev = link(prev, obs, 0.0, False, False)
        chain.append(prev.hex())
        for action in actions:
            obs, r, term, trunc, _ = env.step(action)
            prev = link(prev, obs, r, term, trunc)
            chain.append(prev.hex())
            if term or trunc:
                obs, _ = env.reset()
                prev = link(prev, obs, 0.0, False, False, b"reset")
        obs, _ = env.reset(seed=seed)
        again = obs_digest(obs)
    env.close()
    return {"chain": chain, "reset_first": first, "reset_after": again}


def phase_perf(args):
    venv = args["venv"]
    out = {"sps": [], "reset_ms": []}

    def save():
        with open(args["result"], "w", encoding="utf-8") as fh:
            json.dump({"ok": True, "data": out}, fh)

    if venv == "pettingzoo":
        env = make_pz(args["module"])
        env.reset(seed=0)
        spaces = {a: env.action_space(a) for a in env.possible_agents}
        for s in spaces.values():
            s.seed(0)
        for _ in range(PERF_RUNS):
            started = time.perf_counter()
            for _s in range(PERF_STEPS):
                obs, _r, term, trunc, info = env.last()
                if term or trunc:
                    action = None
                else:
                    mask = obs.get("action_mask") if isinstance(obs, dict) else info.get("action_mask")
                    action = sample(spaces[env.agent_selection], mask)
                env.step(action)
                if not env.agents:
                    env.reset()
            out["sps"].append(PERF_STEPS / (time.perf_counter() - started))
            save()
        for i in range(RESET_SAMPLES):
            started = time.perf_counter()
            env.reset(seed=i)
            out["reset_ms"].append((time.perf_counter() - started) * 1000)
    else:
        env = make_gym(venv, args["id"])
        env.action_space.seed(0)
        actions = [env.action_space.sample() for _ in range(PERF_STEPS)]
        env.reset(seed=0)
        for _ in range(PERF_RUNS):
            started = time.perf_counter()
            for action in actions:
                _o, _r, term, trunc, _i = env.step(action)
                if term or trunc:
                    env.reset()
            out["sps"].append(PERF_STEPS / (time.perf_counter() - started))
            save()
        for i in range(RESET_SAMPLES):
            started = time.perf_counter()
            env.reset(seed=i)
            out["reset_ms"].append((time.perf_counter() - started) * 1000)
    env.close()
    return out


PHASES = {"list": phase_list, "probe": phase_probe, "check": phase_check, "gen": phase_gen, "replay": phase_replay, "perf": phase_perf}


def worker(phase, spec, result):
    warnings.simplefilter("ignore")
    args = json.loads(spec)
    args["result"] = result
    try:
        data = PHASES[phase](args)
        payload = {"ok": True, "data": data}
    except BaseException as ex:
        last = traceback.format_exception_only(type(ex), ex)[-1]
        frame = traceback.extract_tb(ex.__traceback__)[-1]
        payload = {"ok": False, "error_type": type(ex).__name__, "error": clean(f"{last.strip()} at {frame.filename}:{frame.lineno}")}
    with open(result, "w", encoding="utf-8") as fh:
        json.dump(payload, fh)
    os._exit(0)  # some native libraries hang on interpreter shutdown


# Orchestrator side, standard library only.

class Runner:
    def __init__(self, venvs, work, log):
        self.venvs, self.work, self.log = venvs, work, log
        self.counter = 0

    def call(self, phase, spec, hashseed=0, timeout=TIMEOUT):
        self.counter += 1
        result = self.work / f"{phase}-{os.getpid()}-{time.monotonic_ns()}-{self.counter}.json"
        python = self.venvs / spec["venv"] / "bin" / "python"
        env = {**os.environ, **COMMON_ENV, "PYTHONHASHSEED": str(hashseed)}
        cmd = [str(python), __file__, "--worker", phase, json.dumps(spec), str(result)]
        started = time.monotonic()
        try:
            proc = subprocess.run(cmd, env=env, capture_output=True, timeout=timeout)
            status = f"rc={proc.returncode}"
        except subprocess.TimeoutExpired:
            status = "timeout"
            proc = None
        seconds = time.monotonic() - started
        payload = None
        if result.exists():
            try:
                payload = json.loads(result.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                payload = None
            result.unlink(missing_ok=True)
        self.log(f"{spec.get('id', spec['venv'])}\t{phase}\thashseed={hashseed}\t{status}\t{seconds:.2f}s")
        if status == "timeout":
            return {"timeout": True, "partial": payload}
        if payload is None:
            stderr = proc.stderr.decode(errors="replace").strip().splitlines() if proc else []
            return {"ok": False, "error_type": "ProcessDied", "error": clean(stderr[-1] if stderr else f"process exited, {status}")}
        return payload


def not_run_reason(res):
    kind, text = res.get("error_type", ""), res.get("error", "")
    low = text.lower()
    if kind == "TypeError" and ("missing" in low or "required" in low):
        return "not run: needs arguments"
    if kind in ("ModuleNotFoundError", "ImportError", "DependencyNotInstalled") or "not installed" in low or "shimmy" in low or "jax" in low:
        return f"not run: dependency not installed ({text[:120]})"
    return f"not run: make failed ({text[:160]})"


def audit_env(runner, row):
    venv, spec = row["venv"], {k: row[k] for k in ("venv", "id", "module") if k in row}
    out = dict(row)
    out.update({"check": "", "check_detail": "", "replay": "", "first_diverging_step": "", "diverged_seeds": "",
                "reset_restore": "", "steps_per_second": "", "reset_latency_ms": "", "status": "", "error": ""})
    probe = runner.call("probe", spec)
    if probe.get("timeout"):
        out["status"] = "timeout"
        out["error"] = "probe timed out"
        return out
    if not probe.get("ok"):
        out["status"] = not_run_reason(probe)
        out["error"] = probe.get("error", "")
        return out
    out.update(probe["data"])

    chk = runner.call("check", spec)
    if chk.get("timeout"):
        out["check"] = "timeout"
    elif chk.get("ok"):
        out.update(chk["data"])
    else:
        out["check"], out["check_detail"] = "fail", chk.get("error", "")

    diverged, first_step, restored, problems = [], None, [], []
    timed_out = False
    for seed in SEEDS:
        path = str(runner.work / f"actions-{row['slug']}-{seed}.pkl")
        gen = runner.call("gen", {**spec, "seed": seed, "path": path})
        if gen.get("timeout"):
            timed_out = True
            break
        if not gen.get("ok"):
            problems.append(f"action generation failed: {gen.get('error', '')}")
            break
        runs = []
        for hashseed in (1, 2):
            res = runner.call("replay", {**spec, "seed": seed, "path": path}, hashseed=hashseed)
            if res.get("timeout"):
                timed_out = True
                break
            if not res.get("ok"):
                problems.append(f"replay failed: {res.get('error', '')}")
                break
            runs.append(res["data"])
        pathlib.Path(path).unlink(missing_ok=True)
        if timed_out or len(runs) < 2:
            break
        a, b = runs
        if a["chain"] != b["chain"]:
            diverged.append(seed)
            step = next((i for i, (x, y) in enumerate(zip(a["chain"], b["chain"])) if x != y), min(len(a["chain"]), len(b["chain"])))
            first_step = step if first_step is None else min(first_step, step)
        restored.append(a["reset_first"] == a["reset_after"] and b["reset_first"] == b["reset_after"])

    if timed_out:
        out["replay"] = "timeout"
    elif problems:
        out["replay"], out["error"] = "error", problems[0]
    else:
        out["replay"] = "diverged" if diverged else "identical"
        out["reset_restore"] = "restored" if all(restored) else "differs"
    if diverged:
        out["first_diverging_step"] = first_step
        out["diverged_seeds"] = " ".join(map(str, diverged))

    if out["replay"] == "timeout" or out["check"] == "timeout":
        out["status"] = "timeout"
    elif out["replay"] == "error":
        out["status"] = "check_failed"
    elif out["replay"] == "diverged":
        out["status"] = "diverged"
    elif out["check"] == "fail":
        out["status"] = "check_failed"
    else:
        out["status"] = "identical"
    return out


def perf_env(runner, row):
    spec = {k: row[k] for k in ("venv", "id", "module") if k in row}
    res = runner.call("perf", spec)
    data = None
    if res.get("timeout"):
        partial = res.get("partial") or {}
        data = partial.get("data") if isinstance(partial, dict) else None
        row["perf_note"] = "timeout"
    elif res.get("ok"):
        data = res["data"]
    else:
        row["perf_note"] = f"failed: {res.get('error', '')}"
    if data and len(data.get("sps", [])) == PERF_RUNS:
        row["steps_per_second"] = round(statistics.median(data["sps"]), 1)
    else:
        row["steps_per_second"] = "not measured"
    if data and len(data.get("reset_ms", [])) == RESET_SAMPLES:
        row["reset_latency_ms"] = round(statistics.median(data["reset_ms"]), 3)
    else:
        row["reset_latency_ms"] = "not measured"


def slug(text):
    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--venvs", default=".audit")
    ap.add_argument("--out", default="files")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0, help="audit only the first N of each venv, for a trial")
    a = ap.parse_args()
    venvs, outdir = pathlib.Path(a.venvs).resolve(), pathlib.Path(a.out).resolve()
    outdir.mkdir(parents=True, exist_ok=True)
    log_lines = []

    def log(line):
        log_lines.append(f"{time.strftime('%Y-%m-%dT%H:%M:%S')}\t{line}")

    started = time.strftime("%Y-%m-%dT%H:%M:%S")
    with tempfile.TemporaryDirectory(dir=venvs) as tmp:
        runner = Runner(venvs, pathlib.Path(tmp), log)
        rows, skipped = [], []
        for venv in VENVS:
            listed = runner.call("list", {"venv": venv})
            if not listed.get("ok"):
                skipped.append({"venv": venv, "status": f"not run: {listed.get('error', 'could not list environments')}"})
                continue
            items = listed["data"][: a.limit] if a.limit else listed["data"]
            rows += [{**item, "venv": venv, "slug": slug(item["id"])} for item in items]
        seen = {}
        for row in rows:
            seen[row["slug"]] = seen.get(row["slug"], 0) + 1
            if seen[row["slug"]] > 1:
                row["slug"] = f"{row['slug']}-{row['venv']}"
        with cf.ThreadPoolExecutor(a.workers) as pool:
            done = list(pool.map(lambda r: audit_env(runner, r), rows))
        for row in done:  # serial, so timings do not compete with the replay stage
            if row["status"].startswith("not run") or row["replay"] in ("error", ""):
                row["steps_per_second"] = row["reset_latency_ms"] = "not measured"
                continue
            perf_env(runner, row)
    done.sort(key=lambda r: (r["venv"], r["family"], r["id"].lower()))
    meta = {"started": started, "finished": time.strftime("%Y-%m-%dT%H:%M:%S"), "seeds": SEEDS, "replay_steps": REPLAY_STEPS,
            "perf_steps": PERF_STEPS, "perf_runs": PERF_RUNS, "reset_samples": RESET_SAMPLES, "timeout_seconds": TIMEOUT,
            "python": "3.13", "skipped_venvs": skipped}
    (outdir / "results.json").write_text(json.dumps({"meta": meta, "environments": done}, indent=1, ensure_ascii=False) + "\n", encoding="utf-8")
    fields = ["id", "slug", "venv", "family", "package", "package_version", "licence", "observation_space", "action_space", "max_episode_steps",
              "check", "check_detail", "replay", "first_diverging_step", "diverged_seeds", "reset_restore", "steps_per_second", "reset_latency_ms", "status", "error"]
    with open(outdir / "results.csv", "w", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=fields, extrasaction="ignore")
        w.writeheader()
        w.writerows(done)
    (outdir / "run-log.txt").write_text("\n".join(log_lines) + "\n", encoding="utf-8")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--worker":
        worker(sys.argv[2], sys.argv[3], sys.argv[4])
    else:
        main()