"""One environment per question.

Every task gets its own kubeadm cluster: a kind cluster whose control-plane node is the host the
task tells you to `ssh` to, as the real exams do since 2025. Nothing is shared between tasks, so a
task may break its kubelet, roll etcd back or lock the API server down without touching any other
task - and resetting a task is deleting its cluster and making it again.

Graders stay on this machine: each environment has its own kubeconfig, and its nodes are reached
with `podman exec`. A task's files live in /opt/course on its host, which is a directory of this
machine mounted into the control-plane node, so a grader reads them as plain files.
"""
import os
import re
import shutil
import subprocess
import time

RUNTIME = "podman"
NETWORK = "ksim"                 # podman network shared by every exam host and the desktop
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
STATE_DIR = os.path.join(ROOT, ".state")
ENV_DIR = os.path.join(STATE_DIR, "envs")
SSH_DIR = os.path.join(STATE_DIR, "ssh")
COURSE = "/opt/course"           # a task's files, as its host sees them
# A host is <track><question id>: cka0037. kind names the cluster after it, so its nodes are
# cka0037-control-plane, cka0037-worker, ... Only names of this shape are ever deleted.
HOST_RE = re.compile(r"^(cka|ckad|cks)\d{4}$")

# The kind node image each exam host image is built from, one per Kubernetes minor an exam runs.
# Pinned by digest: kind only guarantees an image built for its own release (these are v0.33.0's).
BASE_IMAGES = {
    "v1.35": "docker.io/kindest/node:v1.35.8@sha256:07b2536e30b803ed61d1677a79df6115f798ce64c80f9e22f6ed45afd09323c0",
    "v1.37": "docker.io/kindest/node:v1.37.0@sha256:a1ed56cfb0e7b93589bdf97c8cd566405a265939e3620fc4f5de89adff580ae5",
}
HELM_VERSION = "v4.3.0"
DESKTOP_IMAGE = "localhost/ksim-desktop:latest"
# Pull-through cache for Docker Hub that every exam host's containerd uses as a mirror
# (images/node/rootfs/etc/containerd/certs.d). Its storage is a named volume, so it outlives
# `./ksim down`.
REGISTRY = "ksim-registry"
REGISTRY_IMAGE = "docker.io/library/registry:2"
# Memory a node costs once its cluster is up and seeded, for the check before a session is
# provisioned. Measured on a 20-task CKA mock (20 control planes, 30 workers): 8 GB in all, about
# 320 MB and 60 MB a node; these leave room for the tasks' own Pods.
MB_CONTROL_PLANE, MB_WORKER = 400, 100

# Exam hosts that are virtual machines (see images/vm): for tasks that need a kernel of their own -
# Falco's eBPF probe, AppArmor. One Ubuntu VM per task, run by QEMU in a container that carries
# the host's name. The base disk is built once per Kubernetes minor from the Ubuntu cloud image.
VM_IMAGE = "localhost/ksim-vm:latest"
VM_DIR = os.path.join(STATE_DIR, "vm")
UBUNTU_IMAGE = "https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img"
MB_VM = 2300                     # what QEMU is given (2 GB) and its own overhead


class EnvError(Exception):
    pass


def _env():
    return dict(os.environ, KIND_EXPERIMENTAL_PROVIDER=RUNTIME, KIND_EXPERIMENTAL_PODMAN_NETWORK=NETWORK)


def sh(args, stdin=None, timeout=600):
    """Run a command; (returncode, combined output). Never raises on a non-zero exit."""
    try:
        p = subprocess.run(args, input=stdin, capture_output=True, text=True, timeout=timeout, env=_env(),
                           stdin=None if stdin is not None else subprocess.DEVNULL)
    except subprocess.TimeoutExpired:
        return 124, f"timed out after {timeout}s: {' '.join(args[:4])}"
    return p.returncode, (p.stdout + p.stderr).strip()


