"""
ARIA LLM-Updater (AI-Box) — aktualisiert das llama-swap-Image (mit llama.cpp).

Der einzige Dienst mit Docker-Socket. Darum kann er bewusst NUR drei feste
Dinge, alle fuer genau einen Container (aria-llama-swap):

  check     installierte Version + Image-Pruefsumme, dazu die neueste aus der
            Registry (Pruefsumme des Tags) und das letzte llama-swap-Release
  update    altes Image sichern (Tag vor-update-<Zeit>), neues ziehen,
            Container neu anlegen, Selbsttest — scheitert der, automatisch zurueck
  rollback  das gesicherte Image wieder einsetzen, Container neu anlegen, Selbsttest

Kein freier Befehl, keine anderen Container. Erreichbar nur im Compose-Netz
(kein Port nach aussen), jede Anfrage braucht den Token, den der llm-adapter
in /models/.updater_token ablegt.

Neu angelegt wird ueber `docker compose … up --no-deps --force-recreate
llama-swap` mit der Compose-Datei des Stacks — so bleibt alles, was die
Compose-Datei festlegt (GPU-Runtime, Volumes, LD_LIBRARY_PATH …), gleich.
Relative Pfade loest Compose gegen --project-directory auf; das ist der
Host-Pfad aus dem Label des laufenden Containers, er muss im Container hier
nicht existieren.
"""

from __future__ import annotations

import http.server
import json
import logging
import os
import re
import subprocess
import threading
import time
import urllib.error
import urllib.request

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
logger = logging.getLogger("llm-updater")

CONTAINER = os.getenv("LLAMA_CONTAINER", "aria-llama-swap")
SERVICE = os.getenv("LLAMA_SERVICE", "llama-swap")
STACK_DIR = os.getenv("STACK_DIR", "/stack")
TOKEN_FILE = os.getenv("TOKEN_FILE", "/models/.updater_token")
GEN_CONFIG = os.getenv("LLAMA_GEN_CONFIG", "/models/llama-swap.config.yaml")
LLAMA_URL = os.getenv("LLAMA_URL", "http://llama-swap:8080").rstrip("/")
RELEASES_URL = os.getenv("RELEASES_URL", "https://api.github.com/repos/mostlygeek/llama-swap/releases/latest")
BACKUP_REPO = "llama-swap"            # gesicherte Images: llama-swap:vor-update-<Zeit>
KEEP_BACKUPS = 1
READY_TIMEOUT_S = int(os.getenv("READY_TIMEOUT_S", "900"))   # Modell(e) wieder geladen
PORT = int(os.getenv("PORT", "8090"))

_job = {"id": "", "op": "", "state": "idle", "log": [], "result": None}
_job_lock = threading.Lock()


# ── kleine Helfer ───────────────────────────────────────────────

def _run(args: list[str], timeout: int = 120) -> tuple[int, str]:
    p = subprocess.run(args, capture_output=True, text=True, timeout=timeout)
    return p.returncode, (p.stdout + p.stderr).strip()


def _log(msg: str) -> None:
    logger.info(msg)
    with _job_lock:
        _job["log"].append(f"{time.strftime('%H:%M:%S')} {msg}")
        _job["log"] = _job["log"][-60:]


def _inspect(ref: str, kind: str = "container") -> dict:
    rc, out = _run(["docker", kind, "inspect", ref])
    if rc != 0:
        return {}
    try:
        return json.loads(out)[0]
    except Exception:
        return {}


def _image_ref() -> str:
    return (_inspect(CONTAINER).get("Config") or {}).get("Image", "")


def _local_digest(ref: str) -> str:
    """Pruefsumme, unter der das Image gezogen wurde (Index/Manifest des Tags)."""
    repo = ref.rsplit(":", 1)[0] if ":" in ref.split("/")[-1] else ref
    for d in _inspect(ref, "image").get("RepoDigests") or []:
        if d.startswith(repo + "@"):
            return d.split("@", 1)[1]
    return ""


