# Databricks notebook source
# MAGIC %md
# MAGIC # Bronze ingestion: Redshift foreign catalog -> Delta bronze (Trakt example)
# MAGIC The Redshift login lives in a Unity Catalog **connection**; the **foreign catalog**
# MAGIC built on it is what shows up in the Catalog Explorer UI. No secrets, JDBC URL or S3
# MAGIC tempdir in this notebook. Foreign tables are read-only, so we land into a native Delta catalog.

# COMMAND ----------

import json
import uuid
from datetime import datetime, timezone
from pyspark.sql import functions as F

dbutils.widgets.text("source_catalog", "redshift_bronze")   # foreign catalog (Redshift connection)
dbutils.widgets.text("source_schema", "trakt_raw")
dbutils.widgets.text("target_catalog", "dev")               # native Delta catalog
dbutils.widgets.dropdown("mode", "incremental", ["incremental", "full"])

SRC_CAT = dbutils.widgets.get("source_catalog")
SRC_SCHEMA = dbutils.widgets.get("source_schema")
TGT_CAT = dbutils.widgets.get("target_catalog")
MODE = dbutils.widgets.get("mode")

BATCH_ID = str(uuid.uuid4())
RUN_TS = datetime.now(timezone.utc)

BRONZE = f"{TGT_CAT}.bronze"
WM = f"{TGT_CAT}.ops.ingestion_watermark"

TABLES = {
    "watch_history": {"key": ["id"], "cursor": "watched_at"},
    "ratings":       {"key": ["trakt_id", "rated_at"], "cursor": "rated_at"},
    "watchlist":     {"key": ["trakt_id", "listed_at"], "cursor": "listed_at"},
}

# COMMAND ----------
# MAGIC %md ## Sanity check the federation (what you see in the UI, via SQL)

# COMMAND ----------

spark.sql(f"SHOW TABLES IN {SRC_CAT}.{SRC_SCHEMA}").show(truncate=False)
spark.sql(f"DESCRIBE EXTENDED {SRC_CAT}.{SRC_SCHEMA}.watch_history").show(50, truncate=False)

# COMMAND ----------
# MAGIC %md ## Target objects

# COMMAND ----------

spark.sql(f"CREATE SCHEMA IF NOT EXISTS {BRONZE}")
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {TGT_CAT}.ops")
spark.sql(f"""
    CREATE TABLE IF NOT EXISTS {WM} (
        source_table STRING, last_cursor TIMESTAMP,
        last_batch_id STRING, updated_at TIMESTAMP
    ) USING DELTA
""")
spark.conf.set("spark.databricks.delta.schema.autoMerge.enabled", "true")

# COMMAND ----------
# MAGIC %md ## Helpers

# COMMAND ----------


def get_watermark(name):
    if MODE == "full":
        return None
    row = spark.sql(f"SELECT last_cursor FROM {WM} WHERE source_table = '{name}'").first()
    return row["last_cursor"] if row else None


def set_watermark(name, value):
    if value is None:
        return
    ts = value.strftime("%Y-%m-%d %H:%M:%S.%f")
    spark.sql(f"""
        MERGE INTO {WM} t
        USING (SELECT '{name}' AS source_table, TIMESTAMP'{ts}' AS last_cursor,
                      '{BATCH_ID}' AS last_batch_id, current_timestamp() AS updated_at) s
        ON t.source_table = s.source_table
        WHEN MATCHED THEN UPDATE SET *
        WHEN NOT MATCHED THEN INSERT *
    """)


def add_audit(df, name):
    return (df.withColumn("_source_table", F.lit(f"{SRC_CAT}.{SRC_SCHEMA}.{name}"))
              .withColumn("_batch_id", F.lit(BATCH_ID))
              .withColumn("_ingested_at", F.lit(RUN_TS))
              .withColumn("_ingest_date", F.to_date(F.lit(RUN_TS))))


def merge_insert_only(df, target, keys):
    if not spark.catalog.tableExists(target):
        df.write.format("delta").option("delta.enableChangeDataFeed", "true").saveAsTable(target)
        return
    df.createOrReplaceTempView("incoming")
    on_clause = " AND ".join(f"t.{k} <=> s.{k}" for k in keys)
    spark.sql(f"""
        MERGE INTO {target} t USING incoming s ON {on_clause}
        WHEN NOT MATCHED THEN INSERT *
    """)


def dq_gate(df, name, keys):
    total = df.count()
    null_keys = df.filter(" OR ".join(f"{k} IS NULL" for k in keys)).count()
    if null_keys:
        raise ValueError(f"[{name}] {null_keys} NULL business keys; failing batch {BATCH_ID}")
    return {"rows_read": total, "null_keys": null_keys}

# COMMAND ----------
# MAGIC %md ## Run

# COMMAND ----------

results = {}

for name, cfg in TABLES.items():
    keys, cursor = cfg["key"], cfg["cursor"]
    source = f"{SRC_CAT}.{SRC_SCHEMA}.{name}"
    target = f"{BRONZE}.trakt_{name}"

    wm = get_watermark(name)
    src_df = spark.table(source)                      # federated read, no connector options
    if wm is not None:
        src_df = src_df.filter(F.col(cursor) >= F.lit(wm))   # should push down to Redshift

    df = add_audit(src_df, name).cache()
    stats = dq_gate(df, name, keys)
    new_wm = df.agg(F.max(cursor)).first()[0] if stats["rows_read"] else None

    merge_insert_only(df, target, keys)
    set_watermark(name, new_wm)

    stats.update({"target": target, "watermark_after": str(new_wm)})
    results[name] = stats
    df.unpersist()

# COMMAND ----------
# MAGIC %md ## Verify pushdown (look for the filter inside the federated scan, not above it)

# COMMAND ----------

spark.table(f"{SRC_CAT}.{SRC_SCHEMA}.watch_history") \
     .filter("watched_at >= '2026-01-01'") \
     .select("id", "watched_at") \
     .explain("formatted")

# COMMAND ----------

dbutils.notebook.exit(json.dumps({"batch_id": BATCH_ID, "mode": MODE, "results": results}))