#!/usr/bin/env python3
"""Opt-in legacy controller for one explicitly selected existing JAIF allocation.

No allocation is requested. Scheduler JSON and generated SSH commands are
captured internally; only allowlisted status fields are printed or persisted.
Interrupted queue waits preserve the request. An observed active allocation is
released on success, failure, timeout or handled interruption.
"""
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 shlex
import signal
import subprocess
import sys
import time
from urllib.parse import urlsplit

ROOT = Path(__file__).resolve().parents[1]
OWNED_SHELL = os.environ.get("HESSIAN_NEMO_SHELL_ID")
WORKSPACE_ID = 123
ARCHIVE_SHA256 = "289ecc42e53d3fa6ac5b8532b447977f53e571837fbc4fae133ba42cf5437cfe"
ACTIVE_STATES = {"RUNNING", "READY", "STARTING", "PULLING", "TERMINATING"}
READY_STATES = {"RUNNING", "READY"}
TERMINAL_STATES = {"TERMINATED", "TERMINATING", "CANCELED", "CANCELLED", "ERROR", "COMPLETED", "STOPPED"}


class LaunchError(RuntimeError):
    def __init__(self, code):
        super().__init__(code)
        self.code = code


class Interrupted(BaseException):
    def __init__(self, signum):
        self.signum = signum


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


def require_owned_shell():
    if not isinstance(OWNED_SHELL, str) or not re.fullmatch(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}", OWNED_SHELL):
        raise LaunchError("explicit_owned_shell_id_required")
    return OWNED_SHELL


def owned_snapshot(raw):
    """Never return arbitrary scheduler keys (shells can contain privateKey)."""
    require_owned_shell()
    if not isinstance(raw, list):
        raise LaunchError("unexpected_scheduler_schema")
    matches = [row for row in raw if isinstance(row, dict) and row.get("id") == OWNED_SHELL]
    if len(matches) != 1:
        raise LaunchError("owned_shell_missing_or_ambiguous")
    row = matches[0]
    state = str(row.get("state", "")).removeprefix("STATE_")
    username = row.get("username")
    if not re.fullmatch(r"[A-Z_]+", state) or not isinstance(username, str) or not re.fullmatch(r"[A-Za-z0-9_.@-]+", username):
        raise LaunchError("invalid_scheduler_identity")
    if row.get("workspaceId") != WORKSPACE_ID:
        raise LaunchError("owned_shell_must_belong_to_jaif")
    if row.get("resourcePool") != "42_Compute":
        raise LaunchError("unexpected_owned_resource_pool")
    return {"shell_id": OWNED_SHELL, "state": state, "username": username,
            "resource_pool": "42_Compute", "workspace_id": WORKSPACE_ID}


def scheduler_snapshot(environment):
    require_owned_shell()
    # Without --all the installed Determined CLI filters to the authenticated
    # user's own shells. Do not broaden this query or print the returned object.
    try:
        result = subprocess.run(["det", "shell", "list", "--json"], env=environment,
                                capture_output=True, text=True, timeout=30, check=False)
        if result.returncode or len(result.stdout) > 8_000_000:
            raise LaunchError("scheduler_query_failed")
        raw = json.loads(result.stdout)
    except (OSError, subprocess.TimeoutExpired, ValueError):
        raise LaunchError("scheduler_query_failed") from None
    return owned_snapshot(raw)