def _versions() -> dict:
    v = {}
    rc, out = _run(["docker", "exec", CONTAINER, "llama-swap", "-version"], 30)
    m = re.search(r"version:\s*(\S+)", out)
    if rc == 0 and m:
        v["llamaSwap"] = m.group(1)
        b = re.search(r"built at\s*(\S+)", out)
        if b:
            v["llamaSwapBuilt"] = b.group(1)
    rc, out = _run(["docker", "exec", CONTAINER, "llama-server", "--version"], 30)
    m = re.search(r"version:\s*(.+)", out)
    if rc == 0 and m:
        v["llamaCpp"] = m.group(1).strip()
    return v


def _backups() -> list[dict]:
    rc, out = _run(["docker", "images", BACKUP_REPO, "--format", "{{.Tag}}\t{{.ID}}\t{{.CreatedAt}}"])
    res = []
    for line in out.splitlines() if rc == 0 else []:
        tag, iid, created = (line.split("\t") + ["", "", ""])[:3]
        if tag.startswith("vor-update-"):
            res.append({"tag": f"{BACKUP_REPO}:{tag}", "id": iid, "created": created})
    return sorted(res, key=lambda b: b["tag"], reverse=True)


def _http(url: str, method: str = "GET", headers: dict | None = None, timeout: int = 20):
    req = urllib.request.Request(url, method=method, headers=headers or {})
    return urllib.request.urlopen(req, timeout=timeout)


def _remote_digest(ref: str) -> str:
    """Pruefsumme des Tags in der Registry (ohne Docker, anonymer Token)."""
    name, tag = ref.rsplit(":", 1) if ":" in ref.split("/")[-1] else (ref, "latest")
    parts = name.split("/")
    if "." in parts[0] or ":" in parts[0]:
        host, path = parts[0], "/".join(parts[1:])
    else:
        host, path = "registry-1.docker.io", name if "/" in name else f"library/{name}"
    url = f"https://{host}/v2/{path}/manifests/{tag}"
    accept = ", ".join([
        "application/vnd.oci.image.index.v1+json",
        "application/vnd.docker.distribution.manifest.list.v2+json",
        "application/vnd.oci.image.manifest.v1+json",
        "application/vnd.docker.distribution.manifest.v2+json",
    ])
    try:
        r = _http(url, "HEAD", {"Accept": accept})
        return r.headers.get("Docker-Content-Digest", "")
    except urllib.error.HTTPError as e:
        if e.code != 401:
            raise
        chal = e.headers.get("WWW-Authenticate", "")
    p = dict(re.findall(r'(\w+)="([^"]*)"', chal))
    tok_url = f"{p['realm']}?service={p.get('service', '')}&scope={p.get('scope', f'repository:{path}:pull')}"
    tok = json.load(_http(tok_url))
    token = tok.get("token") or tok.get("access_token") or ""
    r = _http(url, "HEAD", {"Accept": accept, "Authorization": f"Bearer {token}"})
    return r.headers.get("Docker-Content-Digest", "")


def _latest_release() -> dict:
    try:
        d = json.load(_http(RELEASES_URL, headers={"Accept": "application/vnd.github+json"}))
        return {"tag": d.get("tag_name", ""), "published": d.get("published_at", ""),
                "url": d.get("html_url", "")}
    except Exception as e:
        return {"error": str(e)[:120]}


def _status(remote: bool) -> dict:
    ref = _image_ref()
    if not ref:
        return {"ok": False, "error": f"Container {CONTAINER} nicht gefunden"}
    img = _inspect(ref, "image")
    st = {
        "ok": True, "image": ref,
        "installed": {**_versions(), "digest": _local_digest(ref),
                      "created": img.get("Created", "")},
        "backups": _backups(),
        "job": {k: _job[k] for k in ("id", "op", "state")},
    }
    if remote:
        try:
            st["latestDigest"] = _remote_digest(ref)
        except Exception as e:
            st["latestError"] = f"Registry nicht erreichbar: {str(e)[:120]}"
        st["release"] = _latest_release()
        if st.get("latestDigest") and st["installed"]["digest"]:
            st["updateAvailable"] = st["latestDigest"] != st["installed"]["digest"]
    return st


# ── Neu anlegen + Selbsttest ────────────────────────────────────

