"""Add dataset lineage for the Airflow DAGs into OpenMetadata.

- Generator DAGs produce a source table  -> pipeline:table edge
- ETL DAGs transform source -> target     -> table:table edge carrying the pipeline
"""
import json
import os
import urllib.error
import urllib.parse
import urllib.request

OM = "http://openmetadata-server:8585/api"
TOKEN = os.environ["OM_TOKEN"]
HDRS = {"Authorization": "Bearer " + TOKEN}


def req(method, path, body=None):
    data = json.dumps(body).encode() if body is not None else None
    h = dict(HDRS)
    if data is not None:
        h["Content-Type"] = "application/json"
    last = "?"
    for _ in range(3):
        r = urllib.request.Request(OM + path, data=data, headers=h, method=method)
        try:
            with urllib.request.urlopen(r, timeout=60) as resp:
                return resp.status, json.loads(resp.read().decode() or "{}")
        except urllib.error.HTTPError as e:
            return e.code, e.read().decode()
        except Exception as e:  # noqa: BLE001 - timeouts / transient
            last = repr(e)
    return 0, last


def get_id(kind, fqn):
    base = "/v1/pipelines/name/" if kind == "pipeline" else "/v1/tables/name/"
    st, d = req("GET", base + urllib.parse.quote(fqn, safe=""))
    if st == 200:
        return d["id"]
    print("  ! cannot resolve", kind, fqn, "->", st)
    return None


def edge(from_kind, from_fqn, to_fqn, pipeline_fqn=None):
    fid = get_id(from_kind, from_fqn)
    tid = get_id("table", to_fqn)
    if not fid or not tid:
        return
    body = {
        "edge": {
            "fromEntity": {"id": fid, "type": from_kind},
            "toEntity": {"id": tid, "type": "table"},
        }
    }
    if pipeline_fqn:
        pid = get_id("pipeline", pipeline_fqn)
        if pid:
            body["edge"]["lineageDetails"] = {
                "pipeline": {"id": pid, "type": "pipeline"}
            }
    st, res = req("PUT", "/v1/lineage", body)
    label = (pipeline_fqn or from_fqn).split(".")[-1]
    print("  %-22s %s -> %s" % (label, from_fqn.split(".")[-1], to_fqn.split(".")[-1]),
          "OK" if st in (200, 201) else ("ERR %s %s" % (st, res)), flush=True)


PG = "atc_postgres.postgres.public.sales_orders"
MY = "atc_mysql.default.hr.employee_events"
MO = "atc_mongodb.supplychain.supplychain.events"

print("generators:")
edge("pipeline", "atc_airflow.gen_postgres", PG)
edge("pipeline", "atc_airflow.gen_mysql", MY)
edge("pipeline", "atc_airflow.gen_mongodb", MO)
edge("pipeline", "atc_airflow.generate_data_all_databases", PG)
edge("pipeline", "atc_airflow.generate_data_all_databases", MY)
edge("pipeline", "atc_airflow.generate_data_all_databases", MO)

print("ETL (table -> table via pipeline):")
edge("table", PG, "atc_trino.iceberg.curated_masked.sales_orders_masked",
     "atc_airflow.mask_to_curated")
edge("table", MY, "atc_trino.iceberg.curated_masked.employee_events_masked",
     "atc_airflow.mask_to_curated")
edge("pipeline", "atc_airflow.hadoop_to_trino",
     "atc_trino.iceberg.hadoop.historical_sales_hdfs")