from delta.tables import DeltaTable
from pyspark.sql.functions import *

# Step 1: Prepare source with hash for change detection
source_df = spark.table("staging_customer").withColumn(
    "hash_value", md5(concat_ws("||", col("name"), col("email"), col("address")))
)

# Step 2: Get current target records
target = DeltaTable.forName(spark, "dim_customer")
target_df = target.toDF().filter("is_current = true")

# Step 3: Identify changes (new + changed records)
changes = source_df.alias("s").join(
    target_df.alias("t"),
    col("s.customer_id") == col("t.customer_id"),
    "left"
).filter(
    col("t.customer_id").isNull() |           # New records
    (col("s.hash_value") != col("t.hash_value"))  # Changed records
).select("s.*")

# Step 4: THE MERGE KEY TRICK
# For closing old records: use real customer_id as merge_key (matches target)
# For inserting new current records: set merge_key = NULL (never matches → goes to NOT MATCHED)

# Rows that need to CLOSE existing records (update is_current = false)
close_records = changes.filter(
    col("customer_id").isin(
        [row.customer_id for row in target_df.join(changes, "customer_id", "inner")
         .select("customer_id").distinct().collect()]
    )
).withColumn("merge_key", col("customer_id"))

# Rows that need to INSERT as new current records
insert_records = changes.withColumn("merge_key", lit(None).cast("string"))

# Combine: close_records will MATCH (update), insert_records will NOT MATCH (insert)
staged = close_records.unionByName(insert_records)

# Step 5: Execute MERGE
target.alias("t").merge(
    staged.alias("s"),
    "t.customer_id = s.merge_key AND t.is_current = true"
).whenMatchedUpdate(set={
    "is_current": lit(False),
    "effective_end": current_timestamp()
}).whenNotMatchedInsert(values={
    "customer_id": col("s.customer_id"),
    "name": col("s.name"),
    "email": col("s.email"),
    "address": col("s.address"),
    "hash_value": col("s.hash_value"),
    "effective_start": current_timestamp(),
    "effective_end": lit("9999-12-31 23:59:59").cast("timestamp"),
    "is_current": lit(True)
}).execute()

# Step 1: Close old records
target.alias("t").merge(
    changes.alias("s"),
    "t.customer_id = s.customer_id AND t.is_current = true"
).whenMatchedUpdate(set={
    "is_current": lit(False),
    "effective_end": current_timestamp()
}).execute()

# Step 2: Insert new current records
changes.withColumn("effective_start", current_timestamp()) \
    .withColumn("effective_end", lit("9999-12-31").cast("timestamp")) \
    .withColumn("is_current", lit(True)) \
    .write.format("delta").mode("append").saveAsTable("dim_customer")

import dlt
from pyspark.sql.functions import col, expr

# Bronze: raw CDC events
@dlt.table(comment="Raw customer CDC events")
def bronze_customer_cdc():
    return (
        spark.readStream.format("cloudFiles")
        .option("cloudFiles.format", "json")
        .option("cloudFiles.schemaLocation", "/schema/customer_cdc")
        .load("/landing/customer_cdc/")
    )

# Silver: SCD Type 2 using apply_changes
dlt.create_streaming_table("dim_customer_scd2")

dlt.apply_changes(
    target="dim_customer_scd2",
    source="bronze_customer_cdc",
    keys=["customer_id"],
    sequence_by=col("updated_at"),
    apply_as_deletes=expr("operation = 'DELETE'"),
    except_column_list=["operation", "_rescued_data"],
    stored_as_scd_type=2                 # <-- This is the magic
)
# DLT automatically manages __START_AT, __END_AT, __IS_CURRENT columns

# Optimized implementation
from delta.tables import DeltaTable
from pyspark.sql.functions import *

# Deduplicate source (keep latest)
w = Window.partitionBy("customer_id").orderBy(col("updated_at").desc())
source = spark.table("staging_customer") \
    .withColumn("rn", row_number().over(w)) \
    .filter(col("rn") == 1).drop("rn") \
    .withColumn("hash_value", md5(concat_ws("||", col("name"), col("email"), col("address"))))

