#!/usr/bin/env python3
"""Prometheus-file_sd-Ziele für die GPU-VMs erzeugen (läuft auf ops).

Die GPU-VMs bekommen ihre Adressen per DHCP (D5), deshalb stehen sie nicht fest in prometheus.yml.
Eingabe ist `tofu output -json gpu_vms` (stdin), dazu die Registry für die Slot-Labels des Healers.

    tofu output -json gpu_vms | python3 monitoring-targets.py --registry sites/knaster/registry.yaml --out DIR

Schreibt DIR/node.json (node_exporter, Port 9100), DIR/gpu.json (nvidia_gpu_exporter, Port 9835) und
DIR/llamacpp.json (llama-server /metrics je Dienst, Task 6.1).

Labels je Ziel: instance (VM-Name, bleibt bei IP-Wechsel gleich), vm, pve_node, vmid, slot, slots.
`slot` ist der erste Slot der Registry mit dieser vmid (Reihenfolge der Registry), `slots` alle,
kommagetrennt. VMs ohne Slot bekommen slot="none" (der Healer markiert dann keinen echten Slot).

llama.cpp: Welche VM welches Profil fährt, steht im Ansible-Inventar (Gruppe gpu_llamacpp), nicht in der
Registry (die kennt je Slot nur eine VM, die zweite Embedding-Instanz fehlt dort). Übergabe als JSON:
`--llamacpp '{"gpu-pve44": ["embed"], ...}'`. Die Slots eines Dienstes kommen aus der Registry
(engine llamacpp, gleicher Container). Labels zusätzlich docker_host (= VM) und container, damit der
Healer die tatsächliche VM trifft. Diese Slots zählen auch für node/gpu der VM.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import yaml

PORTS = {"node": 9100, "gpu": 9835}
# Dienste je Profil wie in deploy/compose/gpu-llamacpp (feste Ports, Host-Netz)
LLAMACPP_PROFILES = {
    "embed": [("llamacpp-embed", 8081), ("llamacpp-rerank", 8082)],
    "chat": [("llamacpp-chat", 8080)],
}


def slots_by_vmid(registry: dict) -> dict[str, list[str]]:
    out: dict[str, list[str]] = {}
    for name, slot in (registry.get("slots") or {}).items():
        vmid = str((slot.get("ops") or {}).get("vmid", ""))
        if vmid:
            out.setdefault(vmid, []).append(name)
    return out


def llamacpp_slots_by_container(registry: dict) -> dict[str, list[str]]:
    out: dict[str, list[str]] = {}
    for name, slot in (registry.get("slots") or {}).items():
        container = (slot.get("ops") or {}).get("container")
        if slot.get("engine") == "llamacpp" and container:
            out.setdefault(container, []).append(name)
    return out


def build_targets(gpu_vms: dict, registry: dict,
                  llamacpp: dict[str, list[str]] | None = None) -> dict[str, list[dict]]:
    by_vmid = slots_by_vmid(registry)
    by_container = llamacpp_slots_by_container(registry)
    llamacpp = llamacpp or {}
    result: dict[str, list[dict]] = {job: [] for job in [*PORTS, "llamacpp"]}
    for name in sorted(gpu_vms):
        vm = gpu_vms[name]
        if not vm.get("ipv4"):
            print(f"Warnung: {name} hat keine IPv4 (Gast-Agent?), übersprungen", file=sys.stderr)
            continue
        vmid = str(vm["vmid"])
        services = [svc for profile in llamacpp.get(name, []) for svc in LLAMACPP_PROFILES[profile]]
        slots = list(by_vmid.get(vmid, []))
        for container, _ in services:
            slots += [s for s in by_container.get(container, []) if s not in slots]
        labels = {
            "instance": name,
            "vm": name,
            "pve_node": vm["node"],
            "vmid": vmid,
            "slot": slots[0] if slots else "none",
            "slots": ",".join(slots),
        }
        for job, port in PORTS.items():
            result[job].append({"targets": [f"{vm['ipv4']}:{port}"], "labels": labels})
        for container, port in services:
            cslots = by_container.get(container, [])
            result["llamacpp"].append({"targets": [f"{vm['ipv4']}:{port}"], "labels": {
                **labels, "docker_host": name, "container": container,
                "slot": cslots[0] if cslots else "none", "slots": ",".join(cslots)}})
    return result


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--registry", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--llamacpp", default="{}", help='JSON {"<vm>": ["embed"|"chat", ...]} aus dem Inventar')
    args = ap.parse_args()

    gpu_vms = json.load(sys.stdin)
    registry = yaml.safe_load(args.registry.read_text(encoding="utf-8"))
    llamacpp = json.loads(args.llamacpp)
    unknown = {p for profiles in llamacpp.values() for p in profiles} - set(LLAMACPP_PROFILES)
    if unknown:
        ap.error(f"unbekannte llama.cpp-Profile: {sorted(unknown)}")
    targets = build_targets(gpu_vms, registry, llamacpp)
    args.out.mkdir(parents=True, exist_ok=True)
    for job, entries in targets.items():
        path = args.out / f"{job}.json"
        tmp = path.with_suffix(".tmp")
        tmp.write_text(json.dumps(entries, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        tmp.replace(path)
        print(f"{path}: {len(entries)} Ziele")
    return 0 if targets["node"] else 1


if __name__ == "__main__":
    sys.exit(main())