#!/usr/bin/env python3
"""Read the running LLM containers and write Kubernetes manifests.

Secrets stay on the server. This script prints service names, ports, and
hostnames only.
"""
import json
import os
import re
import subprocess
from pathlib import Path

OUT = Path("/usr/cluster/k8s/manifests")
NODE = "prod102"
NODE_IP = "10.0.1.47"
NS_BY_PROJECT = {
    "dify": "dify",
    "langfuse": "langfuse",
    "observability": "observability",
    "new-api": "llm",
}
GPU_CONTAINERS = {"vllm-qwen3.6-27b", "xinference-bge", "dcgm-exporter"}
SKIP_PREFIXES = ("k8s_",)


def run(args):
    return subprocess.check_output(args, text=True)


def inspect(name):
    return json.loads(run(["docker", "inspect", name]))[0]


def dns1123(name):
    out = name.lower().replace("_", "-")
    out = re.sub(r"[^a-z0-9-]", "-", out).strip("-")
    return out[:63] or "app"


def port_num(key):
    return int(key.split("/")[0])


def collect_ports(data):
    bindings = data["HostConfig"].get("PortBindings") or {}
    exposed = list((data["Config"].get("ExposedPorts") or {}).keys())
    rows = []
    seen = set()
    for key, binds in bindings.items():
        cport = port_num(key)
        seen.add(cport)
        if not binds:
            rows.append({"containerPort": cport, "protocol": "TCP"})
            continue
        for bind in binds:
            row = {"containerPort": cport, "protocol": "TCP"}
            if bind.get("HostPort"):
                row["hostPort"] = int(bind["HostPort"])
            if bind.get("HostIp") not in (None, "", "0.0.0.0"):
                row["hostIP"] = bind["HostIp"]
            rows.append(row)
    for key in exposed:
        cport = port_num(key)
        if cport not in seen:
            rows.append({"containerPort": cport, "protocol": "TCP"})
            seen.add(cport)
    return rows


def listen_ports(name):
    try:
        text = run(["docker", "exec", name, "sh", "-c", "ss -lnt 2>/dev/null || netstat -lnt"])
    except subprocess.CalledProcessError:
        return []
    found = []
    for line in text.splitlines():
        match = re.search(r":(\d+)\s", line)
        if not match:
            continue
        port = int(match.group(1))
        if port > 0 and port not in found and port not in (22,):
            found.append(port)
    return found


def namespace_for(data, name):
    labels = data["Config"].get("Labels") or {}
    project = labels.get("com.docker.compose.project")
    if project in NS_BY_PROJECT:
        return NS_BY_PROJECT[project]
    if name.startswith("vllm") or name.startswith("xinference") or name == "llm-nav":
        return "llm"
    return "llm"


SHORT = {
    "vllm-qwen3.6-27b": "vllm",
    "xinference-bge": "xinference",
    "llm-nav": "nav",
}


def service_name(data, name):
    labels = data["Config"].get("Labels") or {}
    raw = labels.get("com.docker.compose.service") or SHORT.get(name) or name
    return raw, dns1123(raw)


def env_map(data, name):
    env = {}
    for item in data["Config"].get("Env") or []:
        if "=" not in item:
            continue
        key, value = item.split("=", 1)
        if key in ("PATH", "LD_LIBRARY_PATH", "NVIDIA_REQUIRE_CUDA"):
            continue
        value = value.replace("172.17.0.1", NODE_IP).replace("host.docker.internal", NODE_IP)
        env[key] = value
    if name in GPU_CONTAINERS:
        env["NVIDIA_VISIBLE_DEVICES"] = "all"
        env.setdefault("NVIDIA_DRIVER_CAPABILITIES", "all")
    return env


def hostnames(env):
    found = set()
    for value in env.values():
        for match in re.findall(r"://([A-Za-z0-9_.-]+)", value):
            if match.replace(".", "").isdigit():
                continue
            found.add(match.split(":")[0])
    return sorted(found)


def volumes_for(mounts):
    volumes = []
    mounts_out = []
    for index, mount in enumerate(mounts):
        source = mount.get("Source")
        target = mount.get("Destination")
        if not source or not target:
            continue
        vol_name = "vol%d" % index
        host = {"path": source}
        if os.path.isfile(source):
            host["type"] = "File"
        elif os.path.isdir(source):
            host["type"] = "Directory"
        else:
            host["type"] = "DirectoryOrCreate"
        volumes.append({"name": vol_name, "hostPath": host})
        item = {"name": vol_name, "mountPath": target}
        if mount.get("RW") is False:
            item["readOnly"] = True
        mounts_out.append(item)
    return volumes, mounts_out