def connection_config(command_text, cache):
    require_owned_shell()
    """Validate the native CLI's generated command, then emit a private config."""
    plain = re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", command_text)
    lines = [line.strip() for line in plain.splitlines() if line.strip().startswith("ssh ")]
    if len(lines) != 1:
        raise LaunchError("ssh_command_missing_or_ambiguous")
    try:
        tokens = shlex.split(lines[0])
    except ValueError:
        raise LaunchError("invalid_ssh_command") from None
    options, key, target, index = {}, None, None, 1
    while index < len(tokens):
        token = tokens[index]
        if token == "-tt":
            index += 1
        elif token in {"-o", "-i"} and index + 1 < len(tokens):
            value = tokens[index + 1]
            if token == "-i":
                if key is not None:
                    raise LaunchError("duplicate_ssh_key")
                key = Path(value)
            else:
                name, separator, argument = value.partition("=")
                if not separator or name not in {"ProxyCommand", "StrictHostKeyChecking", "IdentitiesOnly"} or name in options:
                    raise LaunchError("unexpected_ssh_option")
                options[name] = argument
            index += 2
        elif not token.startswith("-") and index == len(tokens) - 1:
            target, index = token, index + 1
        else:
            raise LaunchError("unexpected_ssh_argument")
    match = re.fullmatch(r"([A-Za-z0-9_.-]+)@" + re.escape(OWNED_SHELL), target or "")
    if (not match or key is None or not key.is_absolute() or key.is_symlink() or not key.is_file()
            or key.stat().st_mode & 0o077 or not key.resolve().is_relative_to(cache.resolve())):
        raise LaunchError("ssh_target_or_key_outside_owned_cache")
    if options.get("StrictHostKeyChecking") != "no" or options.get("IdentitiesOnly") != "yes":
        raise LaunchError("unexpected_native_ssh_policy")
    proxy = shlex.split(options.get("ProxyCommand", ""))
    if (len(proxy) < 5 or not Path(proxy[0]).is_absolute() or not Path(proxy[0]).name.startswith("python")
            or proxy[1:3] != ["-m", "determined.cli.tunnel"] or proxy[4] != "%h"):
        raise LaunchError("unexpected_ssh_proxy")
    master = urlsplit(proxy[3])
    if master.scheme not in {"http", "https"} or not master.hostname or master.username or master.password:
        raise LaunchError("invalid_proxy_master")
    extras = proxy[5:]
    if len(extras) % 2 or any(extras[index] not in {"--cert-file", "--cert-name"} for index in range(0, len(extras), 2)):
        raise LaunchError("unexpected_proxy_arguments")
    if any("\n" in value or "\r" in value or "\0" in value for value in proxy):
        raise LaunchError("invalid_proxy_control_character")
    return "\n".join([
        "Host hessian-nemo-long-owned", f"  HostName {OWNED_SHELL}", f"  User {match.group(1)}",
        f"  IdentityFile {json.dumps(str(key))}", f"  ProxyCommand {shlex.join(proxy)}",
        "  IdentitiesOnly yes", "  BatchMode yes", "  RequestTTY no", "  StrictHostKeyChecking no",
        f"  UserKnownHostsFile {json.dumps(str(cache / 'known_hosts'))}", "  ConnectTimeout 20",
        "  ServerAliveInterval 15", "  ServerAliveCountMax 3", "",
    ])


def prepare_connection(environment, cache):
    try:
        result = subprocess.run(["det", "shell", "show-ssh-command", OWNED_SHELL], env=environment,
                                capture_output=True, text=True, timeout=30, check=False)
        if result.returncode:
            raise LaunchError("ssh_configuration_failed")
    except (OSError, subprocess.TimeoutExpired):
        raise LaunchError("ssh_configuration_failed") from None
    configuration = connection_config(result.stdout, cache)
    path = cache / "owned-shell.ssh.conf"
    descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(descriptor, "w") as stream:
        stream.write(configuration)
    path.chmod(0o600)
    return path


def validate_tar_member(member):
    """The pinned archive contains only ordinary files/directories."""
    path = PurePosixPath(member.name)
    if path.is_absolute() or ".." in path.parts or not (member.isfile() or member.isdir()):
        raise ValueError("unsafe staging archive member")


NODE_PROGRAM = r'''
import hashlib, json, os, shutil, socket, subprocess, sys, tarfile, tempfile, time
from pathlib import Path, PurePosixPath

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 validate_tar_member(member):
    path = PurePosixPath(member.name)
    if path.is_absolute() or ".." in path.parts or not (member.isfile() or member.isdir()):
        raise ValueError("unsafe staging archive member")

root, plan = Path(configuration["root"]), Path(configuration["plan"])
if digest(plan) != configuration["plan_sha256"]:
    raise RuntimeError("frozen plan digest changed before worker staging")
stage = Path(tempfile.mkdtemp(prefix="hessian-nemo-long-", dir="/tmp"))
archive = stage / "site.tar"
shutil.copyfile(root / ".cache/agent-rl-site.tar", archive)
if digest(archive) != configuration["archive_sha256"]:
    raise RuntimeError("worker staging archive hash mismatch")
site = stage / "site-packages"
site.mkdir()
with tarfile.open(archive, "r:") as bundle:
    members = bundle.getmembers()
    for member in members:
        validate_tar_member(member)
    bundle.extractall(site, members=members, filter="data")
archive.unlink()
shared_site = root / ".venv-gpu/lib/python3.11/site-packages"
links = {}
for relative in ("torch/lib", "nvidia", "triton"):
    source, target = shared_site / relative, site / relative
    if not source.is_dir() or target.exists() or target.is_symlink():
        raise RuntimeError("invalid GPU library staging location")
    target.parent.mkdir(parents=True, exist_ok=True)
    target.symlink_to(source, target_is_directory=True)
    links[relative] = str(source)
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)
sys.path.insert(0, str(root / "src"))
from hessian_recipes.gym_experiment import verify_plan
frozen = verify_plan(plan)
if digest(plan) != configuration["plan_sha256"]:
    raise RuntimeError("frozen plan digest changed on worker")
run = Path(frozen["root"])
metadata = {"shell_id": configuration["shell_id"], "hostname": socket.gethostname(),
            "plan_sha256": configuration["plan_sha256"], "archive_sha256": configuration["archive_sha256"],
            "overlay": str(site), "shared_library_links": links, "launcher_sha256": configuration["launcher_sha256"]}
(run / ("launcher-node-" + str(time.time_ns()) + ".json")).write_text(json.dumps(metadata, indent=2) + "\n")
python = str(root / ".venv-gpu/bin/python")
# This CUDA/BF16 check precedes every possible Qwen model-consuming stage.
subprocess.run([python, str(root / "scripts/gpu-preflight.py"), "--output",
                str(run / ("launcher-gpu-preflight-" + str(time.time_ns()) + ".json"))],
               env=environment, check=True, timeout=180)
command = [python, "-m", "hessian_recipes.gym_experiment", "run", "--plan", str(plan),
           "--gym-python", configuration["gym_python"]]
if (run / "experiment-status.json").exists():
    command.append("--resume")
os.execvpe(python, command, environment)
'''


