#!/usr/bin/env python3
"""OpenMetadata API helper: login, fetch ingestion-bot JWT, create admin users."""
import base64
import json
import sys
import urllib.request
import urllib.error

OM = "http://10.0.21.47:8585"


def req(method, path, token=None, body=None):
    url = OM + path
    data = json.dumps(body).encode() if body is not None else None
    r = urllib.request.Request(url, data=data, method=method)
    r.add_header("Content-Type", "application/json")
    if token:
        r.add_header("Authorization", "Bearer " + token)
    try:
        with urllib.request.urlopen(r, timeout=20) as resp:
            raw = resp.read().decode()
            return resp.status, (json.loads(raw) if raw else {})
    except urllib.error.HTTPError as e:
        raw = e.read().decode()
        try:
            return e.code, json.loads(raw)
        except Exception:
            return e.code, {"raw": raw}


def login(email, password):
    st, d = req("POST", "/api/v1/users/login", body={
        "email": email,
        "password": base64.b64encode(password.encode()).decode(),
    })
    return st, d


def main():
    action = sys.argv[1] if len(sys.argv) > 1 else "bootstrap"
    st, d = login("admin@open-metadata.org", "admin")
    if st != 200 or "accessToken" not in d:
        print("LOGIN_FAILED", st, json.dumps(d)[:300])
        sys.exit(1)
    token = d["accessToken"]
    print("ADMIN_LOGIN_OK")

    if action in ("bootstrap", "bottoken"):
        # Reuse the ingestion-bot's existing JWT (do NOT regenerate to avoid breaking internals).
        st, b = req("GET", "/api/v1/bots/name/ingestion-bot", token=token)
        uid = (b.get("botUser") or {}).get("id")
        jwt = None
        if uid:
            st2, u = req("GET", f"/api/v1/users/{uid}?include=all", token=token)
            am = u.get("authenticationMechanism") or {}
            jwt = (am.get("config") or {}).get("JWTToken")
        if not jwt and uid and "--regen" in sys.argv:
            st3, t = req("PUT", f"/api/v1/users/generateToken/{uid}",
                         token=token, body={"JWTTokenExpiry": "Unlimited"})
            jwt = t.get("JWTToken")
        if jwt:
            print("INGESTION_BOT_JWT=" + jwt)
        else:
            print("BOT_JWT_NOT_FOUND", json.dumps(b)[:200])

    if action == "lineage":
        pairs = [
            ("atc_postgres.postgres.public.sales_orders", "atc_trino.iceberg.curated_masked.sales_orders_masked"),
            ("atc_mysql.default.hr.employee_events", "atc_trino.iceberg.curated_masked.employee_events_masked"),
        ]
        for src_fqn, dst_fqn in pairs:
            s1, src = req("GET", f"/api/v1/tables/name/{src_fqn}", token=token)
            s2, dst = req("GET", f"/api/v1/tables/name/{dst_fqn}", token=token)
            sid, did = src.get("id"), dst.get("id")
            if not sid or not did:
                print(f"SKIP {src_fqn}->{dst_fqn} (missing id s={s1} d={s2})")
                continue
            st, r = req("PUT", "/api/v1/lineage", token=token, body={
                "edge": {
                    "fromEntity": {"id": sid, "type": "table"},
                    "toEntity": {"id": did, "type": "table"},
                    "lineageDetails": {"description": "PII masking via mask_to_curated (Trino SQL hash/redact/generalize)"},
                },
            })
            print(f"LINEAGE {src_fqn} -> {dst_fqn}: {st}")
        return

    if action == "tags":
        for fqn in sys.argv[2:]:
            st, t = req("GET", f"/api/v1/tables/name/{fqn}?fields=columns,tags", token=token)
            cols = t.get("columns", [])
            print(f"== {fqn} ({st}) cols={len(cols)} ==")
            for c in cols:
                tg = [x.get("tagFQN") for x in (c.get("tags") or [])]
                if tg:
                    print(f"   {c['name']}: {tg}")
        return

    if action == "bootstrap":
        for name in ("mo", "bart"):
            st, d2 = req("POST", "/api/v1/users/signup", body={
                "firstName": name.capitalize(), "lastName": "ATC",
                "email": f"{name}@open-metadata.org",
                "password": "Dell2026!",
            })
            print(f"USER_{name}: {st} {json.dumps(d2)[:140]}")


if __name__ == "__main__":
    main()