#!/usr/bin/env python3
import base64
import json
import os
import subprocess
import time

KUBECTL = ["kubectl", "-n", "vault"]
INIT_PATH = "/usr/cluster/vault/init.json"
NAMESPACES = ["llm", "dify", "langfuse", "observability"]


def run(args, input_bytes=None, check=True):
    result = subprocess.run(args, input=input_bytes, capture_output=True)
    if check and result.returncode != 0:
        err = result.stderr.decode("utf-8", "replace")[-1500:]
        raise SystemExit(err)
    return result


def vault_status():
    result = run(KUBECTL + ["exec", "vault-0", "--", "vault", "status", "-format=json"], check=False)
    text = result.stdout.decode("utf-8", "replace").strip()
    if not text:
        raise SystemExit(result.stderr.decode("utf-8", "replace")[-1500:])
    return json.loads(text)


def wait_running():
    for _ in range(60):
        result = run(
            ["kubectl", "-n", "vault", "get", "pod", "vault-0", "-o", "jsonpath={.status.phase}"],
            check=False,
        )
        phase = result.stdout.decode().strip()
        print("phase", phase or "missing", flush=True)
        if phase == "Running":
            status = vault_status()
            print("initialized", status.get("initialized"), "sealed", status.get("sealed"), flush=True)
            return status
        time.sleep(5)
    run(["kubectl", "-n", "vault", "describe", "pod", "vault-0"])
    raise SystemExit("vault pod did not start")


def ensure_init(status):
    if status.get("initialized"):
        if not os.path.exists(INIT_PATH):
            raise SystemExit("vault is initialized but the local init file is missing")
        return
    result = run(
        KUBECTL
        + [
            "exec",
            "vault-0",
            "--",
            "vault",
            "operator",
            "init",
            "-key-shares=1",
            "-key-threshold=1",
            "-format=json",
        ]
    )
    os.makedirs("/usr/cluster/vault", exist_ok=True)
    with open(INIT_PATH, "w", encoding="utf-8") as handle:
        handle.write(result.stdout.decode())
    os.chmod(INIT_PATH, 0o600)
    print("init file written", flush=True)


def unseal_if_needed():
    status = vault_status()
    if not status.get("sealed"):
        print("already unsealed", flush=True)
        return
    init = json.load(open(INIT_PATH, encoding="utf-8"))
    key = init["unseal_keys_b64"][0]
    result = run(KUBECTL + ["exec", "vault-0", "--", "vault", "operator", "unseal", key])
    if b'"sealed": false' not in result.stdout and b"Sealed          false" not in result.stdout:
        # status is the source of truth; unseal output format varies
        status = vault_status()
        if status.get("sealed"):
            raise SystemExit("vault remained sealed")
    print("unsealed", flush=True)


def vault_cmd(*args):
    script = "export VAULT_ADDR=http://127.0.0.1:8200; export VAULT_TOKEN=$(cat /tmp/vt); vault " + " ".join(
        json.dumps(arg) for arg in args
    )
    return run(KUBECTL + ["exec", "vault-0", "--", "sh", "-c", script])


def install_token():
    init = json.load(open(INIT_PATH, encoding="utf-8"))
    run(
        KUBECTL + ["exec", "-i", "vault-0", "--", "sh", "-c", "cat > /tmp/vt && chmod 600 /tmp/vt"],
        input_bytes=init["root_token"].encode(),
    )


def enable_kv():
    listed = vault_cmd("secrets", "list", "-format=json")
    mounts = json.loads(listed.stdout.decode())
    if "secret/" not in mounts:
        vault_cmd("secrets", "enable", "-path=secret", "kv-v2")
        print("kv v2 enabled", flush=True)
    else:
        print("kv v2 present", flush=True)


def import_secrets():
    copied = 0
    for ns in NAMESPACES:
        raw = run(["kubectl", "-n", ns, "get", "secret", "-o", "json"]).stdout
        for item in json.loads(raw)["items"]:
            name = item["metadata"]["name"]
            if item.get("type") != "Opaque":
                continue
            if name.startswith("sh.helm.release"):
                continue
            data = item.get("data") or {}
            if not data:
                print(f"SKIP {ns}/{name} empty", flush=True)
                continue
            payload = {}
            for key, value in data.items():
                blob = base64.b64decode(value)
                try:
                    payload[key] = blob.decode("utf-8")
                except UnicodeDecodeError:
                    payload[key] = blob.decode("latin-1")
                    print(f"LATIN1 {ns}/{name} {key}", flush=True)
            local = f"/tmp/vault-payload-{ns}-{name}.json"
            with open(local, "w", encoding="utf-8") as handle:
                json.dump(payload, handle, ensure_ascii=False)
            os.chmod(local, 0o600)
            run(KUBECTL + ["cp", local, f"vault-0:/tmp/payload.json"])
            os.remove(local)
            vault_cmd("kv", "put", "-mount=secret", f"apps/{ns}/{name}", "@/tmp/payload.json")
            run(KUBECTL + ["exec", "vault-0", "--", "rm", "-f", "/tmp/payload.json"])
            check = vault_cmd("kv", "get", "-format=json", "-mount=secret", f"apps/{ns}/{name}")
            stored = json.loads(check.stdout.decode())["data"]["data"]
            if set(stored) != set(payload):
                raise SystemExit(f"key set mismatch {ns}/{name}")
            print(f"PUT {ns}/{name} keys={len(payload)}", flush=True)
            copied += 1
    print(f"COPIED {copied}", flush=True)


def cleanup_token():
    run(KUBECTL + ["exec", "vault-0", "--", "rm", "-f", "/tmp/vt"], check=False)


def main():
    status = wait_running()
    ensure_init(status)
    unseal_if_needed()
    install_token()
    try:
        enable_kv()
        import_secrets()
    finally:
        cleanup_token()


if __name__ == "__main__":
    main()