"""Driving the training host over SSH, including getting the GPU back from the LLMs.

Training runs on a separate machine -- a Ryzen AI MAX+ 395 with a Radeon 8060S, whose
124 GB of memory is *unified*: the iGPU addresses it through GTT rather than owning
dedicated VRAM.  That is why the box can train on frames this large at all, and also why
it can hold nothing else at the time: llama.cpp keeps 115.6 GB of the 124 GB resident, so
a training run does not get a smaller allocation, it gets ``Memory in use`` and dies.

So something has to give up the GPU for the duration.  These services belong to the user,
not to us, which sets the bar for this module: stop as little as will do, put back
whatever it stopped -- including when the training run crashes -- and be able to say
afterwards whether that succeeded.

Stopping as little as will do is worth the extra code.  Measured on this host, almost
all of that 115.6 GB is a 120B model that ``llama-swap`` had loaded on demand; with the
unit stopped, GTT falls to 24.8 GB and roughly 99 GB is free -- ample for training,
without touching anything else.  So :func:`stop_llms` works in stages and escalates only
if the first stage leaves too little.

The two kinds of server are not alike:

* ``llama-swap.service``, a systemd --user unit.  Stopping and starting it is exact, and
  it is where the large on-demand models live.  Almost always sufficient on its own.
* a **standalone** ``llama-server``, launched by hand and reparented to init.  There is
  no unit to restart, so the only way back is to record its argv and working directory
  before killing it and re-exec them.  This is the fragile one, only touched when the
  first stage did not free enough, and the reason :func:`llms_paused` verifies the
  restore rather than assuming it.
"""

import json
import os
import shlex
import subprocess
import time
from contextlib import contextmanager

HOST = os.environ.get("SUVI_GPU_HOST", "htpc@192.168.1.66")
#: Read from a file rather than embedded, so the credential is not in the source tree.
PASSWORD_FILE = os.environ.get(
    "SUVI_GPU_PASSWORD_FILE",
    os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                 "gpu_system_ssh_info"),
)
#: systemd --user units that hold the GPU.
UNITS = ("llama-swap.service",)
#: Seconds to wait for GTT to drain / refill before giving up.
SETTLE_TIMEOUT = 120
#: Unified memory the iGPU can address, from mem_info_gtt_total on this host.
GTT_TOTAL = 115 * 1024**3
#: Free unified memory a training run needs before it will start.
GTT_FREE_BYTES = 48 * 1024**3
GTT_USED = "/sys/class/drm/card1/device/mem_info_gtt_used"


def _password():
    with open(PASSWORD_FILE) as handle:
        lines = [line.strip() for line in handle if line.strip()]
    if len(lines) < 2:
        raise RuntimeError(f"{PASSWORD_FILE}: expected an ssh line and a password line")
    return lines[1]


def run(command, check=True, timeout=600, capture=True):
    """Run a shell command on the training host."""
    argv = [
        "sshpass", "-p", _password(), "ssh",
        "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10",
        HOST, command,
    ]
    result = subprocess.run(
        argv, timeout=timeout,
        stdout=subprocess.PIPE if capture else None,
        stderr=subprocess.STDOUT if capture else None,
    )
    output = result.stdout.decode(errors="replace") if capture else ""
    if check and result.returncode != 0:
        raise RuntimeError(f"remote command failed ({result.returncode}): {command}\n{output}")
    return output


def gtt_used():
    """Bytes of unified memory currently mapped to the GPU."""
    try:
        return int(run(f"cat {GTT_USED}", check=False).strip() or 0)
    except ValueError:
        return 0


# --------------------------------------------------------------- pausing the LLMs


def _standalone_servers():
    """Every llama-server not owned by one of our units, with enough to restart it.

    A server spawned *by* llama-swap comes back when llama-swap does, so recording it
    would restart it twice.  Only the ones reparented to init are ours to restore.
    """
    script = r"""
for p in $(pgrep -x llama-server); do
  ppid=$(ps -o ppid= -p $p 2>/dev/null | tr -d ' ')
  [ "$ppid" = "1" ] || continue
  cwd=$(readlink /proc/$p/cwd)
  argv=$(tr '\0' '\n' < /proc/$p/cmdline | sed 's/"/\\"/g' | awk '{printf "\"%s\",", $0}')
  echo "{\"pid\": $p, \"cwd\": \"$cwd\", \"argv\": [${argv%,}]}"
done
"""
    found = []
    for line in run(script, check=False).splitlines():
        line = line.strip()
        if line.startswith("{"):
            try:
                found.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return found


