"""mask_to_curated — ETL with PII masking into the curated Iceberg layer.

Reads PII-bearing rows from the source databases via Trino (postgres_sales,
mysql_hr catalogs) and writes MASKED copies into iceberg.curated_masked.*.
Masking is expressed as Trino SQL (hash / partial-redact / generalize) so no PII
ever lands in the curated layer. Triggered by the Command Center (movement
'mask_to_curated') or autonomously by the ETL agent.
"""
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
import json
import os
import urllib.request

TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/")
TRINO_USER = os.getenv("TRINO_USER", "mo")
LIMIT = int(os.getenv("MASK_LIMIT", "20000"))

default_args = {"owner": "airflow", "retries": 0}


def trino(sql: str):
    """Run a Trino statement, following nextUri pages. Returns all data rows."""
    req = urllib.request.Request(
        f"{TRINO_URL}/v1/statement", data=sql.encode(),
        headers={"X-Trino-User": TRINO_USER, "Content-Type": "text/plain"},
    )
    rows = []
    with urllib.request.urlopen(req, timeout=30) as resp:
        d = json.loads(resp.read().decode())
    while True:
        rows += d.get("data") or []
        err = d.get("error")
        if err:
            raise Exception(f"Trino error: {err.get('message')}")
        nxt = d.get("nextUri")
        if not nxt:
            break
        with urllib.request.urlopen(nxt, timeout=30) as resp:
            d = json.loads(resp.read().decode())
    return rows


# Masked projections. PII columns are hashed / partially redacted / generalized.
SALES_SQL = f"""
CREATE TABLE iceberg.curated_masked.sales_orders_masked
WITH (format = 'PARQUET') AS
SELECT
  order_id, customer_id, product_id, region, sales_channel,
  CAST(order_ts AS timestamp(6)) AS order_ts, amount, currency, order_status,
  to_hex(md5(to_utf8(customer_name))) AS customer_name_hash,
  regexp_replace(customer_email, '(^.)[^@]*(@.*)$', '$1***$2') AS customer_email_masked,
  concat('***', substr(customer_phone, -4)) AS customer_phone_masked,
  concat(substr(billing_iban, 1, 6), '****') AS billing_iban_masked,
  regexp_replace(customer_ip, '\\.\\d+$', '.0') AS customer_ip_masked,
  element_at(split(shipping_address, ', '), 2) AS shipping_area
FROM postgres_sales.public.sales_orders
WHERE customer_email IS NOT NULL
LIMIT {LIMIT}
"""

EMP_SQL = f"""
CREATE TABLE iceberg.curated_masked.employee_events_masked
WITH (format = 'PARQUET') AS
SELECT
  event_id, employee_id, department, role_name, region, event_type, salary_change,
  CAST(event_ts AS timestamp(6)) AS event_ts,
  to_hex(md5(to_utf8(employee_name))) AS employee_name_hash,
  regexp_replace(employee_email, '(^.)[^@]*(@.*)$', '$1***$2') AS employee_email_masked,
  concat('***', substr(employee_phone, -4)) AS employee_phone_masked,
  concat('***', substr(national_id, -3)) AS national_id_masked,
  element_at(split(home_address, ', '), 2) AS home_area,
  CAST(year(CAST(date_of_birth AS date)) AS varchar) AS birth_year
FROM mysql_hr.hr.employee_events
WHERE employee_email IS NOT NULL
LIMIT {LIMIT}
"""


def _mask(**context):
    trino("CREATE SCHEMA IF NOT EXISTS iceberg.curated_masked")
    for name, sql in [("sales_orders_masked", SALES_SQL), ("employee_events_masked", EMP_SQL)]:
        print(f"Refreshing iceberg.curated_masked.{name} ...")
        trino(f"DROP TABLE IF EXISTS iceberg.curated_masked.{name}")
        trino(sql)
        cnt = trino(f"SELECT count(*) FROM iceberg.curated_masked.{name}")
        print(f"  {name}: {cnt[0][0] if cnt else '?'} masked rows")
    print("mask_to_curated complete")


dag = DAG(
    dag_id="mask_to_curated",
    default_args=default_args,
    description="Mask PII from sources into iceberg.curated_masked (Trino SQL masking)",
    schedule=None,
    start_date=datetime(2025, 1, 1),
    catchup=False,
    tags=["etl", "masking", "pii", "curated"],
)

PythonOperator(task_id="mask_to_curated", python_callable=_mask, dag=dag)