# ----------------------------------------------------------- images ----
def node_image(k8s, variant=None):
    """The exam host image of a Kubernetes minor; `variant` names a track's own flavour of it
    (images/node-<variant>, built on top of the common one) - CKS adds its tools that way."""
    return f"localhost/ksim-node:{k8s}" + (f"-{variant}" if variant else "")


def image_exists(name):
    return sh([RUNTIME, "image", "exists", name])[0] == 0


def build_node_image(k8s, log=print):
    base = BASE_IMAGES.get(k8s)
    if not base:
        raise EnvError(f"no kind node image is pinned for Kubernetes {k8s} (BASE_IMAGES in sim/envs.py)")
    log(f"  {node_image(k8s)}: from {base.split('@')[0]}")
    # etcdctl/etcdutl come out of the etcd image this kubeadm would deploy, so they match the server
    rc, out = sh([RUNTIME, "run", "--rm", "--entrypoint", "sh", base, "-c",
                  'kubeadm config images list --kubernetes-version "$(cat /kind/version)" 2>/dev/null | grep /etcd:'])
    etcd = out.strip().splitlines()[-1] if rc == 0 and out.strip() else ""
    if not etcd:
        raise EnvError(f"could not read the etcd image of {base}:\n{out[-400:]}")
    rc, out = sh([RUNTIME, "build", "-t", node_image(k8s), "--build-arg", f"BASE={base}", "--build-arg", f"ETCD_IMAGE={etcd}",
                  "--build-arg", f"HELM_VERSION={HELM_VERSION}", os.path.join(ROOT, "images", "node")], timeout=1800)
    if rc != 0:
        raise EnvError(f"building {node_image(k8s)} failed:\n{out[-1500:]}")
    log(f"  {node_image(k8s)}: built (etcd tools from {etcd})")


def build_variant_image(k8s, variant, log=print):
    """A track's flavour of the exam host image, built FROM the common one (which must exist)."""
    ctx = os.path.join(ROOT, "images", f"node-{variant}")
    if not os.path.isdir(ctx):
        raise EnvError(f"no image directory for the '{variant}' exam host (images/node-{variant})")
    rc, out = sh([RUNTIME, "run", "--rm", "--entrypoint", "cat", node_image(k8s), "/kind/version"])
    version = out.strip().splitlines()[-1].lstrip("v") if rc == 0 and out.strip() else ""
    if not re.match(r"^\d+\.\d+\.\d+$", version):
        raise EnvError(f"could not read the Kubernetes version of {node_image(k8s)}:\n{out[-300:]}")
    name = node_image(k8s, variant)
    log(f"  {name}: from {node_image(k8s)} (Kubernetes {version} as packages, {variant.upper()} tools)")
    rc, out = sh([RUNTIME, "build", "-t", name, "--build-arg", f"BASE={node_image(k8s)}",
                  "--build-arg", f"K8S_VERSION={version}", "--build-arg", f"K8S_MINOR={k8s}", ctx], timeout=3600)
    if rc != 0:
        raise EnvError(f"building {name} failed:\n{out[-1500:]}")
    log(f"  {name}: built")


def vm_base(k8s):
    """The disk every VM exam host of a Kubernetes minor starts from."""
    return os.path.join(VM_DIR, f"base-{k8s}.qcow2")


def vm_ready(k8s):
    """VM hosts can be made: the base disk is there (the runner image is checked when one is made)."""
    return os.path.exists(vm_base(k8s))