def container_spec(data, name, app, ports):
    ep = data["Config"].get("Entrypoint") or None
    cmd = data["Config"].get("Cmd") or None
    container = {
        "name": app,
        "image": data["Config"]["Image"],
        "imagePullPolicy": "IfNotPresent",
        "envFrom": [{"secretRef": {"name": app + "-env"}}],
        "ports": ports,
    }
    if ep:
        container["command"] = ep
        if cmd:
            container["args"] = cmd
    elif cmd:
        container["command"] = cmd
    workdir = data["Config"].get("WorkingDir")
    if workdir:
        container["workingDir"] = workdir
    volumes, mounts = volumes_for(data.get("Mounts") or [])
    image = data["Config"].get("Image") or ""
    patch = Path("/usr/cluster/k8s/patches/langfuse_trace.py")
    patch_path = Path("/usr/cluster/k8s/patches/langfuse_trace.path")
    if "dify-api" in image and patch.is_file() and patch_path.is_file():
        target = patch_path.read_text(encoding="utf-8").strip()
        if target:
            volumes.append({
                "name": "lfpatch",
                "hostPath": {"path": str(patch), "type": "File"},
            })
            mounts.append({"name": "lfpatch", "mountPath": target})
    shm = int(data["HostConfig"].get("ShmSize") or 0)
    if shm > 64 * 1024 * 1024:
        volumes.append({
            "name": "dshm",
            "emptyDir": {"medium": "Memory", "sizeLimit": "8Gi"},
        })
        mounts.append({"name": "dshm", "mountPath": "/dev/shm"})
    if mounts:
        container["volumeMounts"] = mounts
    caps = data["HostConfig"].get("CapAdd") or []
    if caps:
        container["securityContext"] = {"capabilities": {"add": caps}}
    return container, volumes


def workload(ns, app, raw_service, data, name, ports):
    container, volumes = container_spec(data, name, app, ports)
    pod_spec = {
        "nodeName": NODE,
        "enableServiceLinks": False,
        "hostAliases": [{
            "ip": NODE_IP,
            "hostnames": ["host.docker.internal"],
        }],
        "containers": [container],
    }
    if volumes:
        pod_spec["volumes"] = volumes
    return {
        "apiVersion": "apps/v1",
        "kind": "Deployment",
        "metadata": {"name": app, "namespace": ns, "labels": {"app": app}},
        "spec": {
            "replicas": 1,
            "strategy": {"type": "Recreate"},
            "selector": {"matchLabels": {"app": app}},
            "template": {
                "metadata": {
                    "labels": {"app": app},
                    "annotations": {"llm.local/compose-service": raw_service},
                },
                "spec": pod_spec,
            },
        },
    }


def service(ns, app, raw_service, ports):
    sport = []
    seen = set()
    for row in ports:
        cport = row["containerPort"]
        if cport in seen:
            continue
        seen.add(cport)
        sport.append({
            "name": "p%d" % cport,
            "port": cport,
            "targetPort": cport,
            "protocol": "TCP",
        })
    if not sport:
        return None
    return {
        "apiVersion": "v1",
        "kind": "Service",
        "metadata": {
            "name": app,
            "namespace": ns,
            "annotations": {"llm.local/compose-service": raw_service},
        },
        "spec": {
            "selector": {"app": app},
            "ports": sport,
        },
    }


def secret(ns, app, env):
    return {
        "apiVersion": "v1",
        "kind": "Secret",
        "metadata": {"name": app + "-env", "namespace": ns},
        "type": "Opaque",
        "stringData": env,
    }


def main():
    names = [n for n in run(["docker", "ps", "--format", "{{.Names}}"]).split() if n and not n.startswith(SKIP_PREFIXES)]
    by_ns = {}
    rewrites = []
    summary = []
    for name in names:
        data = inspect(name)
        ns = namespace_for(data, name)
        raw, app = service_name(data, name)
        ports = collect_ports(data)
        if not ports and raw == "sandbox":
            ports.append({"containerPort": 8194, "protocol": "TCP"})
        elif not ports:
            for port in listen_ports(name):
                if 1 < port < 10000:
                    ports.append({"containerPort": port, "protocol": "TCP"})
        env = env_map(data, name)
        by_ns.setdefault(ns, {"secrets": [], "services": [], "workloads": []})
        by_ns[ns]["secrets"].append(secret(ns, app, env))
        svc = service(ns, app, raw, ports)
        if svc:
            by_ns[ns]["services"].append(svc)
        by_ns[ns]["workloads"].append(workload(ns, app, raw, data, name, ports))
        if raw != app:
            rewrites.append("%s.%s.svc.cluster.local %s.%s.svc.cluster.local" % (raw, ns, app, ns))
        hosts = ",".join(hostnames(env))
        published = ",".join(str(p.get("hostPort")) for p in ports if p.get("hostPort"))
        summary.append("%s ns=%s service=%s ports=%s hosts=%s" % (name, ns, app, published or "-", hosts or "-"))

    if OUT.exists():
        for child in OUT.rglob("*"):
            if child.is_file():
                child.unlink()
    OUT.mkdir(parents=True, exist_ok=True)
    ns_docs = [{
        "apiVersion": "v1",
        "kind": "Namespace",
        "metadata": {"name": ns},
    } for ns in sorted(by_ns)]
    (OUT / "namespaces.json").write_text(json.dumps({"apiVersion": "v1", "kind": "List", "items": ns_docs}), encoding="utf-8")
    for ns, groups in by_ns.items():
        folder = OUT / ns
        folder.mkdir(parents=True, exist_ok=True)
        for kind in ("secrets", "services", "workloads"):
            items = groups[kind]
            if not items:
                continue
            path = folder / (kind + ".json")
            path.write_text(json.dumps({"apiVersion": "v1", "kind": "List", "items": items}), encoding="utf-8")
            if kind == "secrets":
                os.chmod(path, 0o600)
    (OUT / "dns-rewrites.txt").write_text("\n".join(rewrites) + ("\n" if rewrites else ""), encoding="utf-8")
    (OUT / "summary.txt").write_text("\n".join(summary) + "\n", encoding="utf-8")
    print("\n".join(summary))
    print("rewrites", len(rewrites))


if __name__ == "__main__":
    main()