"""Evaluate every available actor checkpoint in the final training window.

Each training job is one subprocess containing all of its final-window models.
This avoids starting thousands of Python processes while preserving one result
record per actor checkpoint.  Planned checkpoints that were not frozen by the
rollout-granular trainer remain explicit missing work in ``status.json``.
"""

import argparse
import json
import os
import re
import subprocess
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

import yaml


REPO_ROOT = Path(__file__).resolve().parents[2]


def _now():
    return datetime.now().astimezone().isoformat(timespec="seconds")


def _load_json(path):
    try:
        return json.loads(Path(path).read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        return None


def _atomic_json(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    temporary.replace(path)


def _atomic_yaml(path, payload):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")
    temporary.write_text(
        yaml.safe_dump(payload, sort_keys=False, allow_unicode=True),
        encoding="utf-8",
    )
    temporary.replace(path)


def _parse_named_paths(values, multiple=False):
    parsed = {}
    for value in values:
        if "=" not in value:
            raise ValueError(f"expected NAME=PATH, got {value!r}")
        name, path = value.split("=", 1)
        if multiple:
            parsed.setdefault(name, []).append(Path(path).resolve())
        else:
            parsed[name] = Path(path).resolve()
    return parsed


class FinalWindowQueue:
    def __init__(
        self,
        output_dir,
        stages,
        recovery_stages,
        cuda,
        max_workers,
        fraction,
        eval_episodes,
        eval_seed,
        job_pattern=None,
    ):
        self.root = Path(output_dir).resolve()
        self.status_path = self.root / "status.json"
        self.stages = stages
        self.recovery_stages = recovery_stages
        self.cuda = int(cuda)
        self.max_workers = int(max_workers)
        self.fraction = float(fraction)
        self.eval_episodes = int(eval_episodes)
        self.eval_seed = int(eval_seed)
        self.job_pattern = re.compile(job_pattern) if job_pattern else None
        self.lock = threading.Lock()
        self.state = _load_json(self.status_path) or {
            "name": "final-window-v1",
            "created_at": _now(),
            "state": "initializing",
            "phase": "final_window",
            "cuda": self.cuda,
            "max_workers": self.max_workers,
            "fraction": self.fraction,
            "eval_episodes_per_checkpoint": self.eval_episodes,
            "eval_seed": self.eval_seed,
            "tasks": {},
        }
        # Requeue interrupted/failed per-job suites.  A completed suite is
        # immutable and will be reused on process restart.
        for task in self.state.get("tasks", {}).values():
            if task.get("state") in ("pending", "running", "failed"):
                task.update(state="pending")
                for key in ("pid", "return_code", "finished_at", "started_at"):
                    task.pop(key, None)
        self._write_state()

    def _write_state(self):
        tasks = list(self.state.get("tasks", {}).values())
        planned = sum(int(task.get("planned_units", 0)) for task in tasks)
        submitted = sum(int(task.get("submitted_units", 0)) for task in tasks)
        completed = sum(int(task.get("completed_units", 0)) for task in tasks)
        self.state.update(
            updated_at=_now(),
            counts={
                key: sum(task.get("state") == key for task in tasks)
                for key in ("pending", "running", "completed", "failed")
            },
            planned_units=planned,
            submitted_units=submitted,
            completed_units=completed,
            missing_units=max(0, planned - submitted),
        )
        _atomic_json(self.status_path, self.state)

    def _stage_jobs(self, stage):
        status = _load_json(self.stages[stage] / "status.json") or {}
        jobs = dict(status.get("jobs", {}))
        for recovery_dir in self.recovery_stages.get(stage, []):
            recovery = _load_json(recovery_dir / "status.json") or {}
            for name, job in recovery.get("jobs", {}).items():
                if job.get("state") == "completed":
                    jobs[name] = job
        return jobs

    def _discover(self):
        discovered = []
        for stage in self.stages:
            for job_name, job in self._stage_jobs(stage).items():
                if job.get("state") != "completed":
                    continue
                if self.job_pattern and not self.job_pattern.search(job_name):
                    continue
                config_path = Path(job["config"])
                config = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}
                max_episode = int(config["max_train_episodes"])
                first_episode = int(max_episode * (1.0 - self.fraction)) + 1
                planned_episodes = list(range(first_episode, max_episode + 1))
                models_dir = Path(job["run_dir"]) / "models"
                models = []
                missing = []
                for episode in planned_episodes:
                    checkpoint = (
                        models_dir / "actor.pt"
                        if episode == max_episode
                        else models_dir / f"actor_ep{episode:06d}.pt"
                    )
                    if checkpoint.is_file():
                        models.append(
                            {
                                "name": f"ep{episode:06d}",
                                "checkpoint": str(checkpoint),
                                "overrides": {},
                            }
                        )
                    else:
                        missing.append(episode)
                task_id = f"{stage}:{job_name}"
                task = self.state["tasks"].setdefault(task_id, {})
                task.update(
                    task_id=task_id,
                    stage=stage,
                    job_name=job_name,
                    config_path=str(config_path),
                    run_dir=str(job["run_dir"]),
                    first_episode=first_episode,
                    last_episode=max_episode,
                    planned_checkpoints=len(planned_episodes),
                    available_checkpoints=len(models),
                    missing_checkpoint_episodes=missing,
                    planned_units=len(planned_episodes) * self.eval_episodes,
                    submitted_units=len(models) * self.eval_episodes,
                    models=models,
                )
                task.setdefault("state", "pending")
                task.setdefault("completed_units", 0)
                discovered.append(task_id)
        self.state["state"] = "running"
        self._write_state()
        return discovered

    def _prepare(self, task):
        task_dir = self.root / task["stage"] / task["job_name"]
        config = yaml.safe_load(
            Path(task["config_path"]).read_text(encoding="utf-8")
        ) or {}
        progress_path = task_dir / "progress.json"
        config.update(
            cuda=self.cuda,
            use_eval=True,
            model_dir=None,
            n_eval_rollout_threads=1,
            eval_episodes=self.eval_episodes,
            eval_seed=self.eval_seed,
            eval_reset_mode="full_random",
            eval_deterministic=True,
            progress_file=str(progress_path),
        )
        config_path = task_dir / "config.yaml"
        manifest_path = task_dir / "manifest.yaml"
        _atomic_yaml(config_path, config)
        _atomic_yaml(
            manifest_path,
            {
                "name": f"final-window-{task['stage']}-{task['job_name']}",
                "output_root": str(task_dir / "results"),
                "models": task["models"],
            },
        )
        return task_dir, config_path, manifest_path, progress_path

    def _update_progress(self, task_id, log_path, progress_path):
        try:
            started_models = log_path.read_text(
                encoding="utf-8", errors="replace"
            ).count("Evaluating ")
        except OSError:
            started_models = 0
        current_episodes = 0
        progress = _load_json(progress_path)
        if progress:
            current_episodes = min(
                self.eval_episodes,
                max(0, int(progress.get("completed_episodes", 0))),
            )
        completed_units = (
            max(0, started_models - 1) * self.eval_episodes + current_episodes
        )
        with self.lock:
            task = self.state["tasks"][task_id]
            task["completed_units"] = min(
                int(task["submitted_units"]), completed_units
            )
            task["heartbeat_at"] = _now()
            self._write_state()

    def _run_task(self, task_id):
        with self.lock:
            task = self.state["tasks"][task_id]
            task.update(state="running", started_at=_now())
            self._write_state()
            task = dict(task)
        task_dir, config_path, manifest_path, progress_path = self._prepare(task)
        log_path = task_dir / "eval.log"
        command = [
            sys.executable,
            str(REPO_ROOT / "scripts" / "gpu_thermal_guard.py"),
            "--gpu", str(self.cuda),
            "--pause-at", "93",
            "--resume-at", "85",
            "--poll-seconds", "2",
            "--safe-readings", "3",
            "--query-failures", "2",
            "--status-file", str(task_dir / "thermal_status.json"),
            "--",
            sys.executable,
            str(REPO_ROOT / "train" / "eval.py"),
            "--config_file", str(config_path),
            "--models_file", str(manifest_path),
            "--output_root", str(task_dir / "results"),
            "--progress_file", str(progress_path),
        ]
        with log_path.open("w", encoding="utf-8") as log:
            process = subprocess.Popen(
                command,
                cwd=REPO_ROOT,
                stdin=subprocess.DEVNULL,
                stdout=log,
                stderr=subprocess.STDOUT,
                env={**os.environ, "PYTHONUNBUFFERED": "1"},
            )
            with self.lock:
                self.state["tasks"][task_id]["pid"] = process.pid
                self._write_state()
            while process.poll() is None:
                self._update_progress(task_id, log_path, progress_path)
                time.sleep(10)
            return_code = process.returncode

        summaries = sorted(
            (task_dir / "results").glob("*/summary.json"),
            key=lambda path: path.stat().st_mtime,
        )
        records = []
        summary_path = None
        if return_code == 0 and summaries:
            summary_path = summaries[-1]
            payload = _load_json(summary_path) or {}
            records = payload.get("models") or []
        ok = return_code == 0 and len(records) == int(task["available_checkpoints"])
        metrics = {}
        if records:
            for key in (
                "success", "task_completion_ratio", "sent_percentage",
                "episode_length", "episode_reward",
            ):
                values = [float(record[key]) for record in records if key in record]
                if values:
                    metrics[f"{key}_mean"] = sum(values) / len(values)
                    metrics[f"{key}_min"] = min(values)
                    metrics[f"{key}_max"] = max(values)
        with self.lock:
            target = self.state["tasks"][task_id]
            target.update(
                state="completed" if ok else "failed",
                return_code=return_code,
                finished_at=_now(),
                completed_checkpoints=len(records),
                completed_units=len(records) * self.eval_episodes,
                summary_path=str(summary_path) if summary_path else None,
                metrics=metrics,
            )
            self._write_state()
        return ok

    def run(self):
        task_ids = self._discover()
        pending = [
            task_id for task_id in task_ids
            if self.state["tasks"][task_id].get("state") != "completed"
        ]
        results = []
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = [executor.submit(self._run_task, task_id) for task_id in pending]
            for future in as_completed(futures):
                results.append(future.result())
        with self.lock:
            failed = any(
                task.get("state") != "completed"
                for task in self.state["tasks"].values()
            )
            self.state.update(
                state="partial_failure" if failed else "completed",
                finished_at=_now(),
            )
            self._write_state()
        return 1 if failed else 0


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--stage", action="append", required=True, metavar="NAME=PATH")
    parser.add_argument("--recovery-stage", action="append", default=[], metavar="NAME=PATH")
    parser.add_argument("--cuda", type=int, default=1)
    parser.add_argument("--max-workers", type=int, default=12)
    parser.add_argument("--fraction", type=float, default=0.1)
    parser.add_argument("--eval-episodes", type=int, default=8)
    parser.add_argument("--eval-seed", type=int, default=30000)
    parser.add_argument("--job-pattern")
    args = parser.parse_args(argv)
    stages = _parse_named_paths(args.stage)
    recovery = _parse_named_paths(args.recovery_stage, multiple=True)
    return FinalWindowQueue(
        args.output_dir,
        stages,
        recovery,
        args.cuda,
        args.max_workers,
        args.fraction,
        args.eval_episodes,
        args.eval_seed,
        args.job_pattern,
    ).run()


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