# Get only changed records (compare hashes)
target_current = spark.table("dim_customer").filter("is_current = true") \
    .select("customer_id", col("hash_value").alias("target_hash"))

changes = source.join(target_current, "customer_id", "left") \
    .filter(col("target_hash").isNull() | (col("hash_value") != col("target_hash"))) \
    .drop("target_hash")

# Execute SCD Type 2 merge (using the merge-key trick)
# ... (same pattern as Q3)

1 minute") \
    .start()"># Bronze: Ingest raw CDC events from Kafka
raw_cdc = spark.readStream.format("kafka") \
    .option("kafka.bootstrap.servers", "broker:9092") \
    .option("subscribe", "oracle.schema.customers") \
    .load() \
    .select(
        from_json(col("value").cast("string"), cdc_schema).alias("data"),
        col("timestamp").alias("kafka_timestamp")
    ).select("data.*", "kafka_timestamp")

raw_cdc.writeStream.format("delta") \
    .outputMode("append") \
    .option("checkpointLocation", "/checkpoints/bronze_cdc") \
    .table("bronze_customer_cdc")

# Silver: Apply CDC using foreachBatch
def apply_cdc_changes(batch_df, batch_id):
    from delta.tables import DeltaTable

    # Deduplicate within batch (keep latest per key)
    w = Window.partitionBy("after.id").orderBy(col("ts_ms").desc())
    deduped = batch_df.withColumn("rn", row_number().over(w)) \
        .filter("rn = 1").drop("rn")

    target = DeltaTable.forName(spark, "silver_customers")

    target.alias("t").merge(
        deduped.alias("s"),
        "t.customer_id = s.after.id"
    ).whenMatchedDelete(
        condition="s.op = 'd'"                    # Delete
    ).whenMatchedUpdate(
        condition="s.op = 'u'",                   # Update
        set={
            "name": "s.after.name",
            "email": "s.after.email",
            "address": "s.after.address",
            "updated_at": "s.ts_ms"
        }
    ).whenNotMatchedInsert(
        condition="s.op IN ('c', 'r')",           # Create / Read (snapshot)
        values={
            "customer_id": "s.after.id",
            "name": "s.after.name",
            "email": "s.after.email",
            "address": "s.after.address",
            "created_at": "s.ts_ms",
            "updated_at": "s.ts_ms"
        }
    ).execute()

spark.readStream.table("bronze_customer_cdc") \
    .writeStream.foreachBatch(apply_cdc_changes) \
    .option("checkpointLocation", "/checkpoints/silver_customers") \
    .trigger(processingTime="1 minute") \
    .start()

# By version range
changes = spark.read.format("delta") \
    .option("readChangeFeed", "true") \
    .option("startingVersion", 5) \
    .option("endingVersion", 10) \
    .table("my_table")

# By timestamp range
changes = spark.read.format("delta") \
    .option("readChangeFeed", "true") \
    .option("startingTimestamp", "2025-01-01") \
    .option("endingTimestamp", "2025-01-31") \
    .table("my_table")

# Streaming (incremental)
changes = spark.readStream.format("delta") \
    .option("readChangeFeed", "true") \
    .option("startingVersion", 5) \
    .table("my_table")

# Design decisions:
# - Append-only (NEVER update or delete raw data)
# - Include ingestion metadata
# - Partition by ingestion date (not business date)
# - Schema: flexible (can store as strings or use schema inference)
# - Retention: long (compliance, reprocessing)

bronze_df = spark.readStream.format("cloudFiles") \
    .option("cloudFiles.format", "json") \
    .option("cloudFiles.schemaLocation", schema_loc) \
    .load(landing_path) \
    .withColumn("_ingested_at", current_timestamp()) \
    .withColumn("_source_file", input_file_name()) \
    .withColumn("_batch_id", lit(batch_id))

