#!/usr/bin/env python3
"""初始化 Vault、解封，并把四个命名空间里已有的 Opaque Secret 抄进 KV v2。

现网已经完成。再次执行时如果 init.json 已存在，只做解封和补齐，不会重新 init。
脚本不打印口令、令牌和 Secret 的值。
"""
import base64
import json
import os
import subprocess

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:
        raise SystemExit(result.stderr.decode("utf-8", "replace")[-800:])
    return result


def status():
    result = run(KUBECTL + ["exec", "vault-0", "--", "vault", "status", "-format=json"], check=False)
    return json.loads(result.stdout.decode())


def ensure_init():
    current = status()
    if current.get("initialized"):
        if not os.path.exists(INIT_PATH):
            raise SystemExit("Vault 已初始化，但本机没有 " + INIT_PATH)
        print("already initialized")
        return
    result = run(
        KUBECTL
        + ["exec", "vault-0", "--", "vault", "operator", "init", "-key-shares=1", "-key-threshold=1", "-format=json"]
    )
    with open(INIT_PATH, "w", encoding="utf-8") as handle:
        handle.write(result.stdout.decode())
    os.chmod(INIT_PATH, 0o600)
    print("init file written")


def unseal():
    if not status().get("sealed"):
        print("already unsealed")
        return
    key = json.load(open(INIT_PATH, encoding="utf-8"))["unseal_keys_b64"][0]
    run(KUBECTL + ["exec", "vault-0", "--", "vault", "operator", "unseal", key])
    if status().get("sealed"):
        raise SystemExit("unseal failed")
    print("unsealed")


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 with_token():
    token = json.load(open(INIT_PATH, encoding="utf-8"))["root_token"]
    run(
        KUBECTL + ["exec", "-i", "vault-0", "--", "sh", "-c", "cat > /tmp/vt && chmod 600 /tmp/vt"],
        input_bytes=token.encode(),
    )


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


def import_secrets():
    copied = 0
    for ns in NAMESPACES:
        raw = json.loads(run(["kubectl", "-n", ns, "get", "secret", "-o", "json"]).stdout)
        for item in raw["items"]:
            name = item["metadata"]["name"]
            if item.get("type") != "Opaque" or name.startswith("sh.helm.release"):
                continue
            data = item.get("data") or {}
            if not data:
                print(f"SKIP {ns}/{name} empty")
                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")
            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, "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"])
            print(f"PUT {ns}/{name} keys={len(payload)}")
            copied += 1
    print(f"COPIED {copied}")


def main():
    ensure_init()
    unseal()
    with_token()
    try:
        enable_kv()
        import_secrets()
    finally:
        run(KUBECTL + ["exec", "vault-0", "--", "rm", "-f", "/tmp/vt"], check=False)


if __name__ == "__main__":
    main()