#!/usr/bin/env python3
"""Opt-in legacy OpenEnv controller after an explicitly selected JAIF predecessor.

The check mode is offline. Launch waits for an acknowledged NeMo release and a
terminal/missing scheduler row before submitting one bounded command job. Native
command exit releases its GPU; no SSH keys, model weights or datasets are uploaded.
"""
from __future__ import annotations

import argparse
from datetime import datetime, timezone
import fcntl
import hashlib
import json
import os
from pathlib import Path, PurePosixPath
import re
import shutil
import signal
import subprocess
import sys
import tarfile
import tempfile
import time
import uuid

ROOT = Path(__file__).resolve().parents[1]
NEMO_ID = os.environ.get("HESSIAN_NEMO_SHELL_ID")
WORKSPACE_ID = 123
ARCHIVE_SHA256 = "289ecc42e53d3fa6ac5b8532b447977f53e571837fbc4fae133ba42cf5437cfe"
CONFIG = ROOT / "cluster/determined/german-qwen-openenv-command.yaml"
TERMINAL = {"TERMINATED", "CANCELED", "CANCELLED", "ERROR", "COMPLETED", "STOPPED"}
ACTIVE = {"RUNNING", "READY", "STARTING", "PULLING", "TERMINATING"}
sys.path.insert(0, str(ROOT / "src"))


def digest(path):
    value = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for block in iter(lambda: stream.read(1048576), b""):
            value.update(block)
    return value.hexdigest()


def read(path):
    from hessian_recipes.training import _unique_pairs

    return json.loads(Path(path).read_text(), object_pairs_hook=_unique_pairs,
                      parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value)))


def write(path, value):
    from hessian_recipes.gym_experiment import _write

    _write(path, value)


def verify(plan_path, expected):
    from hessian_recipes.openenv_experiment import verify_plan

    path = Path(plan_path).resolve()
    if (not re.fullmatch(r"[0-9a-f]{64}", expected)
            or not path.is_relative_to(ROOT / "runs") or digest(path) != expected):
        raise ValueError("OpenEnv plan path/digest changed")
    plan = verify_plan(path)
    if (plan["max_gpu_wall_seconds"] != 28800 or plan.get("backend") != "openenv"
            or Path(plan["sdk"]["python"]) != ROOT / ".venv-openenv/bin/python"):
        raise ValueError("Unexpected OpenEnv budget/backend/SDK")
    if digest(ROOT / ".cache/agent-rl-site.tar") != ARCHIVE_SHA256:
        raise ValueError("GPU overlay archive changed")
    return plan


def scheduler(kind):
    result = subprocess.run(["det", kind, "list", "--json"], capture_output=True,
                            text=True, timeout=30, check=True)
    rows = json.loads(result.stdout)
    if not isinstance(rows, list):
        raise ValueError("Unexpected scheduler list schema")
    # Do not retain or print arbitrary scheduler fields: shells contain keys.
    return [{key: row.get(key) for key in
             ("id", "state", "description", "resourcePool", "username", "workspaceId")}
            for row in rows if isinstance(row, dict)]


def state_name(row):
    return str(row.get("state", "")).removeprefix("STATE_")


def require_predecessor():
    if not isinstance(NEMO_ID, str) or not re.fullmatch(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}", NEMO_ID):
        raise ValueError("Explicit JAIF predecessor shell ID is required")
    return NEMO_ID


def predecessor_released(plan, rows):
    require_predecessor()
    status_path = Path(plan["nemo_origin"]["path"]).parent / "queue-launcher-status.json"
    if not status_path.exists():
        return False
    status = read(status_path)
    if (status.get("shell_id") != NEMO_ID or status.get("workspace_id") != WORKSPACE_ID
            or status.get("plan_sha256") != plan["nemo_origin"]["sha256"]
            or status.get("event") != "allocation_release"
            or status.get("release_acknowledged") is not True):
        return False
    matching = [row for row in rows if row.get("id") == NEMO_ID]
    if len(matching) > 1:
        raise ValueError("Ambiguous predecessor allocation")
    if matching and matching[0].get("workspaceId") != WORKSPACE_ID:
        raise ValueError("Predecessor shell must belong to JAIF")
    return not matching or state_name(matching[0]) in TERMINAL


def event(path, record, name, **fields):
    record.update(event=name, updated_at=datetime.now(timezone.utc).isoformat(), **fields)
    write(path, record)
    print(json.dumps({"event": name, **fields}), flush=True)