bronze_df.writeStream.format("delta") \
    .partitionBy("_ingested_date") \
    .option("checkpointLocation", checkpoint) \
    .table("bronze_orders")

# Design decisions:
# - Apply data quality checks
# - Deduplicate
# - Standardize data types and values
# - Join reference data
# - Apply SCD logic here
# - Enable CDF for downstream
# - Partition by business date

silver_df = spark.readStream.table("bronze_orders") \
    .dropDuplicates(["order_id"]) \
    .filter(col("order_id").isNotNull()) \
    .filter(col("amount") > 0) \
    .withColumn("amount", col("amount").cast("decimal(10,2)")) \
    .withColumn("order_date", to_date(col("order_date_str"), "yyyy-MM-dd"))

# Design decisions:
# - Pre-aggregated for specific use cases
# - Denormalized (star schema or wide tables)
# - Heavily optimized (Z-ORDER, OPTIMIZE, caching)
# - Lower volume, higher query performance

gold_df = spark.table("silver_orders") \
    .groupBy("order_date", "category", "region") \
    .agg(
        count("*").alias("order_count"),
        sum("amount").alias("total_revenue"),
        avg("amount").alias("avg_order_value"),
        countDistinct("customer_id").alias("unique_customers")
    )

# Pattern: Quarantine bad records, pass good records

from pyspark.sql.functions import col, when

bronze_df = spark.read.table("bronze_orders")

# Define quality rules
quality_rules = {
    "valid_order_id": col("order_id").isNotNull(),
    "positive_amount": col("amount") > 0,
    "valid_date": col("order_date").isNotNull() & (col("order_date") >= "2020-01-01"),
    "valid_status": col("status").isin("pending", "completed", "cancelled", "shipped"),
}

# Apply all rules
quality_check = bronze_df
for rule_name, condition in quality_rules.items():
    quality_check = quality_check.withColumn(
        f"_rule_{rule_name}", when(condition, True).otherwise(False)
    )

# All rules must pass
all_rules = [f"_rule_{name}" for name in quality_rules]
quality_check = quality_check.withColumn(
    "_is_valid",
    reduce(lambda a, b: a & b, [col(r) for r in all_rules])
)

# Route records
good_records = quality_check.filter("_is_valid = true").drop(*all_rules, "_is_valid")
bad_records = quality_check.filter("_is_valid = false")

good_records.write.format("delta").mode("append").saveAsTable("silver_orders")
bad_records.write.format("delta").mode("append").saveAsTable("quarantine_orders")

df = spark.readStream.format("cloudFiles") \
    .option("cloudFiles.format", "json") \
    .option("cloudFiles.schemaLocation", "/schema/orders") \
    .option("cloudFiles.inferColumnTypes", "true") \
    .option("cloudFiles.useNotifications", "true")  # File notification mode \
    .load("/landing/orders/")

# Schema evolution modes
df = spark.readStream.format("cloudFiles") \
    .option("cloudFiles.format", "json") \
    .option("cloudFiles.schemaLocation", "/schema/orders") \
    .option("cloudFiles.schemaEvolutionMode", "addNewColumns")  # or rescue, failOnNewColumns, none
    .load("/landing/orders/")

.option("cloudFiles.schemaHints", "order_id LONG, amount DECIMAL(10,2)")

import dlt

@dlt.table
@dlt.expect("valid_id", "order_id IS NOT NULL")                    # Log violation, keep row
@dlt.expect_or_drop("positive_amount", "amount > 0")               # Drop violating rows
@dlt.expect_or_fail("valid_status", "status IN ('A','B','C')")     # Fail entire pipeline
def silver_orders():
    return dlt.read_stream("bronze_orders")

# Streaming Live Table (Bronze — append-only)
@dlt.table
def bronze_events():
    return spark.readStream.format("cloudFiles").load(path)

# Materialized View (Gold — automatically incremental)
@dlt.table
def gold_daily_metrics():
    return dlt.read("silver_events") \
        .groupBy("date").agg(count("*").alias("event_count"))

