#!/usr/bin/env python
# =============================================================================
# eBdesk SiCABE — JALUR B: Data Quality NATIVE OpenMetadata (OM yang menjalankan)
# =============================================================================
# Bedanya dengan Jalur A (Great Expectations):
#   A: GE yang MENJALANKAN ekspektasi → hasilnya DIKIRIM ke OM
#      (test case ber-platform "GreatExpectations").
#   B: OPENMETADATA yang MENJALANKAN test (definisi bawaan OM:
#      tableRowCountToBeBetween, columnValuesToBeNotNull, columnValuesToBeBetween)
#      → hasilnya langsung tersimpan di OM (platform "OpenMetadata").
#
# Catatan implementasi (dua hal yang harus begini, sudah diverifikasi):
#   1. HARUS lewat `TestRunner.for_table`/`from_yaml` — `TestRunner(...)` mentah
#      tidak memuat entity tabel sehingga run() gagal "Table entity not provided".
#   2. Koneksi data HARUS dibawa eksplisit lewat workflow (serviceConnection +
#      verify=<sertifikat Trino>). Bila mengandalkan koneksi yang TERSIMPAN di OM,
#      eksekusi gagal 401 'Invalid credentials' pada OM 1.13.1. `metadata ingest`
#      CLI juga bukan jalur eksekusi TestSuite (registry source-nya terpisah).
#
# Jalankan DI DALAM container Airflow:
#   docker exec -e OPENMETADATA_ADMIN_PASSWORD=... \
#     sicabe-airflow-webserver python /opt/airflow/scripts/run-om-native-dq.py
#
# Exit code: 0 = semua test lulus; 1 = ada test gagal/aborted.
# =============================================================================
import base64
import json
import os
import sys
import urllib.request

from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    OpenMetadataConnection,
)
from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
    OpenMetadataJWTClientConfig,
)
from metadata.ingestion.ometa.ometa_api import OpenMetadata
from metadata.sdk.data_quality.runner import TestRunner
from metadata.sdk.data_quality.tests import (
    ColumnValuesToBeBetween,
    ColumnValuesToBeNotNull,
    TableRowCountToBeBetween,
)

OM_API_URL = os.getenv("OM_API_URL", "http://openmetadata-server:8585/api")
SERVICE = os.getenv("OM_DATABASE_SERVICE", "trino_sicabe")
CATALOG = os.getenv("OM_DATABASE_NAME", "iceberg")
TRINO_HOSTPORT = os.getenv("OM_TRINO_HOSTPORT", "trino-coordinator:8443")
TRINO_CERT = os.getenv("TRINO_CA_BUNDLE", "/etc/ssl/trino/trino.crt")

# Workflow TestSuite: membawa koneksi Trino eksplisit (username + password +
# verify=sertifikat) supaya eksekusi terautentikasi dan TLS tetap DIVERIFIKASI.
WORKFLOW_TEMPLATE = """
source:
  type: TestSuite
  serviceName: {service}
  sourceConfig:
    config:
      type: TestSuite
      entityFullyQualifiedName: {table}
      # WAJIB 'serviceConnections' (plural, list) — DQ source membaca field ini via
      # _load_yaml_service_connections(). Memakai 'serviceConnection' (tunggal)
      # diabaikan, sehingga eksekusi jatuh ke koneksi TERSIMPAN di OM yang
      # kehilangan secret → HTTP 401 'Invalid credentials' (lihat TODO upstream
      # open-metadata#21259). Dengan bentuk ini koneksi dari YAML yang dipakai.
      serviceConnections:
        - serviceName: {service}
          serviceConnection:
            config:
              type: Trino
              scheme: trino
              hostPort: {hostport}
              username: svc_om
              authType:
                password: "{password}"
              catalog: {catalog}
              verify: {cert}
sink:
  type: metadata-rest
  config: {{}}
workflowConfig:
  loggerLevel: INFO
  openMetadataServerConfig:
    hostPort: {om_api}
    authProvider: openmetadata
    securityConfig:
      jwtToken: "{token}"
"""