def build_vm_image(k8s, log=print):
    """The VM runner image, and the base disk: the Ubuntu cloud image booted once with
    images/vm/guest/provision.sh as its first-boot script. Ten minutes or so, and about 4 GB."""
    m = re.search(r":v([\d.]+)", BASE_IMAGES.get(k8s, ""))
    if not m:
        raise EnvError(f"no Kubernetes version is pinned for {k8s} (BASE_IMAGES in sim/envs.py)")
    if not os.path.exists("/dev/kvm") or not os.access("/dev/kvm", os.R_OK | os.W_OK):
        raise EnvError("VM exam hosts need KVM: /dev/kvm is missing or not accessible to you")
    log(f"  {VM_IMAGE}: QEMU, to run exam hosts that have to be real machines")
    rc, out = sh([RUNTIME, "build", "-t", VM_IMAGE, os.path.join(ROOT, "images", "vm")], timeout=1800)
    if rc != 0:
        raise EnvError(f"building {VM_IMAGE} failed:\n{out[-1500:]}")
    os.makedirs(VM_DIR, exist_ok=True)
    cloud = os.path.join(VM_DIR, "noble.img")
    if not os.path.exists(cloud):
        log(f"  downloading the Ubuntu 24.04 cloud image (about 600 MB)")
        rc, out = sh(["curl", "-fsSL", "-o", cloud + ".part", UBUNTU_IMAGE], timeout=3600)
        if rc != 0:
            raise EnvError(f"could not download {UBUNTU_IMAGE}:\n{out[-300:]}")
        os.replace(cloud + ".part", cloud)
    build = os.path.join(VM_DIR, "build")
    shutil.rmtree(build, ignore_errors=True)
    shutil.copytree(os.path.join(ROOT, "images", "vm", "guest"), build)
    with open(os.path.join(build, "env"), "w") as f:
        f.write(f"K8S_VERSION={m.group(1)}\nK8S_MINOR={k8s}\n")
    ensure_network()
    ensure_registry()
    log(f"  {os.path.relpath(vm_base(k8s), ROOT)}: Ubuntu 24.04, Kubernetes {m.group(1)}, Falco, AppArmor - provisioning in a VM, about ten minutes")
    rc, out = sh([RUNTIME, "run", "--rm", "--name", "ksim-vm-build", "--network", NETWORK, "--device", "/dev/kvm",
                  "--security-opt", "label=disable", "-e", f"KSIM_BASE={os.path.basename(vm_base(k8s))}",
                  "-v", f"{VM_DIR}:/base", "-v", f"{build}:/build:ro", VM_IMAGE, "build"], timeout=3000)
    shutil.rmtree(build, ignore_errors=True)
    if rc != 0 or not os.path.exists(vm_base(k8s)):
        raise EnvError(f"building the VM base disk failed (the VM's console is in {os.path.join(VM_DIR, 'build.log')}):\n{out[-1500:]}")
    log(f"  {os.path.relpath(vm_base(k8s), ROOT)}: built ({os.path.getsize(vm_base(k8s)) / 2**30:.1f} GB)")


def build_desktop_image(log=print):
    log(f"  {DESKTOP_IMAGE}: XFCE, terminal, Firefox, noVNC")
    rc, out = sh([RUNTIME, "build", "-t", DESKTOP_IMAGE, os.path.join(ROOT, "images", "desktop")], timeout=1800)
    if rc != 0:
        raise EnvError(f"building {DESKTOP_IMAGE} failed:\n{out[-1500:]}")
    log(f"  {DESKTOP_IMAGE}: built")
    if not image_exists(REGISTRY_IMAGE):
        rc, out = sh([RUNTIME, "pull", "-q", REGISTRY_IMAGE], timeout=600)
        log(f"  {REGISTRY_IMAGE}: " + ("pulled (the Docker Hub cache)" if rc == 0 else f"could not be pulled - hosts will pull from Docker Hub directly\n{out[-200:]}"))


# --------------------------------------------------- network and keys ----
def ensure_network():
    if sh([RUNTIME, "network", "exists", NETWORK])[0] != 0:
        rc, out = sh([RUNTIME, "network", "create", NETWORK])
        if rc != 0 and "already exists" not in out:
            raise EnvError(f"could not create the podman network '{NETWORK}': {out}")