def _wait_for_gtt(below, timeout=SETTLE_TIMEOUT):
    deadline = time.time() + timeout
    while time.time() < deadline:
        if gtt_used() < below:
            return True
        time.sleep(2)
    return False


def stop_llms(needed_bytes=None):
    """Free the GPU, stopping as little as will do.

    Stage one stops the units, which is where the big on-demand models live and which
    restarts exactly.  Only if that leaves less than `needed_bytes` free does stage two
    kill hand-launched servers, whose restoration is a re-exec rather than a restart.

    Returns the state :func:`start_llms` needs to undo this.
    """
    needed = GTT_TOTAL - (needed_bytes or GTT_FREE_BYTES)
    state = {"units": [], "standalone": []}

    for unit in UNITS:
        active = run(f"systemctl --user is-active {unit}", check=False).strip()
        state["units"].append({"unit": unit, "was_active": active == "active"})
        if active == "active":
            run(f"systemctl --user stop {unit}", check=False)
    if _wait_for_gtt(needed, timeout=30):
        return state

    # Stage two: the units alone were not enough.
    state["standalone"] = _standalone_servers()
    for server in state["standalone"]:
        run(f"kill {int(server['pid'])}", check=False)
    if not _wait_for_gtt(needed):
        for server in state["standalone"]:
            run(f"kill -9 {int(server['pid'])}", check=False)
        _wait_for_gtt(needed)
    return state


def start_llms(state):
    """Put back exactly what :func:`stop_llms` took away.

    Reports what it could not restore rather than raising: this runs in a `finally`,
    and masking the training error with a restore error would lose the more useful of
    the two.
    """
    failures = []
    for index, server in enumerate(state.get("standalone", [])):
        argv = " ".join(shlex.quote(arg) for arg in server["argv"])
        # systemd-run rather than nohup: a backgrounded process still holds the SSH
        # channel open, so ssh blocks until it exits -- which for a server is never.
        # A transient unit forks away cleanly and returns at once, and --collect means
        # it leaves nothing behind when it stops.
        started = run(
            f"systemd-run --user --collect --unit=suvi-restored-llama-{index} "
            f"--property=WorkingDirectory={shlex.quote(server['cwd'])} {argv}",
            check=False, timeout=60,
        )
        if "Running as unit" not in started and "Failed" in started:
            failures.append(f"standalone {os.path.basename(server['argv'][0])}")

    for entry in state.get("units", []):
        if not entry["was_active"]:
            continue
        run(f"systemctl --user start {entry['unit']}", check=False)
        if run(f"systemctl --user is-active {entry['unit']}", check=False).strip() != "active":
            failures.append(entry["unit"])
    return failures


#: States stopped but not yet restored.  Module level so the exit hooks can reach them.
_paused = []
_hooks_installed = False


def _restore_all(verbose=True):
    """Put back every outstanding pause.  Safe to call more than once."""
    while _paused:
        state = _paused.pop()
        # Our own containers first: restarting the LLMs while a training container still
        # holds 60 GB of GTT just moves the memory exhaustion onto the user's services.
        try:
            stop_containers(verbose=verbose)
        except (RuntimeError, OSError):
            pass
        failures = start_llms(state)
        if failures:
            print(f"  WARNING: could not restart: {', '.join(failures)}")
            print("           The GPU host is missing services it had before this run.")
        elif verbose:
            print(f"  restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used")


def _install_hooks():
    """Restore on the ways out that `finally` does not cover.

    A `finally` block handles a return or an exception.  It does not handle SIGTERM,
    which is what a timeout, a `kill`, or a parent shell giving up actually sends -- the
    default disposition terminates the interpreter without unwinding.  That is not
    hypothetical: it happened here, and left the user's LLM proxy stopped with nothing
    scheduled to start it again.  So the same belt-and-braces the archive traversals
    use (see :class:`suvi.vfs.Reliever`): atexit for orderly exits, explicit handlers
    for the signals.
    """
    global _hooks_installed
    if _hooks_installed:
        return
    import atexit
    import signal

    atexit.register(_restore_all, verbose=False)
    for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP):
        previous = signal.getsignal(signum)

        def handler(number, frame, previous=previous):
            _restore_all()
            if callable(previous) and previous not in (signal.SIG_IGN, signal.SIG_DFL):
                return previous(number, frame)
            raise SystemExit(128 + number)

        try:
            signal.signal(signum, handler)
        except (ValueError, OSError):  # not the main thread, or no such signal
            pass
    _hooks_installed = True


