#!/usr/bin/env python
# =============================================================================
# eBdesk SiCABE — seed contoh Data Quality NATIVE OpenMetadata (Jalur B)
# =============================================================================
# Mengapa skrip ini ada:
#   Test case native TIDAK bisa diandalkan lewat TestRunner SDK yang berjalan di
#   klien: jalur itu memakai koneksi TERSIMPAN yang kehilangan secret sehingga
#   selalu berakhir `Aborted` dengan `401 Access Denied` (bug upstream, masih ada
#   di OM 2.0.2 — kredensial sudah dipastikan valid).
#   Jalur yang bekerja: definisi test case dibuat lewat REST API, lalu eksekusi
#   DIPICU OpenMetadata lewat pipeline service client (plugin → Airflow). Koneksi
#   Trino dibawa EKSPLISIT di `sourceConfig.serviceConnections`, sehingga tidak
#   bergantung pada koneksi tersimpan.
#
# Idempoten: test case diperiksa berdasarkan nama; pipeline diperiksa berdasarkan
# nama. Menjalankan ulang tidak menggandakan apa pun.
#
# Dipanggil oleh: scripts/seed-native-dq.sh (dari host) — atau manual:
#   docker exec -e OM_SERVICE=trino_sicabe ... sicabe-airflow-webserver \
#     python /opt/airflow/scripts/seed-om-native-dq.py [--trigger]
#
# Exit code: 0 = semua langkah yang diminta berhasil.
# =============================================================================
import base64
import json
import os
import sys
import urllib.error
import urllib.request

OM_API_URL = os.getenv("OM_API_URL", "http://openmetadata-server:8585/api")
ADMIN_EMAIL = os.getenv("OPENMETADATA_ADMIN_EMAIL", "admin@open-metadata.org")
SERVICE = os.getenv("OM_DATABASE_SERVICE", "trino_sicabe")
CATALOG = os.getenv("OM_DATABASE_NAME", "iceberg")
TABLE_FQN = os.getenv("OM_DQ_SEED_TABLE", f"{SERVICE}.{CATALOG}.gold.orders_gold")
PIPELINE_NAME = os.getenv("OM_DQ_SEED_PIPELINE", "sicabe_dq_gold_native")
SCHEDULE = os.getenv("OM_DQ_SEED_SCHEDULE", "0 3 * * *")
TRINO_HOSTPORT = os.getenv("OM_TRINO_HOSTPORT", "trino-coordinator:8443")
TRINO_USER = os.getenv("DBT_TRINO_USER", "svc_om")
TRINO_CERT = os.getenv("TRINO_CA_BUNDLE", "/etc/ssl/trino/trino.crt")

# Contoh test native. `parameterValues` HARUS memakai nama parameter definisi
# (tableRowCountToBeBetween: minValue/maxValue) dan columnValuesToBeNotNull
# tidak punya parameter sama sekali — kolomnya ditentukan lewat entityLink.
TEST_CASES = [
    {
        "name": "native_row_count_between",
        "testDefinition": "tableRowCountToBeBetween",
        "entityLink": f"<#E::table::{TABLE_FQN}>",
        "parameterValues": [
            {"name": "minValue", "value": "1"},
            {"name": "maxValue", "value": "1000000"},
        ],
    },
    {
        "name": "native_komoditas_not_null",
        "testDefinition": "columnValuesToBeNotNull",
        "entityLink": f"<#E::table::{TABLE_FQN}::columns::komoditas>",
    },
]