def ensure_registry():
    """The Docker Hub cache, running. Without it hosts still pull - straight from Docker Hub."""
    ensure_network()
    rc, out = sh([RUNTIME, "inspect", "-f", "{{.State.Running}}", REGISTRY])
    if rc == 0 and out.strip() == "true":
        return True
    if rc == 0:
        return sh([RUNTIME, "start", REGISTRY])[0] == 0
    rc, out = sh([RUNTIME, "run", "-d", "--name", REGISTRY, "--network", NETWORK, "-v", f"{REGISTRY}:/var/lib/registry",
                  "-e", "REGISTRY_PROXY_REMOTEURL=https://registry-1.docker.io", REGISTRY_IMAGE], timeout=300)
    return rc == 0


def stop_registry():
    sh([RUNTIME, "rm", "-f", REGISTRY])


def ensure_keys():
    """The key pair the desktop logs in to every exam host with. Returns (private, public) text."""
    key = os.path.join(SSH_DIR, "id_ed25519")
    if not os.path.exists(key):
        os.makedirs(SSH_DIR, mode=0o700, exist_ok=True)
        rc, out = sh(["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-C", "ksim", "-f", key])
        if rc != 0:
            raise EnvError(f"ssh-keygen failed: {out}")
    with open(key) as f, open(key + ".pub") as g:
        return f.read(), g.read().strip()


# ------------------------------------------------------------- names ----
def host_of(track_id, qid):
    return f"{track_id}{int(qid):04d}"


def node_names(host, workers):
    """CP/W1/W2/W3 -> node (and container) name. Workers the host does not have map to ''."""
    names = {"CP": f"{host}-control-plane"}
    for i in (1, 2, 3):
        names[f"W{i}"] = (f"{host}-worker" + ("" if i == 1 else str(i))) if i <= workers else ""
    return names


def env_dir(host):
    return os.path.join(ENV_DIR, host)


def kubeconfig(host):
    return os.path.join(env_dir(host), "kubeconfig")


def course_dir(host):
    return os.path.join(env_dir(host), "course")


def memory_mb(workers, vm=False):
    return MB_VM if vm else MB_CONTROL_PLANE + MB_WORKER * workers


def is_vm(host):
    """The host is a virtual machine (create_vm leaves a marker in its directory)."""
    return os.path.exists(os.path.join(env_dir(host), "vm"))


def available_mb():
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemAvailable:"):
                    return int(line.split()[1]) // 1024
    except OSError:
        pass
    return None


# --------------------------------------------------------- lifecycle ----
def existing():
    """Exam hosts that exist right now: the clusters kind knows, and the VM hosts' containers."""
    rc, out = sh(["kind", "get", "clusters"])
    hosts = {n for n in out.split() if HOST_RE.match(n)} if rc == 0 else set()
    rc, out = sh([RUNTIME, "ps", "-a", "--filter", "label=ksim.vm", "--format", '{{index .Labels "ksim.vm"}}'])
    if rc == 0:
        hosts |= {n for n in out.split() if HOST_RE.match(n)}
    return sorted(hosts)


def running(host):
    rc, out = sh([RUNTIME, "inspect", "-f", "{{.State.Running}}", f"{host}-control-plane"])
    return rc == 0 and out.strip() == "true"


def start(host):
    """Start the stopped nodes of a host - rootless podman does not bring them back after a reboot."""
    if is_vm(host):      # the VM boots from its disk again: same machine, same cluster
        return sh([RUNTIME, "start", f"{host}-control-plane"])[0] == 0
    rc, out = sh(["kind", "get", "nodes", "--name", host])
    nodes = [n for n in out.split() if n.startswith(host + "-")] if rc == 0 else []
    for n in sorted(nodes, key=lambda n: "control-plane" not in n):
        sh([RUNTIME, "start", n])
    return bool(nodes)