@contextmanager
def llms_paused(needed_bytes=None, verbose=True):
    """Hold the GPU for the duration of the block, then give it back.

    Restoration is registered before the block runs and happens on every exit path --
    return, exception, or signal.  What it cannot restore it names, because a
    hand-launched server that quietly failed to come back would otherwise be discovered
    by the user, later, as a service that is simply gone.
    """
    _install_hooks()
    state = stop_llms(needed_bytes)
    _paused.append(state)
    if verbose:
        units = [e["unit"] for e in state["units"] if e["was_active"]]
        print(f"  paused: units {units or 'none'}"
              + (f", {len(state['standalone'])} standalone server(s)"
                 if state["standalone"] else " (no standalone servers touched)"))
        print(f"  GTT now {gtt_used() / 1024**3:.1f} GiB used")
    try:
        yield state
    finally:
        if state in _paused:
            _paused.remove(state)
            stop_containers(verbose=verbose)
            failures = start_llms(state)
            if failures:
                print(f"  WARNING: could not restart: {', '.join(failures)}")
                print("           The GPU host is missing services it had before this run.")
            elif verbose:
                print(f"  restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used")


# ------------------------------------------------------------------ the container


#: ROCm+PyTorch plus this project's dependencies.  Bazzite's root is read-only and its
#: Python is 3.14, which has no torch wheels, so the toolchain is containerised
#: regardless; the derived tag adds zstandard, opencv, scikit-image and astropy.
IMAGE = os.environ.get("SUVI_GPU_IMAGE", "localhost/suvi-train:latest")
#: Prefix for containers this module starts, so they can be found and stopped again.
CONTAINER_PREFIX = "suvi-run-"


def stop_containers(verbose=True):
    """Stop every container this module started.

    Necessary because a container does *not* die with the SSH session that launched it.
    A run that times out client-side leaves the container running, holding tens of GB of
    unified memory; two such orphans were what put GTT at 89.6 GB with the GPU idle.
    """
    names = [
        line.strip()
        for line in run(
            f"podman ps --filter name={CONTAINER_PREFIX} --format '{{{{.Names}}}}'",
            check=False,
        ).splitlines()
        if line.strip().startswith(CONTAINER_PREFIX)
    ]
    for name in names:
        if verbose:
            print(f"  stopping orphaned container {name}")
        run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60)
    return names


def launch(command, mounts=(), name=None, log="/var/home/htpc/suvi/run.log"):
    """Start a container detached on the host and return (name, log path).

    Long jobs must not be tied to an SSH channel.  Streaming a training run through
    ``ssh`` means the run dies with the connection, or -- worse, and observed here --
    the container finishes while the local client stays blocked reading a pipe that
    will never close.  A transient systemd unit is owned by the host, survives the
    client entirely, and can be polled with :func:`tail`.
    """
    name = name or f"{CONTAINER_PREFIX}{int(time.time())}"
    binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts)
    inner = (
        f"podman run --rm --name {shlex.quote(name)} "
        f"--device /dev/kfd --device /dev/dri "
        f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}"
    )
    run(f"rm -f {shlex.quote(log)}", check=False)
    run(
        f"systemd-run --user --collect --unit={shlex.quote(name)} "
        f"--property=StandardOutput=append:{shlex.quote(log)} "
        f"--property=StandardError=append:{shlex.quote(log)} "
        f"/bin/sh -c {shlex.quote(inner)}",
        timeout=60,
    )
    return name, log


def tail(log, lines=40):
    return run(f"tail -n {int(lines)} {shlex.quote(log)} 2>/dev/null", check=False)


def running(name):
    return run(f"podman ps --filter name={shlex.quote(name)} --format '{{{{.Names}}}}'",
               check=False).strip() != ""


def podman(command, mounts=(), timeout=None, capture=True, name=None):
    """Run `command` inside the training container on the GPU host.

    ``/dev/kfd`` and ``/dev/dri`` are both world-accessible on this host, so no group
    mapping is needed; ``seccomp=unconfined`` is what ROCm needs to issue its ioctls.

    The container is named and stopped in a `finally`, so a client-side timeout cannot
    leave it running -- ``--rm`` only covers containers that actually exit.
    """
    name = name or f"{CONTAINER_PREFIX}{os.getpid()}"
    binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts)
    try:
        return run(
            f"podman run --rm --name {shlex.quote(name)} "
            f"--device /dev/kfd --device /dev/dri "
            f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}",
            timeout=timeout, capture=capture,
        )
    finally:
        run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60)