def node_program(configuration):
    # JSON is inserted as a Python string literal, never into a shell command.
    return "import json\nconfiguration = json.loads(" + repr(json.dumps(configuration)) + ")\n" + NODE_PROGRAM


def verify_inputs(plan_path, expected, gym_python):
    if not re.fullmatch(r"[0-9a-f]{64}", expected):
        raise LaunchError("invalid_expected_plan_digest")
    plan_path = Path(plan_path).resolve()
    if not plan_path.is_relative_to(ROOT / "runs") or plan_path.name != "experiment-plan.json":
        raise LaunchError("plan_outside_repository_runs")
    if digest(plan_path) != expected:
        raise LaunchError("frozen_plan_digest_mismatch")
    sys.path.insert(0, str(ROOT / "src"))
    from hessian_recipes.gym_experiment import verify_plan
    plan = verify_plan(plan_path)
    if plan["max_gpu_wall_seconds"] > 8 * 3600:
        raise LaunchError("gpu_budget_exceeds_eight_hours")
    # Preserve venv symlink spelling: resolving bin/python would discard its venv.
    gym_python = Path(gym_python).absolute()
    if gym_python != ROOT / ".venv-nemo-gym/bin/python" or not gym_python.is_file():
        raise LaunchError("unexpected_gym_python")
    archive = ROOT / ".cache/agent-rl-site.tar"
    if digest(archive) != ARCHIVE_SHA256:
        raise LaunchError("local_staging_archive_hash_mismatch")
    site = ROOT / ".venv-gpu/lib/python3.11/site-packages"
    if any(not (site / name).is_dir() for name in ("torch/lib", "nvidia", "triton")):
        raise LaunchError("missing_shared_gpu_libraries")
    return {"root": str(ROOT), "plan": str(plan_path), "plan_sha256": expected,
            "gym_python": str(gym_python), "shell_id": OWNED_SHELL, "archive_sha256": ARCHIVE_SHA256,
            "launcher_sha256": digest(__file__), "budget_seconds": plan["max_gpu_wall_seconds"]}


def write_status(path, state, event, **safe_fields):
    state.update(event=event, updated_at=datetime.now(timezone.utc).isoformat(), **safe_fields)
    temporary = path.with_suffix(".writing")
    temporary.write_text(json.dumps(state, indent=2) + "\n")
    os.replace(temporary, path)
    print(json.dumps({"event": event, "shell_id": OWNED_SHELL, **safe_fields}), flush=True)