def _compose_args() -> list[str]:
    labels = (_inspect(CONTAINER).get("Config") or {}).get("Labels") or {}
    proj = labels.get("com.docker.compose.project", "")
    wdir = labels.get("com.docker.compose.project.working_dir", "")
    if not proj or not wdir:
        raise RuntimeError("Container wurde nicht mit Compose angelegt — Labels fehlen")
    args = ["docker", "compose", "--project-directory", wdir, "-p", proj,
            "-f", os.path.join(STACK_DIR, "docker-compose.yml"), "--profile", "llm"]
    env = os.path.join(STACK_DIR, ".env")
    if os.path.exists(env):
        args += ["--env-file", env]
    return args


def _recreate() -> None:
    rc, out = _run(_compose_args() + ["up", "-d", "--no-deps", "--force-recreate",
                                      "--pull", "never", SERVICE], 300)
    if rc != 0:
        raise RuntimeError(f"Container neu anlegen fehlgeschlagen: {out[-400:]}")
    _log("Container neu angelegt")


def _preloaded() -> list[str]:
    try:
        import yaml
        cfg = yaml.safe_load(open(GEN_CONFIG)) or {}
        return list(((cfg.get("hooks") or {}).get("on_startup") or {}).get("preload") or [])
    except Exception:
        return []


def _selftest() -> tuple[bool, str]:
    """Karten sichtbar (sonst rechnet llama.cpp still auf der CPU) und die
    vorher geladenen Modelle wieder bereit."""
    for _ in range(30):
        if (_inspect(CONTAINER).get("State") or {}).get("Running"):
            break
        time.sleep(2)
    else:
        return False, "Container laeuft nicht"
    rc, out = _run(["docker", "exec", CONTAINER, "llama-server", "--list-devices"], 60)
    cards = [l.strip() for l in out.splitlines() if re.match(r"\s*CUDA\d+", l)]
    if not cards:
        return False, f"llama.cpp sieht keine Grafikkarte (wuerde auf der CPU rechnen): {out[-300:]}"
    _log(f"Karten sichtbar: {len(cards)} ({'; '.join(c[:60] for c in cards)})")
    want = _preloaded()
    if not want:
        return True, f"{len(cards)} Karte(n) sichtbar, kein Modell vorgeladen"
    _log(f"warte, bis {', '.join(want)} geladen ist …")
    deadline = time.time() + READY_TIMEOUT_S
    last = ""
    while time.time() < deadline:
        try:
            run = json.load(_http(f"{LLAMA_URL}/running", timeout=10)).get("running") or []
            states = {r.get("model"): r.get("state") for r in run}
            last = ", ".join(f"{m}={states.get(m, '-')}" for m in want)
            if all(states.get(m) == "ready" for m in want):
                return True, f"{len(cards)} Karte(n), Modell(e) bereit: {', '.join(want)}"
        except Exception as e:
            last = str(e)[:80]
        time.sleep(5)
    return False, f"Modell nach {READY_TIMEOUT_S // 60} min nicht bereit ({last})"


def _drop_untagged(image_id: str) -> None:
    """Ein Image, das nach dem Zuruecksetzen keinen Namen mehr hat, belegt nur
    noch Platz (das llama-swap-Image hat ~14 GB)."""
    if image_id and not (_inspect(image_id, "image").get("RepoTags") or []):
        rc, _ = _run(["docker", "rmi", image_id])
        if rc == 0:
            _log("verworfene Version geloescht")


def _prune_backups() -> None:
    for b in _backups()[KEEP_BACKUPS:]:
        _run(["docker", "rmi", b["tag"]])
        _log(f"alte Sicherung entfernt: {b['tag']}")


# ── Jobs ────────────────────────────────────────────────────────