# Test NATIVE (dijalankan OM) — definisi bawaan OM, bukan ekspektasi GE.
TABLES = {
    f"{SERVICE}.{CATALOG}.bronze.orders": [
        TableRowCountToBeBetween(min_count=1, max_count=1000000),
        ColumnValuesToBeNotNull(column="id"),
        ColumnValuesToBeNotNull(column="komoditas"),
        ColumnValuesToBeBetween(column="harga_per_kg", min_value=0, max_value=1000000),
    ],
    f"{SERVICE}.{CATALOG}.silver.orders_silver": [
        TableRowCountToBeBetween(min_count=1, max_count=1000000),
        ColumnValuesToBeNotNull(column="id"),
        ColumnValuesToBeBetween(column="total_value", min_value=0, max_value=1000000000000),
    ],
    f"{SERVICE}.{CATALOG}.gold.orders_gold": [
        TableRowCountToBeBetween(min_count=1, max_count=1000000),
        ColumnValuesToBeNotNull(column="komoditas"),
        ColumnValuesToBeBetween(column="avg_harga_per_kg", min_value=0, max_value=1000000),
        ColumnValuesToBeBetween(column="total_nilai", min_value=0, max_value=1000000000000),
    ],
}


def _token():
    """Token OM: dari OM_JWT_TOKEN, atau login pakai password admin.

    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 diulang dalam loop: akun admin OM bisa terkunci.

    Catatan: 401 pada kandidat pertama adalah hal NORMAL (mis. step 7 belum jalan),
    jadi urlopen-nya dibungkus try/except — tanpa itu HTTPError 401 menggagalkan
    seluruh skrip.
    """
    if os.getenv("OM_JWT_TOKEN"):
        return os.environ["OM_JWT_TOKEN"]
    email = os.getenv("OPENMETADATA_ADMIN_EMAIL", "admin@open-metadata.org")
    for password in (os.getenv("OPENMETADATA_ADMIN_PASSWORD"), "admin"):
        if not password:
            continue
        body = json.dumps({
            "email": email,
            "password": base64.b64encode(password.encode()).decode(),
        }).encode()
        req = urllib.request.Request(f"{OM_API_URL}/v1/users/login", data=body,
                                    headers={"Content-Type": "application/json"})
        try:
            tok = json.load(urllib.request.urlopen(req, timeout=20)).get("accessToken")
        except Exception:  # noqa: BLE001 — kandidat pertama gagal = hal biasa
            tok = None
        if tok:
            return tok
    return None


def main():
    token = _token()
    trino_password = os.getenv("TRINO_SVC_OM_PASSWORD")
    if not token:
        print("ERROR: token OpenMetadata tidak tersedia (OM_JWT_TOKEN / "
              "OPENMETADATA_ADMIN_PASSWORD)")
        return 1
    if not trino_password:
        print("ERROR: TRINO_SVC_OM_PASSWORD tidak diset — koneksi Trino tidak bisa dibuat")
        return 1

    client = OpenMetadata(
        OpenMetadataConnection(
            hostPort=OM_API_URL,
            authProvider="openmetadata",
            securityConfig=OpenMetadataJWTClientConfig(jwtToken=token),
        )
    )

    failed = 0
    total = 0
    for table_fqn, tests in TABLES.items():
        print(f"[DQ-native] tabel {table_fqn} — {len(tests)} test")
        workflow = WORKFLOW_TEMPLATE.format(
            service=SERVICE, hostport=TRINO_HOSTPORT, password=trino_password,
            catalog=CATALOG, cert=TRINO_CERT, table=table_fqn, om_api=OM_API_URL,
            token=token,
        )
        # use_connection_from_yaml=True → pakai koneksi dari workflow (bukan yang
        # tersimpan di OM), sehingga autentikasi Trino bekerja.
        runner = TestRunner.from_yaml(
            yaml_string=workflow, use_connection_from_yaml=True, client=client
        )
        for test in tests:
            runner.add_test(test)

        for result in runner.run():
            total += 1
            status = str(getattr(result.testCaseResult, "testCaseStatus", None))
            name = getattr(result, "name", None) or getattr(result, "fqn", "test")
            print(f"    - {name}: {status.split('.')[-1]}")
            if status.split(".")[-1] != "Success":
                failed += 1

    print(f"[DQ-native] ringkasan: {total - failed}/{total} test lulus")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())