# View (intermediate, not materialized)
@dlt.view
def enriched_events():
    return dlt.read("bronze_events").join(dim_table, "key")

# DLT approach — 5 lines vs 50+ lines manually
dlt.create_streaming_table("dim_customer_scd2")

dlt.apply_changes(
    target="dim_customer_scd2",
    source="bronze_customer_cdc",
    keys=["customer_id"],                          # Business key
    sequence_by=col("updated_at"),                 # Ordering column
    apply_as_deletes=expr("operation = 'DELETE'"),  # Delete condition
    except_column_list=["operation", "_rescued_data"],
    stored_as_scd_type=2                           # SCD Type 2
)
# DLT auto-manages: __START_AT, __END_AT, __IS_CURRENT columns

5 minutes") \
    .withColumn("txn_count_1h", count("*").over(w_1h)) \
    .withColumn("total_amount_1h", sum("amount").over(w_1h)) \
    .withColumn("txn_count_24h", count("*").over(w_24h)) \
    .withColumn("avg_amount_24h", avg("amount").over(w_24h)) \
    .withColumn("amount_deviation", (col("amount") - col("avg_amount_24h")) / col("avg_amount_24h"))

# 3. ML Scoring (using foreachBatch for model inference)
def score_and_alert(batch_df, batch_id):
    import mlflow
    model = mlflow.pyfunc.load_model("models:/fraud_model/Production")

    # Score
    scored = batch_df.toPandas()
    scored["fraud_score"] = model.predict(scored[feature_cols])

    scored_df = spark.createDataFrame(scored)

    # Write all scored transactions
    scored_df.write.format("delta").mode("append").saveAsTable("silver_scored_transactions")

    # Alert on high-risk
    alerts = scored_df.filter(col("fraud_score") > 0.8)
    alerts.write.format("delta").mode("append").saveAsTable("gold_fraud_alerts")

    # Trigger webhook for real-time blocking
    if alerts.count() > 0:
        send_webhook(alerts.collect())

features.writeStream.foreachBatch(score_and_alert) \
    .option("checkpointLocation", "/checkpoints/fraud") \
    .trigger(processingTime="10 seconds") \
    .start()"># Architecture:
# Transaction stream (Kafka) → Bronze → Feature Engineering → ML Scoring → Alert

# 1. Bronze: Raw transactions
transactions = spark.readStream.format("kafka") \
    .option("subscribe", "transactions") \
    .load() \
    .select(from_json(col("value").cast("string"), txn_schema).alias("txn")) \
    .select("txn.*")

# 2. Feature Engineering: Rolling aggregates
w_1h = Window.partitionBy("card_id").orderBy("event_time").rangeBetween(-3600, 0)
w_24h = Window.partitionBy("card_id").orderBy("event_time").rangeBetween(-86400, 0)

features = transactions \
    .withWatermark("event_time", "5 minutes") \
    .withColumn("txn_count_1h", count("*").over(w_1h)) \
    .withColumn("total_amount_1h", sum("amount").over(w_1h)) \
    .withColumn("txn_count_24h", count("*").over(w_24h)) \
    .withColumn("avg_amount_24h", avg("amount").over(w_24h)) \
    .withColumn("amount_deviation", (col("amount") - col("avg_amount_24h")) / col("avg_amount_24h"))

# 3. ML Scoring (using foreachBatch for model inference)
def score_and_alert(batch_df, batch_id):
    import mlflow
    model = mlflow.pyfunc.load_model("models:/fraud_model/Production")

    # Score
    scored = batch_df.toPandas()
    scored["fraud_score"] = model.predict(scored[feature_cols])

    scored_df = spark.createDataFrame(scored)

    # Write all scored transactions
    scored_df.write.format("delta").mode("append").saveAsTable("silver_scored_transactions")

    # Alert on high-risk
    alerts = scored_df.filter(col("fraud_score") > 0.8)
    alerts.write.format("delta").mode("append").saveAsTable("gold_fraud_alerts")

    # Trigger webhook for real-time blocking
    if alerts.count() > 0:
        send_webhook(alerts.collect())