def _do_update() -> dict:
    ref = _image_ref()
    before = _status(remote=False)
    old_id = _inspect(ref, "image").get("Id", "")
    backup = f"{BACKUP_REPO}:vor-update-{time.strftime('%Y%m%d-%H%M%S')}"
    rc, out = _run(["docker", "tag", old_id, backup])
    if rc != 0:
        raise RuntimeError(f"Sicherung fehlgeschlagen: {out}")
    _log(f"altes Image gesichert als {backup}")
    _log(f"ziehe {ref} … (mehrere GB, dauert)")
    rc, out = _run(["docker", "pull", ref], 3600)
    if rc != 0:
        raise RuntimeError(f"Ziehen fehlgeschlagen: {out[-400:]}")
    new_id = _inspect(ref, "image").get("Id", "")
    if new_id == old_id:
        _run(["docker", "rmi", backup])
        return {"ok": True, "changed": False, "message": "Schon aktuell — nichts geaendert."}
    _recreate()
    ok, why = _selftest()
    if ok:
        _prune_backups()
        after = _status(remote=False)
        return {"ok": True, "changed": True, "before": before.get("installed"),
                "after": after.get("installed"), "selftest": why, "backup": backup}
    _log(f"Selbsttest fehlgeschlagen: {why} — setze zurueck")
    _run(["docker", "tag", backup, ref])
    _recreate()
    _drop_untagged(new_id)
    ok2, why2 = _selftest()
    return {"ok": False, "rolledBack": True, "error": f"Selbsttest fehlgeschlagen: {why}",
            "rollbackSelftest": why2 if ok2 else f"AUCH die alte Version laeuft nicht: {why2}"}


def _do_rollback() -> dict:
    bks = _backups()
    if not bks:
        raise RuntimeError("Keine gesicherte Version vorhanden")
    ref = _image_ref()
    cur_id = _inspect(ref, "image").get("Id", "")
    b = bks[0]
    rc, out = _run(["docker", "tag", b["tag"], ref])
    if rc != 0:
        raise RuntimeError(f"Zuruecksetzen fehlgeschlagen: {out}")
    _log(f"{b['tag']} wieder als {ref} gesetzt")
    _recreate()
    _drop_untagged(cur_id)
    ok, why = _selftest()
    return {"ok": ok, "selftest": why, "restored": b["tag"],
            "after": _status(remote=False).get("installed"),
            **({} if ok else {"error": f"Selbsttest fehlgeschlagen: {why}"})}


def _start_job(op: str) -> dict:
    with _job_lock:
        if _job["state"] == "running":
            return {"ok": False, "error": f"Es laeuft schon: {_job['op']}", "job": _job["id"]}
        _job.update(id=f"{op}-{int(time.time())}", op=op, state="running", log=[], result=None)
        jid = _job["id"]

    def work():
        try:
            res = _do_update() if op == "update" else _do_rollback()
        except Exception as e:
            logger.exception("%s fehlgeschlagen", op)
            res = {"ok": False, "error": str(e)[:500]}
        with _job_lock:
            _job["result"] = res
            _job["state"] = "done"

    threading.Thread(target=work, daemon=True).start()
    return {"ok": True, "job": jid}


# ── HTTP ────────────────────────────────────────────────────────

def _token() -> str:
    try:
        return open(TOKEN_FILE).read().strip()
    except OSError:
        return ""


class Handler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *a):
        pass

    def _reply(self, code: int, body: dict) -> None:
        data = json.dumps(body).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def _authorized(self) -> bool:
        tok = _token()
        if not tok or self.headers.get("X-Updater-Token", "") != tok:
            self._reply(403, {"ok": False, "error": "Token fehlt oder falsch"})
            return False
        return True

    def do_GET(self):
        if not self._authorized():
            return
        try:
            if self.path == "/status":
                self._reply(200, _status(remote=False))
            elif self.path == "/check":
                self._reply(200, _status(remote=True))
            elif self.path == "/job":
                with _job_lock:
                    self._reply(200, dict(_job))
            else:
                self._reply(404, {"ok": False, "error": "unbekannt"})
        except Exception as e:
            self._reply(500, {"ok": False, "error": str(e)[:300]})

    def do_POST(self):
        if not self._authorized():
            return
        if self.path in ("/update", "/rollback"):
            self._reply(200, _start_job(self.path[1:]))
        else:
            self._reply(404, {"ok": False, "error": "unbekannt"})


if __name__ == "__main__":
    logger.info("LLM-Updater auf :%d fuer Container %s", PORT, CONTAINER)
    http.server.ThreadingHTTPServer(("0.0.0.0", PORT), Handler).serve_forever()