def _request(method, path, payload=None, token=None):
    """Panggil REST API OpenMetadata; kembalikan (status, body terurai)."""
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    data = json.dumps(payload).encode() if payload is not None else None
    req = urllib.request.Request(f"{OM_API_URL}{path}", data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            raw = resp.read()
            return resp.status, (json.loads(raw) if raw else {})
    except urllib.error.HTTPError as exc:
        raw = exc.read()[:300]
        try:
            return exc.code, json.loads(raw)
        except Exception:  # noqa: BLE001
            return exc.code, raw.decode(errors="replace")


def login():
    """Token admin: password dikirim base64 (wajib sejak OM 2.0.x).

    Dua kandidat: OPENMETADATA_ADMIN_PASSWORD, lalu 'admin' — default yang dipakai
    OM 2.0.x sampai bootstrap step 7 menyetel hash-nya dari .env. MAKSIMAL DUA
    percobaan dan tidak pernah diulang dalam loop: akun admin OM bisa terkunci.
    """
    last = "tidak ada kandidat password"
    for password in (os.getenv("OPENMETADATA_ADMIN_PASSWORD"), "admin"):
        if not password:
            continue
        body = {"email": ADMIN_EMAIL,
                "password": base64.b64encode(password.encode()).decode()}
        status, res = _request("POST", "/v1/users/login", body)
        if status == 200 and res.get("accessToken"):
            return res["accessToken"]
        last = f"HTTP {status}: {str(res)[:120]}"
    raise RuntimeError(f"login gagal (password .env dan 'admin'): {last}")


def ensure_test_cases(token):
    """Buat test case native yang belum ada. Kembalikan jumlah yang dibuat."""
    status, page = _request("GET", "/v1/dataQuality/testCases?limit=200", token=token)
    if status != 200:
        raise RuntimeError(f"daftar test case tidak terbaca (HTTP {status})")
    existing = {tc.get("name") for tc in page.get("data", [])}

    created = 0
    for spec in TEST_CASES:
        if spec["name"] in existing:
            print(f"  [SKIP] test case sudah ada: {spec['name']}")
            continue
        status, res = _request("POST", "/v1/dataQuality/testCases", spec, token=token)
        if status in (200, 201):
            created += 1
            print(f"  [OK]   test case dibuat: {spec['name']}")
        else:
            raise RuntimeError(f"gagal membuat {spec['name']} (HTTP {status}): {str(res)[:200]}")
    return created


def _pipeline_payload(suite_id):
    """Payload pipeline TestSuite dengan koneksi Trino EKSPLISIT (anti-401)."""
    return {
        "name": PIPELINE_NAME,
        "pipelineType": "TestSuite",
        "service": {"id": suite_id, "type": "testSuite"},
        "sourceConfig": {"config": {
            "type": "TestSuite",
            "entityFullyQualifiedName": TABLE_FQN,
            "serviceConnections": [{
                "serviceName": SERVICE,
                "serviceConnection": {"config": {
                    "type": "Trino", "scheme": "trino", "hostPort": TRINO_HOSTPORT,
                    "username": TRINO_USER,
                    "authType": {"password": os.getenv("TRINO_SVC_OM_PASSWORD", "")},
                    "catalog": CATALOG, "verify": TRINO_CERT,
                }},
            }],
        }},
        "airflowConfig": {"scheduleInterval": SCHEDULE},
    }


def ensure_pipeline(token):
    """Buat pipeline TestSuite bila belum ada. Kembalikan (id, dibuat_baru)."""
    status, suites = _request("GET", "/v1/dataQuality/testSuites?limit=50", token=token)
    if status != 200:
        raise RuntimeError(f"daftar suite tidak terbaca (HTTP {status})")
    suite = next((s for s in suites.get("data", []) if TABLE_FQN.rsplit(".", 1)[0] in s.get("name", "")), None)
    if suite is None:
        raise RuntimeError(f"test suite untuk '{TABLE_FQN}' belum ada — jalankan ingestion dulu")

    status, page = _request("GET", "/v1/services/ingestionPipelines?limit=50", token=token)
    for pipe in (page.get("data") or []) if status == 200 else []:
        if pipe.get("name") == PIPELINE_NAME:
            print(f"  [SKIP] pipeline sudah ada: {PIPELINE_NAME} (id {pipe['id']})")
            return pipe["id"], False

    status, res = _request("POST", "/v1/services/ingestionPipelines",
                           _pipeline_payload(suite["id"]), token=token)
    if status not in (200, 201):
        raise RuntimeError(f"gagal membuat pipeline (HTTP {status}): {str(res)[:200]}")
    print(f"  [OK]   pipeline dibuat: {PIPELINE_NAME} (id {res['id']})")
    return res["id"], True


def deploy_and_trigger(token, pipeline_id):
    """Deploy lalu trigger pipeline; trigger diulang sekali bila DAG belum siap."""
    import time

    status, res = _request("POST", f"/v1/services/ingestionPipelines/deploy/{pipeline_id}", token=token)
    print(f"  [{'OK' if status == 200 else 'GAGAL'}] deploy  -> HTTP {status}")
    if status != 200:
        raise RuntimeError(f"deploy gagal: {str(res)[:200]}")

    for attempt in (1, 2):
        status, res = _request("POST", f"/v1/services/ingestionPipelines/trigger/{pipeline_id}", token=token)
        print(f"  [{'OK' if status == 200 else 'GAGAL'}] trigger -> HTTP {status} (percobaan {attempt})")
        if status == 200:
            return
        # DAG baru dibuat: dag-processor bisa belum memparsingnya.
        time.sleep(20)
    raise RuntimeError("trigger gagal setelah 2 percobaan")


def main():
    trigger = "--trigger" in sys.argv
    print(f"[seed-native-dq] tabel={TABLE_FQN} pipeline={PIPELINE_NAME}")
    token = login()
    print(f"  [OK]   login {ADMIN_EMAIL}")

    created = ensure_test_cases(token)
    pipeline_id, new_pipeline = ensure_pipeline(token)

    if trigger:
        deploy_and_trigger(token, pipeline_id)
    print(f"[seed-native-dq] selesai: {created} test case baru, pipeline "
          f"{'baru' if new_pipeline else 'sudah ada'}"
          f"{', deploy+trigger dijalankan' if trigger else ''}")
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as exc:  # noqa: BLE001
        print(f"ERROR: {exc}")
        sys.exit(1)