"""Run full epochs across successive four-hour salloc allocations.

This process runs OUTSIDE the GPU allocation. It never submits sbatch jobs,
renews only after a TIMEOUT, and stops on cancellation or training failures.
"""

import argparse
import fcntl
import json
import os
import re
import shlex
import signal
import subprocess
import sys
import time
import uuid
from pathlib import Path


def write_json(path, value):
    temporary = path.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(value, indent=2) + "\n")
    os.replace(temporary, path)


class Supervisor:
    def __init__(self, repo, root, initial_job, epochs=1,
                 wandb_entity="Nemo-automodel",
                 wandb_project="speculators-ngram-candidates-validation"):
        if epochs < 1:
            raise ValueError("epochs must be positive")
        # Preserve the shared /lustre/fs1 spelling: this CPU node's canonical
        # mount alias need not exist on GPU nodes or on the login host.
        self.repo, self.root, self.job = repo.absolute(), root.absolute(), initial_job
        self.root.mkdir(parents=True, exist_ok=True)
        self.lock = (self.root / "supervisor.lock").open("a")
        fcntl.flock(self.lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        self.ssh = [
            "ssh",
            "-o",
            "BatchMode=yes",
            "-o",
            "ConnectTimeout=10",
            "-o",
            "ServerAliveInterval=30",
            "-o",
            "ServerAliveCountMax=3",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/tmp/codex_known_hosts",
        ]
        self.login = "huiyingl@oci-hsg-cs-001-login-02"
        self.child = self.holder = self.maintenance = None
        self.source = self.root / "source"
        manifest_path = self.root / "manifest.json"
        if manifest_path.exists():
            self.manifest = json.loads(manifest_path.read_text())
            if (self.manifest["epochs"], self.manifest["wandb_entity"],
                    self.manifest["wandb_project"]) != (epochs, wandb_entity, wandb_project):
                raise ValueError("Existing run configuration differs; use a new run root")
        else:
            if self.source.exists():
                raise RuntimeError("Unowned source snapshot already exists")
            revision = subprocess.check_output(
                ["git", "rev-parse", "HEAD"], cwd=self.repo, text=True
            ).strip()
            run_name = self.root.name
            checkpoint = (
                self.repo.parent
                / "model_weights"
                / "ngram_candidates_qwen3_6-35b-a3b-perfectblend"
                / run_name
                / "checkpoints"
            )
            if checkpoint.exists():
                raise RuntimeError("Use a new run name; checkpoint directory exists")
            subprocess.run(
                [
                    "git",
                    "clone",
                    "--no-hardlinks",
                    "--no-checkout",
                    "--quiet",
                    str(self.repo),
                    str(self.source),
                ],
                check=True,
            )
            subprocess.run(
                ["git", "checkout", "--detach", "--quiet", revision],
                cwd=self.source,
                check=True,
            )
            self.manifest = {
                "run_name": run_name,
                "source_revision": revision,
                "source_snapshot": str(self.source),
                "initial_job": initial_job,
                "checkpoint_root": str(checkpoint),
                "epochs": epochs,
                "data_path": str(
                    self.repo.parent
                    / "datasets/qwen3.6-35b-a3b/qwen3.6-35b-a3b_train_spec_len3072_fullvocab"
                ),
                "wandb_entity": wandb_entity,
                "wandb_project": wandb_project,
                "wandb_run_id": uuid.uuid4().hex[:8],
                "image": str(
                    self.repo / "images/speculators-ngram-py312-vllm0.27.1-arm64.sqsh"
                ),
                "image_sha256": "9c0b5afa769aa3637f1470341d7725f3b752f72246d978e97a17bf793b74f368",
                "topology": "single node; GPU0 vLLM, GPU1-3 DDP",
                "checkpoint_fraction": 0.01,
                "retained_checkpoint_generations_per_epoch": 2,
            }
            write_json(manifest_path, self.manifest)
        (self.root / "supervisor.pid").write_text(str(os.getpid()) + "\n")
        for name in ("segments", "allocations", "maintenance"):
            (self.root / name).mkdir(exist_ok=True)

    def status(self, phase, **extra):
        value = {
            "time": time.time(),
            "phase": phase,
            "job_id": self.job,
            "supervisor_pid": os.getpid(),
            **extra,
        }
        write_json(self.root / "status.json", value)
        print(json.dumps(value), flush=True)

    def remote(self, argv, timeout=30):
        return subprocess.run(
            self.ssh + [self.login, shlex.join(argv)],
            text=True,
            capture_output=True,
            timeout=timeout,
            check=True,
        ).stdout.strip()

    def job_state(self):
        try:
            result = self.remote(["squeue", "-h", "-j", self.job, "-o", "%T"])
        except subprocess.CalledProcessError:
            # squeue exits nonzero once an old job has left the controller;
            # sacct remains authoritative for completed allocations.
            result = ""
        queued = result.splitlines()[0] if result else "UNKNOWN"
        if queued not in {"COMPLETING", "UNKNOWN"}:
            return queued
        # Node cleanup can leave squeue in COMPLETING for several minutes,
        # while accounting already records the actual reason for termination.
        result = self.remote(
            ["sacct", "-n", "-P", "-X", "-j", self.job, "--format=State%40"]
        )
        accounted = (
            result.splitlines()[0].split("|")[0].split()[0]
            if result else "UNKNOWN"
        )
        if accounted in {"RUNNING", "COMPLETING", "UNKNOWN"}:
            return queued
        return accounted

    def wait_for_allocation_exit(self):
        running_deadline = time.monotonic() + 60
        while True:
            self.check_stop()
            try:
                state = self.job_state()
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error:
                # A transient SSH/accounting outage is not evidence of TIMEOUT.
                self.status("waiting_for_allocation_state", error=str(error))
                time.sleep(15)
                continue
            if state not in {"RUNNING", "COMPLETING", "UNKNOWN"}:
                return state
            if state == "RUNNING" and time.monotonic() >= running_deadline:
                # A dead trainer in a healthy allocation is a real failure.
                return state
            self.status("waiting_for_allocation_state", allocation_state=state)
            time.sleep(15)

    def recover_existing_segment(self):
        if not (self.root / "segments" / self.job).exists():
            return
        state = self.wait_for_allocation_exit()
        if not state.startswith("TIMEOUT"):
            raise RuntimeError(
                f"Existing segment allocation state {state}; not renewing or relaunching"
            )
        self.status("expired", recovery=True)
        self.allocate()

    def environment(self, segment):
        m = self.manifest
        return {
            "REPO": str(self.source),
            "ROOT": str(self.repo.parent),
            "NGRAM_RUN_ROOT": str(self.root),
            "RUN_NAME": m["run_name"],
            "LOG_DIR": str(segment),
            "CHECKPOINT_DIR": m["checkpoint_root"],
            "DATA_DIR": m["data_path"],
            "WANDB_ENTITY": m["wandb_entity"],
            "WANDB_PROJECT": m["wandb_project"],
            "WANDB_RUN_ID": m["wandb_run_id"],
            "WANDB_RESUME": "allow",
            "WANDB_MODE": "online",
            "EPOCHS": str(m["epochs"]),
        }

    def container_command(self, segment, argv, cpus):
        environment = self.environment(segment)
        outer = {
            **os.environ,
            "REPO": str(self.source),
            "NGRAM_STEP_CPUS": str(cpus),
            "NGRAM_CONTAINER_IMAGE": self.manifest["image"],
            "NGRAM_CONTAINER_NAME": "ngram-candidates-vllm0271",
        }
        command = [
            "bash",
            str(self.source / "scripts/cluster/ngram_container.sh"),
            self.job,
            "/usr/bin/env",
            *[f"{key}={value}" for key, value in environment.items()],
            *argv,
        ]
        return command, outer

    def start_segment(self):
        segment = self.root / "segments" / self.job
        segment.mkdir(exist_ok=False)
        command, environment = self.container_command(
            segment,
            ["bash", str(self.source / "scripts/cluster/ngram_full_segment.sh")],
            72,
        )
        with (segment / "controller.log").open("w") as output:
            self.child = subprocess.Popen(
                command,
                env=environment,
                stdout=output,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        self.status("training", segment=str(segment), launcher_pid=self.child.pid)
        return segment

    def start_maintenance(self, segment, action):
        if self.maintenance is not None and self.maintenance.poll() is None:
            return
        script = self.source / "scripts/cluster/ngram_wandb_maintenance.py"
        runtime = self.source / "scripts/cluster/ngram_runtime_env.sh"
        command_text = (
            f"source {shlex.quote(str(runtime))}; status=0; "
            f"/opt/speculators_venv/bin/python {shlex.quote(str(script))} {action} || status=$?; "
            'task_cache="/tmp/${USER:-speculators}/speculators-cache"; '
            'if [[ -d "$task_cache" ]]; then mkdir -p "$ROOT/.cache/compiled/ngram-candidates"; '
            'rsync -a "$task_cache/" "$ROOT/.cache/compiled/ngram-candidates/"; fi; exit "$status"'
        )
        command, environment = self.container_command(
            segment, ["bash", "-lc", command_text], 4
        )
        log = self.root / "maintenance" / f"{int(time.time())}-{self.job}-{action}.log"
        with log.open("w") as output:
            self.maintenance = subprocess.Popen(
                command,
                env=environment,
                stdout=output,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        self.status(
            "training" if action == "sync" else "uploading", maintenance=str(log)
        )

    def allocate(self):
        self.status("requesting_interactive")
        shares = self.remote(["sshare", "-n", "-p", "-u", "huiyingl"])
        accounts = []
        for line in shares.splitlines():
            fields = line.split("|")
            if len(fields) > 6 and fields[1].strip() and fields[6].strip():
                accounts.append((float(fields[6]), fields[0].strip()))
        if not accounts:
            raise RuntimeError("Could not resolve an eligible fairshare account")
        account = max(accounts)[1]
        request = self.root / "allocations" / str(int(time.time()))
        request.mkdir()
        allocation = [
            "salloc",
            f"--account={account}",
            "--partition=batch",
            "--qos=interactive",
            "--nodes=1",
            "--ntasks=1",
            "--gpus-per-node=4",
            "--cpus-per-task=72",
            "--mem=0",
            "--time=04:00:00",
            "--job-name=ngram-candidates-full",
            "bash",
            str(self.source / "scripts/cluster/ngram_allocation_holder.sh"),
            str(request),
        ]
        write_json(
            request / "request.json", {"command": allocation, "account": account}
        )
        with (request / "salloc.log").open("w") as output:
            self.holder = subprocess.Popen(
                self.ssh + ["-tt", self.login, shlex.join(allocation)],
                stdin=subprocess.DEVNULL,
                stdout=output,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        while not (request / "job_id").exists():
            self.check_stop()
            if self.holder.poll() is not None:
                raise RuntimeError(
                    f"Interactive request exited {self.holder.returncode}; see {request}"
                )
            time.sleep(15)
        self.job = (request / "job_id").read_text().strip()
        if not re.fullmatch(r"\d+", self.job):
            raise RuntimeError("Invalid allocation ID")
        self.status("interactive_allocated", allocation_request=str(request))

    def check_stop(self):
        if (self.root / "STOP").exists():
            raise KeyboardInterrupt("STOP file requested; do not renew")

    def run(self):
        self.recover_existing_segment()
        while True:
            self.check_stop()
            segment = self.start_segment()
            next_sync = time.monotonic() + 3600
            while self.child.poll() is None:
                self.check_stop()
                if time.monotonic() >= next_sync:
                    self.start_maintenance(segment, "sync")
                    next_sync = time.monotonic() + 3600
                time.sleep(15)
            result = self.child.returncode
            if (self.root / "training_complete.json").exists() and result == 0:
                if self.maintenance is not None:
                    self.maintenance.wait(timeout=1200)
                self.start_maintenance(segment, "sync")
                sync_result = self.maintenance.wait(timeout=1200)
                self.start_maintenance(segment, "upload")
                upload_result = self.maintenance.wait(timeout=4200)
                if sync_result or upload_result:
                    raise RuntimeError(
                        "Training complete, but W&B final sync/upload failed"
                    )
                self.status("complete", checkpoint=self.manifest["checkpoint_root"])
                return
            state = self.wait_for_allocation_exit()
            if not state.startswith("TIMEOUT"):
                raise RuntimeError(
                    f"Training exited {result}, allocation state {state}; not renewing a failed/cancelled job"
                )
            if self.maintenance is not None and self.maintenance.poll() is None:
                self.maintenance.wait(timeout=180)
            self.status("expired", launcher_exit_code=result)
            self.allocate()

    def stop_children(self):
        for child in (self.child, self.maintenance, self.holder):
            if child is not None and child.poll() is None:
                os.killpg(child.pid, signal.SIGTERM)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", type=Path, required=True)
    parser.add_argument("--run-root", type=Path, required=True)
    parser.add_argument("--initial-job", required=True)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--wandb-entity", default="Nemo-automodel")
    parser.add_argument("--wandb-project", default="speculators-ngram-candidates-validation")
    args = parser.parse_args()
    controller = Supervisor(args.repo, args.run_root, args.initial_job, args.epochs,
                            args.wandb_entity, args.wandb_project)
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(
            sig, lambda signum, frame: (_ for _ in ()).throw(KeyboardInterrupt())
        )
    try:
        controller.run()
    except BaseException as error:
        controller.status(
            "stopped" if isinstance(error, KeyboardInterrupt) else "failed",
            error=str(error),
        )
        controller.stop_children()
        raise


if __name__ == "__main__":
    main()