features.writeStream.foreachBatch(score_and_alert) \
    .option("checkpointLocation", "/checkpoints/fraud") \
    .trigger(processingTime="10 seconds") \
    .start()

# Phase 1: Historical Backfill (batch)
# Process month by month to manage memory

from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

start_date = datetime(2023, 1, 1)
end_date = datetime(2025, 1, 1)

current = start_date
while current < end_date:
    next_month = current + relativedelta(months=1)

    # Read one month at a time
    month_df = spark.read.parquet(f"/raw/data/year={current.year}/month={current.month:02d}/")

    # Transform
    transformed = apply_transformations(month_df)

    # Write idempotently using replaceWhere
    transformed.write.format("delta") \
        .mode("overwrite") \
        .option("replaceWhere", f"year = {current.year} AND month = {current.month}") \
        .saveAsTable("silver_table")

    current = next_month

# Optimize after backfill
spark.sql("OPTIMIZE silver_table ZORDER BY (customer_id)")

# Phase 2: Switch to Incremental (streaming)
spark.readStream.format("cloudFiles") \
    .option("cloudFiles.format", "parquet") \
    .load("/raw/data/") \
    .writeStream.format("delta") \
    .option("checkpointLocation", "/checkpoints/silver") \
    .trigger(availableNow=True) \  # Run as scheduled batch
    .table("silver_table")

# Principle: Re-running the pipeline produces the SAME result

# 1. Use MERGE instead of INSERT (handles reruns)
target.alias("t").merge(source.alias("s"), "t.id = s.id") \
    .whenMatchedUpdateAll().whenNotMatchedInsertAll().execute()

# 2. Use replaceWhere for partition-level overwrites
df.write.format("delta").mode("overwrite") \
    .option("replaceWhere", "date = '2025-01-15'") \
    .saveAsTable("orders")
# Re-running replaces the same partition — idempotent

# 3. Streaming: checkpoints ensure exactly-once
stream.writeStream.option("checkpointLocation", "/checkpoints/job1").start()
# Re-starting reads from last committed offset

# 4. Auto Loader: tracks processed files
spark.readStream.format("cloudFiles").load(path)
# Re-starting skips already-processed files

# 5. Batch jobs: use run_date parameter
def process_day(run_date):
    df = spark.read.parquet(f"/raw/{run_date}/")
    df.write.format("delta").mode("overwrite") \
        .option("replaceWhere", f"process_date = '{run_date}'") \
        .saveAsTable("silver_table")

# 1. Row count reconciliation
source_count = spark.read.table("bronze_orders").filter("date = '2025-01-15'").count()
silver_count = spark.read.table("silver_orders").filter("date = '2025-01-15'").count()
gold_count = spark.read.table("gold_daily_orders").filter("date = '2025-01-15'").count()

assert abs(source_count - silver_count) / source_count < 0.01, "Silver count drift > 1%"

# 2. Data freshness SLA
from pyspark.sql.functions import max as _max, current_timestamp, unix_timestamp

freshness = spark.table("silver_orders") \
    .select(
        (unix_timestamp(current_timestamp()) - unix_timestamp(_max("updated_at"))) / 60
    ).collect()[0][0]  # minutes since last update

assert freshness < 60, f"Data is {freshness} minutes stale!"

# 3. Streaming metrics
for query in spark.streams.active:
    progress = query.lastProgress
    print(f"Input rate: {progress['inputRowsPerSecond']}")
    print(f"Processing rate: {progress['processedRowsPerSecond']}")
    print(f"Batch duration: {progress['batchDuration']}")
    # Alert if processing rate < input rate (falling behind)

# 4. DLT expectation metrics (available in DLT event log)
expectations = spark.read.table("event_log") \
    .filter("event_type = 'flow_progress'") \
    .select("timestamp", "details:flow_progress:data_quality:*")