def create_command(plan_path, expected, description):
    command = ["det", "command", "run", "--config-file", str(CONFIG),
               "--workspace-name", "JAIF", "--detach", "--config", "description=" + description,
               "/usr/bin/timeout", "--signal=KILL", "28800s", str(ROOT / ".venv-gpu/bin/python"),
               str(Path(__file__).resolve()), "--worker", "--plan", str(plan_path),
               "--expected-plan-sha256", expected]
    result = subprocess.run(command, capture_output=True, text=True, timeout=90, check=True)
    matches = set(re.findall(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}", result.stdout))
    if len(matches) != 1:
        raise ValueError("Creation response did not identify exactly one command")
    return str(uuid.UUID(matches.pop()))


def find_created(rows, description):
    matches = [row for row in rows if row.get("description") == description]
    if len(matches) > 1:
        raise ValueError("Ambiguous earlier command creation; no duplicate will be submitted")
    if not matches:
        return None
    row = matches[0]
    if row.get("resourcePool") != "42_Compute":
        raise ValueError("Created command is outside the frozen pool")
    return str(uuid.UUID(row["id"]))


def kill_owned(identifier):
    str(uuid.UUID(identifier))
    for attempt in range(3):
        try:
            result = subprocess.run(["det", "command", "kill", identifier], capture_output=True,
                                    timeout=30, check=False)
            if result.returncode == 0:
                return True
        except (OSError, subprocess.TimeoutExpired):
            pass
        if attempt < 2:
            time.sleep(2)
    return False


def experiment_status(root, expected):
    path = root / "experiment-status.json"
    if not path.exists():
        return "no_worker_evidence"
    report = read(path)
    if report.get("plan_sha256") != expected:
        raise ValueError("Worker evidence belongs to another frozen plan")
    return report.get("status")


def check_deadline(current):
    since = current["active_since_epoch"]
    if since is not None:
        elapsed = time.time() - since
        if elapsed < 0 or elapsed >= 28800:
            raise TimeoutError("Eight-hour active GPU budget exhausted or clock moved back")


def launch(plan_path, expected, poll=15):
    require_predecessor()
    if not 1 <= poll <= 30:
        raise ValueError("Poll interval must be in [1,30] seconds")
    plan = verify(plan_path, expected)
    root = Path(plan["root"])
    status_path = root / "command-launcher-status.json"
    description = "hessian-openenv-" + expected[:16]
    with (root / ".command-launcher.lock").open("a") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        current = read(status_path) if status_path.exists() else {
            "plan_sha256": expected, "launcher_sha256": digest(__file__),
            "description": description, "predecessor": NEMO_ID, "command_id": None,
            "creation_started": False, "active_since_epoch": None,
            "command_seen": False, "created_since_epoch": None}
        if (current["plan_sha256"] != expected or current["launcher_sha256"] != digest(__file__)
                or current["description"] != description):
            raise ValueError("Existing controller state belongs to another frozen launch")
        owned = current.get("command_id")
        try:
            if owned is None:
                while not predecessor_released(plan, scheduler("shell")):
                    event(status_path, current, "waiting_for_nemo_release")
                    time.sleep(poll)
                verify(plan_path, expected)
                if current["creation_started"]:
                    owned = find_created(scheduler("command"), description)
                    if owned is None:
                        raise RuntimeError("Earlier creation is unresolved; refusing a duplicate")
                else:
                    event(status_path, current, "creating_one_command", creation_started=True,
                          created_since_epoch=time.time())
                    owned = create_command(plan_path, expected, description)
                event(status_path, current, "command_created", command_id=owned)
            missing_since = None
            while True:
                check_deadline(current)
                rows = [row for row in scheduler("command") if row.get("id") == owned]
                if len(rows) > 1:
                    raise ValueError("Ambiguous owned command state")
                if not rows:
                    # A newly created command may not appear in list immediately.
                    # Do not abandon a possibly live allocation on a transient gap.
                    if missing_since is None:
                        missing_since = time.time()
                    if experiment_status(root, expected) == "completed":
                        # A successful worker may disappear before our first poll;
                        # still obtain an explicit stop acknowledgement on this path.
                        released = kill_owned(owned)
                        event(status_path, current, "worker_completed_command_absent",
                              owned_command_stop_acknowledged=released)
                        return 0
                    if time.time() - missing_since >= 120:
                        raise RuntimeError("Owned command absent from scheduler for two minutes")
                    event(status_path, current, "awaiting_scheduler_visibility")
                    time.sleep(poll)
                    continue
                missing_since = None
                row = rows[0]
                if (row.get("description") != description
                        or row.get("resourcePool") != "42_Compute"):
                    raise ValueError("Owned command identity or pool changed")
                current["command_seen"] = True
                if state_name(row) in TERMINAL:
                    result = experiment_status(root, expected)
                    event(status_path, current, "scheduler_finished", experiment_status=result)
                    return 0 if result == "completed" else 1
                state = state_name(row)
                if state in ACTIVE and current["active_since_epoch"] is None:
                    current["active_since_epoch"] = time.time()
                check_deadline(current)
                event(status_path, current, "scheduler", state=state)
                time.sleep(poll)
        except BaseException as error:
            released = kill_owned(owned) if owned else False
            event(status_path, current, "failed", error_type=type(error).__name__,
                  owned_command_stop_acknowledged=released)
            raise


def validate_member(member):
    path = PurePosixPath(member.name)
    if path.is_absolute() or ".." in path.parts or not (member.isfile() or member.isdir()):
        raise ValueError("Unsafe GPU overlay member")


def worker(plan_path, expected):
    plan = verify(plan_path, expected)
    root = Path(plan["root"])
    stage = Path(tempfile.mkdtemp(prefix="hessian-openenv-", dir="/tmp"))
    archive = stage / "site.tar"
    shutil.copyfile(ROOT / ".cache/agent-rl-site.tar", archive)
    if digest(archive) != ARCHIVE_SHA256:
        raise ValueError("Worker overlay copy changed")
    site = stage / "site-packages"
    site.mkdir()
    with tarfile.open(archive, "r:") as bundle:
        members = bundle.getmembers()
        for member in members:
            validate_member(member)
        bundle.extractall(site, members=members, filter="data")
    archive.unlink()
    shared = ROOT / ".venv-gpu/lib/python3.11/site-packages"
    for name in ("torch/lib", "nvidia", "triton"):
        source, target = shared / name, site / name
        if not source.is_dir() or target.exists() or target.is_symlink():
            raise ValueError("Unexpected GPU overlay library location")
        target.parent.mkdir(parents=True, exist_ok=True)
        target.symlink_to(source, target_is_directory=True)
    environment = dict(os.environ, PYTHONPATH=str(site) + os.pathsep + str(ROOT / "src"),
                       HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", HF_DATASETS_OFFLINE="1",
                       TOKENIZERS_PARALLELISM="false", OMP_NUM_THREADS="4", MKL_NUM_THREADS="4")
    os.chdir(ROOT)
    python = str(ROOT / ".venv-gpu/bin/python")
    subprocess.run([python, str(ROOT / "scripts/gpu-preflight.py"), "--output",
                    str(root / f"command-gpu-preflight-{time.time_ns()}.json")],
                   env=environment, timeout=180, check=True)
    write(root / "command-worker-provenance.json", {
        "plan_sha256": expected, "launcher_sha256": digest(__file__),
        "overlay_archive_sha256": ARCHIVE_SHA256,
        "external_active_deadline_seconds": 28800, "model_data_uploads": False})
    args = [python, "-m", "hessian_recipes.openenv_experiment", "run", "--plan", str(plan_path)]
    if (root / "experiment-status.json").exists():
        args.append("--resume")
    os.execvpe(python, args, environment)


def main():
    global NEMO_ID
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--predecessor-shell-id", default=NEMO_ID, help="explicit existing JAIF predecessor; no historical allocation default")
    parser.add_argument("--plan", type=Path, required=True)
    parser.add_argument("--expected-plan-sha256", required=True)
    modes = parser.add_mutually_exclusive_group(required=True)
    modes.add_argument("--check", action="store_true")
    modes.add_argument("--launch", action="store_true")
    modes.add_argument("--worker", action="store_true")
    parser.add_argument("--poll-seconds", type=int, default=15)
    args = parser.parse_args()
    NEMO_ID = args.predecessor_shell_id
    if args.launch:
        require_predecessor()
    plan = verify(args.plan, args.expected_plan_sha256)
    if args.check:
        print(json.dumps({"status": "checked", "backend": plan["backend"],
                          "cluster_contacted": False, "model_inference": False}))
        return 0
    def interrupted(signum, frame):
        raise KeyboardInterrupt
    for signum in (signal.SIGINT, signal.SIGTERM):
        signal.signal(signum, interrupted)
    if args.worker:
        worker(args.plan, args.expected_plan_sha256)
        return 0
    return launch(args.plan, args.expected_plan_sha256, args.poll_seconds)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except BaseException as error:
        if isinstance(error, SystemExit):
            raise
        print(json.dumps({"status": "failed", "error_type": type(error).__name__}), file=sys.stderr)
        raise SystemExit(1) from None