def node_exec(node, script, timeout=120):
    """Run a bash script as root on a node. For a VM host the node's container only runs QEMU:
    the script goes on into the machine over ssh."""
    into = ["ssh", "vm", "bash -s"] if is_vm(node.rsplit("-control-plane", 1)[0]) else ["bash", "-s"]
    return sh([RUNTIME, "exec", "-i", node, *into], stdin=script, timeout=timeout)


_ACCESS = r"""
set -e
install -d -m700 -o candidate -g candidate /home/candidate/.ssh
install -d -m700 /root/.ssh
cat > /home/candidate/.ssh/authorized_keys <<'KSIM'
%(pub)s
KSIM
cat > /home/candidate/.ssh/id_ed25519 <<'KSIM'
%(priv)sKSIM
cp /home/candidate/.ssh/authorized_keys /home/candidate/.ssh/id_ed25519 /root/.ssh/
chown candidate:candidate /home/candidate/.ssh/authorized_keys /home/candidate/.ssh/id_ed25519
chmod 600 /home/candidate/.ssh/* /root/.ssh/*
if [ -f /etc/kubernetes/admin.conf ]; then
  install -d -o candidate -g candidate /home/candidate/.kube
  install -m600 -o candidate -g candidate /etc/kubernetes/admin.conf /home/candidate/.kube/config
  install -d /root/.kube
  install -m600 /etc/kubernetes/admin.conf /root/.kube/config
fi
systemctl is-active --quiet ssh || systemctl start ssh
"""


def _why(out):
    """The lines of a failed `kind create` that say what went wrong, without kubeadm's stack trace."""
    lines = [l.strip() for l in out.splitlines() if l.strip() and not l.startswith(("\t", " "))
             and not re.match(r"^(github\.com|k8s\.io|sigs\.k8s\.io|runtime\.|main\.|Stack Trace)", l)]
    keep = [l for l in lines if re.search(r"error|fail|timed out|unable|cannot|refused|no such", l, re.I)]
    return "\n".join((keep or lines)[-6:])[-900:]


def create(host, image, workers=0):
    """Make the cluster of one question and open it to `ssh` from the desktop. Idempotent only in
    the sense that a leftover cluster of the same name is replaced."""
    if not HOST_RE.match(host):
        raise EnvError(f"'{host}' is not an exam host name")
    priv, pub = ensure_keys()
    ensure_network()
    destroy(host)
    os.makedirs(course_dir(host), exist_ok=True)
    cfg = os.path.join(env_dir(host), "kind.yaml")
    with open(cfg, "w") as f:
        f.write("kind: Cluster\napiVersion: kind.x-k8s.io/v1alpha4\nnodes:\n"
                "  - role: control-plane\n    extraMounts:\n"
                f"      - hostPath: {course_dir(host)}\n        containerPath: {COURSE}\n"
                + "  - role: worker\n" * workers)
    out = ""
    for attempt in (1, 2):
        rc, out = sh(["kind", "create", "cluster", "--name", host, "--image", image, "--config", cfg,
                      "--kubeconfig", kubeconfig(host), "--wait", "240s"], timeout=600)
        if rc == 0:
            break
        sh(["kind", "delete", "cluster", "--name", host, "--kubeconfig", kubeconfig(host)])
        if attempt == 2:
            raise EnvError(f"kind could not create {host}:\n{_why(out)}")
    names = node_names(host, workers)
    for node in [n for n in names.values() if n]:
        rc, out = node_exec(node, _ACCESS % {"pub": pub, "priv": priv})
        if rc != 0:
            raise EnvError(f"could not open {node} to ssh:\n{out[-400:]}")
    return names