def release_owned(environment):
    require_owned_shell()
    for attempt in range(3):
        print(json.dumps({"event": "allocation_release_attempt", "shell_id": OWNED_SHELL,
                          "attempt": attempt + 1}), flush=True)
        try:
            result = subprocess.run(["det", "shell", "kill", OWNED_SHELL], env=environment,
                                    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 launch(configuration, poll_seconds=15):
    require_owned_shell()
    if not 1 <= poll_seconds <= 30:
        raise LaunchError("poll_interval_must_be_one_to_thirty_seconds")
    cache = ROOT / ".cache/determined-long"
    cache.mkdir(parents=True, exist_ok=True, mode=0o700)
    if cache.is_symlink() or cache.stat().st_uid != os.getuid():
        raise LaunchError("invalid_owned_ssh_cache")
    cache.chmod(0o700)
    environment = dict(os.environ, XDG_CACHE_HOME=str(cache))
    status_path = Path(configuration["plan"]).parent / "queue-launcher-status.json"
    state = {"shell_id": OWNED_SHELL, "workspace_id": WORKSPACE_ID, "plan_sha256": configuration["plan_sha256"],
             "launcher_sha256": configuration["launcher_sha256"], "allocation_requested": False}
    active, active_started, process, result_code = False, None, None, 1
    with (cache / "owned-shell-launcher.lock").open("a") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            raise LaunchError("another_launcher_owns_existing_shell") from None
        try:
            while True:
                snapshot = scheduler_snapshot(environment)
                active = active or snapshot["state"] in ACTIVE_STATES
                if active and active_started is None:
                    active_started = time.monotonic()
                write_status(status_path, state, "scheduler", snapshot=snapshot)
                if active_started is not None and time.monotonic() - active_started >= configuration["budget_seconds"]:
                    raise LaunchError("active_gpu_budget_exhausted")
                if snapshot["state"] in TERMINAL_STATES:
                    raise LaunchError("owned_allocation_ended_before_launch")
                if snapshot["state"] in READY_STATES:
                    break
                time.sleep(poll_seconds)
            started = active_started if active_started is not None else time.monotonic()
            # A source/input change while queued must not consume a GPU run.
            verify_inputs(configuration["plan"], configuration["plan_sha256"], configuration["gym_python"])
            connection = prepare_connection(environment, cache)
            remaining = configuration["budget_seconds"] - (time.monotonic() - started)
            if remaining <= 0:
                raise LaunchError("active_gpu_budget_exhausted")
            command = ["ssh", "-F", str(connection), "hessian-nemo-long-owned",
                       shlex.join([str(ROOT / ".venv-gpu/bin/python"), "-"])]
            logfile = Path(configuration["plan"]).parent / f"launcher-ssh-{time.time_ns()}.log"
            descriptor = os.open(logfile, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
            write_status(status_path, state, "worker_start", log_path=str(logfile))
            with os.fdopen(descriptor, "w") as output:
                process = subprocess.Popen(command, env=environment, stdin=subprocess.PIPE,
                                           stdout=output, stderr=subprocess.STDOUT, text=True, start_new_session=True)
                try:
                    process.communicate(node_program(configuration), timeout=remaining)
                except subprocess.TimeoutExpired:
                    raise LaunchError("active_gpu_budget_exhausted") from None
            if process.returncode:
                raise LaunchError("worker_program_failed")
            result_code = 0
            write_status(status_path, state, "completed", result="completed", remote_exit_code=0)
        except Interrupted as error:
            result_code = 128 + error.signum
            write_status(status_path, state, "interrupted", result="interrupted", before_allocation=not active)
        except BaseException as error:
            write_status(status_path, state, "failed", result="failed", error_code=error.code if isinstance(error, LaunchError)
                         else type(error).__name__)
        finally:
            if process is not None and process.poll() is None:
                process.terminate()
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    process.kill()
                    process.wait(timeout=10)
            if active:
                previous = {signum: signal.getsignal(signum) for signum in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP)}
                try:
                    for signum in previous:
                        signal.signal(signum, signal.SIG_IGN)
                    released = release_owned(environment)
                finally:
                    for signum, handler in previous.items():
                        signal.signal(signum, handler)
                write_status(status_path, state, "allocation_release", release_acknowledged=released)
                if not released:
                    result_code = 1
            else:
                write_status(status_path, state, "queue_preserved", release_acknowledged=False)
    return result_code


def main():
    global OWNED_SHELL
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--shell-id", default=OWNED_SHELL, help="explicit existing JAIF GPU shell; no default allocation")
    parser.add_argument("--plan", type=Path, required=True)
    parser.add_argument("--expected-plan-sha256", required=True)
    parser.add_argument("--gym-python", type=Path, default=ROOT / ".venv-nemo-gym/bin/python")
    parser.add_argument("--poll-seconds", type=int, default=15)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--check", action="store_true", help="local verification only; no scheduler/SSH actions")
    mode.add_argument("--launch", action="store_true", help="opt in to watching the explicitly selected existing owned shell")
    args = parser.parse_args()
    try:
        OWNED_SHELL = args.shell_id
        if args.launch:
            require_owned_shell()
        configuration = verify_inputs(args.plan, args.expected_plan_sha256, args.gym_python)
        if args.check:
            print(json.dumps({"status": "checked", "shell_id": OWNED_SHELL,
                              "plan_sha256": configuration["plan_sha256"], "cluster_contacted": False}))
            return 0
        def interrupted(signum, frame):
            raise Interrupted(signum)
        for signum in (signal.SIGINT, signal.SIGTERM, signal.SIGHUP):
            signal.signal(signum, interrupted)
        return launch(configuration, args.poll_seconds)
    except (LaunchError, OSError, ValueError) as error:
        print(json.dumps({"status": "failed", "error_code": error.code if isinstance(error, LaunchError)
                          else type(error).__name__}), file=sys.stderr)
        return 1


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