# Databricks notebook source
# MAGIC %md
# MAGIC # Bronze ingestion: Redshift -> Unity Catalog (Trakt.tv example)
# MAGIC Assumes Trakt personal data (watch history, ratings, watchlist) was already
# MAGIC exported from the Trakt API into Redshift schema `trakt_raw`.
# MAGIC Pattern: incremental read -> audit columns -> idempotent MERGE -> watermark -> DQ gate.

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

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

# ---- Parameters (set per environment by the job / bundle) ----
dbutils.widgets.text("catalog", "dev")
dbutils.widgets.text("source_schema", "trakt_raw")
dbutils.widgets.dropdown("mode", "incremental", ["incremental", "full"])

CATALOG = dbutils.widgets.get("catalog")
SRC_SCHEMA = dbutils.widgets.get("source_schema")
MODE = dbutils.widgets.get("mode")

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

BRONZE = f"{CATALOG}.bronze"
OPS = f"{CATALOG}.ops"
WM = f"{OPS}.ingestion_watermark"

# ---- Source table config: business key + incremental cursor column ----
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 ## Secrets and connection (no credentials in code)

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

JDBC_URL = dbutils.secrets.get("redshift", "jdbc_url")      # jdbc:redshift://host:5439/db
RS_USER = dbutils.secrets.get("redshift", "user")
RS_PASSWORD = dbutils.secrets.get("redshift", "password")
IAM_ROLE = dbutils.secrets.get("redshift", "unload_iam_role")
TEMPDIR = dbutils.secrets.get("redshift", "tempdir")         # s3://bucket/redshift-tmp/ (UC external location)


def read_redshift(query: str):
    """Redshift connector unloads via S3 (faster than plain JDBC for large reads)."""
    return (
        spark.read.format("redshift")
        .option("url", JDBC_URL)
        .option("user", RS_USER)
        .option("password", RS_PASSWORD)
        .option("query", query)
        .option("tempdir", TEMPDIR)
        .option("aws_iam_role", IAM_ROLE)
        .load()
    )

# COMMAND ----------
# MAGIC %md ## Unity Catalog objects (Spark SQL DDL)

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

spark.sql(f"CREATE SCHEMA IF NOT EXISTS {BRONZE} COMMENT 'Raw landed data, minimal transformation'")
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {OPS} COMMENT 'Pipeline control tables'")

spark.sql(f"""
    CREATE TABLE IF NOT EXISTS {WM} (
        source_table  STRING,
        last_cursor   TIMESTAMP,
        last_batch_id STRING,
        updated_at    TIMESTAMP
    ) USING DELTA
""")

# Allow new source columns to flow into bronze instead of failing the job
spark.conf.set("spark.databricks.delta.schema.autoMerge.enabled", "true")

# COMMAND ----------
# MAGIC %md ## Helpers: watermark, audit columns, idempotent write, DQ

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


def get_watermark(name: str):
    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: str, cursor_value):
    if cursor_value is None:
        return
    ts = cursor_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_columns(df, name: str):
    return (
        df.withColumn("_source_table", F.lit(f"{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: str, keys: list):
    """Idempotent: re-running a batch (or overlapping windows) never duplicates rows."""
    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)  # null-safe equality
    spark.sql(f"""
        MERGE INTO {target} t
        USING incoming s
        ON {on_clause}
        WHEN NOT MATCHED THEN INSERT *
    """)


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

# COMMAND ----------
# MAGIC %md ## Run: one pass per source table

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

results = {}

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

    wm = get_watermark(name)
    query = f"SELECT * FROM {SRC_SCHEMA}.{name}"
    if wm is not None:
        # >= (not >) so late-arriving rows at the boundary aren't missed;
        # the MERGE de-duplicates the overlap
        query += f" WHERE {cursor} >= '{wm.strftime('%Y-%m-%d %H:%M:%S')}'"

    df = read_redshift(query)
    df = add_audit_columns(df, name).cache()

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

    merge_insert_only(df, target, keys)
    set_watermark(name, new_wm)   # only after a successful write

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

# COMMAND ----------
# MAGIC %md ## Governance (Unity Catalog grants) and post-load checks

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

spark.sql(f"GRANT USE SCHEMA, SELECT ON SCHEMA {BRONZE} TO `data-engineers`")

spark.sql(f"""
    SELECT _source_table, _batch_id, COUNT(*) AS rows_loaded
    FROM {BRONZE}.trakt_watch_history
    GROUP BY _source_table, _batch_id
    ORDER BY MAX(_ingested_at) DESC
    LIMIT 5
""").show(truncate=False)

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

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