def create_vm(host, k8s):
    """A VM exam host: a single-node kubeadm cluster on Ubuntu, in a virtual machine of its own.
    Returns its node names, as create() does."""
    if not HOST_RE.match(host):
        raise EnvError(f"'{host}' is not an exam host name")
    if not vm_ready(k8s) or not image_exists(VM_IMAGE):
        raise EnvError("the VM exam host is not built yet:  ./ksim build cks --vm")
    ensure_keys()
    ensure_network()
    destroy(host)
    os.makedirs(course_dir(host), exist_ok=True)
    open(os.path.join(env_dir(host), "vm"), "w").close()
    node = f"{host}-control-plane"
    rc, out = sh([RUNTIME, "run", "-d", "--name", node, "--hostname", node, "--network", NETWORK, "--device", "/dev/kvm",
                  "--security-opt", "label=disable", "--label", f"ksim.vm={host}", "-p", "127.0.0.1::6443",
                  "-e", f"KSIM_HOSTNAME={node}", "-e", f"KSIM_BASE={os.path.basename(vm_base(k8s))}",
                  "-v", f"{VM_DIR}:/base:ro", "-v", f"{env_dir(host)}:/env", "-v", f"{course_dir(host)}:/course",
                  "-v", f"{SSH_DIR}:/keys:ro", VM_IMAGE, "run"])
    if rc != 0:
        raise EnvError(f"could not start the VM of {host}:\n{out[-400:]}")

    def console():
        try:
            with open(os.path.join(env_dir(host), "console.log"), errors="replace") as f:
                return "".join(f.readlines()[-12:])
        except OSError:
            return ""
    deadline = time.time() + 480      # boot, then kubeadm init on the first boot: usually about 90 s
    while True:
        if sh([RUNTIME, "exec", node, "ssh", "vm", "test -f /var/lib/ksim-ready"], timeout=30)[0] == 0:
            break
        if not running(host):
            raise EnvError(f"the VM of {host} stopped while starting:\n{console()}")
        if time.time() > deadline:
            raise EnvError(f"the VM of {host} did not get its cluster up in 8 minutes:\n{console()}")
        time.sleep(3)
    # the graders' kubeconfig: the VM's API server, through the port its container publishes
    rc, port = sh([RUNTIME, "port", node, "6443/tcp"])
    rc2, conf = sh([RUNTIME, "exec", node, "ssh", "vm", "cat /etc/kubernetes/admin.conf"], timeout=30)
    if rc != 0 or rc2 != 0 or "server:" not in conf:
        raise EnvError(f"could not read the kubeconfig of {host}:\n{(port + conf)[-300:]}")
    with open(os.open(kubeconfig(host), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w") as f:
        f.write(re.sub(r"server: https://\S+", f"server: https://{port.strip().splitlines()[0]}", conf) + "\n")
    return node_names(host, 0)


def hand_over(host, namespace=None):
    """After seeding: the task's files belong to `candidate`, and the host's kubeconfig starts in
    `namespace` when the task says to work in "the current namespace"."""
    script = f"chown -R candidate:candidate {COURSE}\n"
    if namespace:
        script += (f"for kc in /home/candidate/.kube/config /root/.kube/config; do\n"
                   f"  kubectl --kubeconfig $kc config set-context --current --namespace={namespace} >/dev/null\ndone\n")
    return node_exec(f"{host}-control-plane", script)


def destroy(host):
    if not HOST_RE.match(host):
        raise EnvError(f"refusing to delete '{host}': not an exam host name")
    sh(["kind", "delete", "cluster", "--name", host, "--kubeconfig", kubeconfig(host)], timeout=180)
    # a VM host is one container that kind knows nothing about; only one that carries this host's
    # label is removed
    rc, out = sh([RUNTIME, "inspect", "-f", '{{index .Config.Labels "ksim.vm"}}', f"{host}-control-plane"])
    if rc == 0 and out.strip() == host:
        sh([RUNTIME, "rm", "-f", "-t", "3", f"{host}-control-plane"], timeout=60)
    d = env_dir(host)
    if os.path.isdir(d):
        # files written by `candidate` on the host belong to a sub-uid here; remove them as the
        # user namespace's root
        sh([RUNTIME, "unshare", "rm", "-rf", d])
        shutil.rmtree(d, ignore_errors=True)