# Structured Streaming with checkpoint = at-least-once by default
# Checkpoint ensures no data loss (no missed messages)
# BUT: if crash occurs between write and checkpoint update → duplicate

query = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "orders")
  .load()
  .writeStream
  .format("parquet")  # non-transactional sink → at-least-once only
  .option("checkpointLocation", "s3://bucket/ckpt/orders/")
  .option("path", "s3://bucket/parquet/orders/")
  .start()
)
# Result: no data loss ✓, but possible duplicates ⚠️

# dropDuplicates() in streaming requires a watermark to bound state
from pyspark.sql import functions as F

query = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "orders")
  .load()
  .selectExpr("CAST(value AS STRING) AS json")
  .select(F.from_json("json", order_schema).alias("d")).select("d.*")
  # Step 1: Add watermark so Spark knows when to expire dedup state
  .withWatermark("event_time", "1 hour")
  # Step 2: Drop duplicate order_ids within the watermark window
  .dropDuplicates(["order_id", "event_time"])
  .writeStream
  .format("delta")
  .option("checkpointLocation", "s3://bucket/ckpt/orders-dedup/")
  .start("s3://bucket/delta/orders/")
)
# dropDuplicates keeps state for all seen order_ids within 1 hour watermark
# If same order_id arrives twice → second one is dropped

from pyspark.sql import functions as F
from delta.tables import DeltaTable

def idempotent_write(batch_df, batch_id):
    """Write a batch idempotently using batchId as dedup key."""

    # Check if this batch was already processed
    control_table_path = "s3://bucket/delta/control/processed_batches"

    try:
        processed = spark.read.format("delta").load(control_table_path)
        already_done = processed.filter(F.col("batch_id") == batch_id).count() > 0
    except:
        already_done = False  # First run, control table doesn't exist yet

    if already_done:
        print(f"Batch {batch_id} already processed — skipping (idempotent)")
        return

    # Write actual data
    (batch_df
      .write
      .format("delta")
      .mode("append")
      .save("s3://bucket/delta/orders/")
    )

    # Record that this batch is done (atomic with the write above if using Delta)
    batch_marker = spark.createDataFrame(
      [(batch_id, F.current_timestamp())],
      ["batch_id", "processed_at"]
    )
    (batch_marker
      .write
      .format("delta")
      .mode("append")
      .save(control_table_path)
    )
    print(f"Batch {batch_id} written and committed ✓")

# Use the idempotent writer in streaming
query = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "orders")
  .load()
  .writeStream
  .foreachBatch(idempotent_write)
  .option("checkpointLocation", "s3://bucket/ckpt/orders/")
  .start()
)

from delta.tables import DeltaTable
from pyspark.sql import functions as F

def upsert_to_delta(batch_df, batch_id):
    """Idempotent MERGE — running this twice = same result as once."""

    target = DeltaTable.forPath(spark, "s3://bucket/delta/orders/")

    # MERGE: match on order_id (business key)
    # If order_id exists → UPDATE (idempotent: same data written again = no change)
    # If order_id doesn't exist → INSERT
    (target.alias("target")
      .merge(
        batch_df.alias("source"),
        "target.order_id = source.order_id"
      )
      .whenMatchedUpdateAll()   # update if key matches
      .whenNotMatchedInsertAll() # insert if new record
      .execute()
    )
    print(f"Batch {batch_id} upserted to Delta ✓")

query = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "orders")
  .load()
  .selectExpr("CAST(value AS STRING) AS json")
  .select(F.from_json("json", order_schema).alias("d")).select("d.*")
  .writeStream
  .foreachBatch(upsert_to_delta)
  .option("checkpointLocation", "s3://bucket/ckpt/orders-upsert/")
  .start()
)

import psycopg2

def jdbc_upsert(batch_df, batch_id):
    """Idempotent JDBC write using ON CONFLICT DO UPDATE."""

    # Collect to driver (only for small batches!)
    rows = batch_df.collect()

    conn = psycopg2.connect(
      host="postgres-host", database="mydb",
      user="user", password="secret"
    )
    cursor = conn.cursor()

    for row in rows:
        cursor.execute("""
          INSERT INTO orders (order_id, amount, status)
          VALUES (%s, %s, %s)
          ON CONFLICT (order_id)            -- idempotent: if exists, update
          DO UPDATE SET
            amount = EXCLUDED.amount,
            status = EXCLUDED.status
        """, (row["order_id"], row["amount"], row["status"]))

    conn.commit()
    cursor.close()
    conn.close()
    print(f"Batch {batch_id}: {len(rows)} rows upserted to PostgreSQL ✓")

query = (df_stream
  .writeStream
  .foreachBatch(jdbc_upsert)
  .option("checkpointLocation", "s3://bucket/ckpt/jdbc-orders/")
  .start()
)

# Delta streaming write — transactional by nature
# Each micro-batch = one Delta transaction
# Either the entire batch commits atomically, or nothing commits

query = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "payments")
  .load()
  .writeStream
  .format("delta")        # transactional ACID sink
  .outputMode("append")
  .option("checkpointLocation", "s3://bucket/ckpt/payments/")
  .start("s3://bucket/delta/payments/")
)

# What happens internally each batch:
# 1. Spark writes parquet files to _delta_log staging area
# 2. Spark writes atomic commit: {"add": [file1, file2], "txnId": "batch-5-abc"}
# 3. If crash before commit → files are orphans → VACUUM cleans them up
# 4. Next restart replays batch 5 → new files + new commit → clean state

# Verify transactions in Delta history
from delta.tables import DeltaTable
dt = DeltaTable.forPath(spark, "s3://bucket/delta/payments/")
dt.history().select("version", "timestamp", "operation", "operationMetrics").show(5)

# Writing to Kafka sink with transactional settings
# These options enable exactly-once Kafka-to-Kafka via transactions

query = (df_processed
  .writeStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("topic", "processed-events")
  # Enable idempotent producer (exactly-once within single producer session)
  .option("kafka.enable.idempotence", "true")
  # Required for transactions: acks must be "all"
  .option("kafka.acks", "all")
  # Transactional ID (Spark auto-generates per partition if not set)
  # For exactly-once Kafka-to-Kafka, Spark manages this internally
  .option("checkpointLocation", "s3://bucket/ckpt/kafka-to-kafka/")
  .start()
)

# Kafka producer configs that matter for exactly-once:
# enable.idempotence=true  → deduplicates retries within a session
# acks=all                 → leader + all replicas acknowledge
# max.in.flight.requests.per.connection=5  → allows pipelining safely
# retries → must be > 0 (default is very high now)

# Full exactly-once Kafka → Spark → Kafka pipeline
from pyspark.sql import functions as F

# Read from source Kafka topic
df_source = (spark
  .readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "raw-events")
  .option("startingOffsets", "latest")
  .load()
)

# Transform
df_processed = (df_source
  .selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
  .withColumn("enriched_value", F.concat(F.col("value"), F.lit("_processed")))
  .select(
    F.col("key"),
    F.col("enriched_value").alias("value")  # Kafka sink needs "value" column
  )
)

# Write to output Kafka topic with exactly-once
query = (df_processed
  .writeStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("topic", "processed-events")
  .option("kafka.enable.idempotence", "true")
  .option("kafka.acks", "all")
  .option("checkpointLocation", "s3://bucket/ckpt/k2k-pipeline/")
  .start()
)

# When reading from a Kafka topic that uses transactions,
# set isolation.level=read_committed to skip aborted transactions

df_committed_only = (spark
  .readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "processed-events")
  # Only read messages from COMMITTED transactions
  .option("kafka.isolation.level", "read_committed")
  .load()
)

# Without read_committed: you might read messages from aborted transactions
# (messages that were written but the transaction was rolled back)
# This would violate exactly-once on the consumer side

# Default is read_uncommitted — must explicitly set read_committed
# for correct exactly-once behavior in transactional pipelines

# The simplest exactly-once pattern — just use Delta + checkpoint
# Delta's transaction log handles idempotency automatically

from pyspark.sql import functions as F

# Parse incoming Kafka events
df_events = (spark
  .readStream.format("kafka")
  .option("kafka.bootstrap.servers", "broker:9092")
  .option("subscribe", "user-events")
  .load()
  .selectExpr("CAST(value AS STRING) AS raw_json")
  .withColumn("data", F.from_json("raw_json", event_schema))
  .select("data.*")
  .withColumn("ingested_at", F.current_timestamp())
)

# Write to Delta — exactly-once via Delta's built-in batchId dedup
query = (df_events
  .writeStream
  .format("delta")
  .outputMode("append")
  .option("checkpointLocation", "s3://bucket/ckpt/user-events/")
  .trigger(processingTime="60 seconds")
  .start("s3://bucket/delta/user-events/")
)
# That's it! No extra idempotency code needed.
# Delta's transaction log guarantees exactly-once writes automatically.

from delta.tables import DeltaTable
from pyspark.sql import functions as F

def write_with_batch_id(batch_df, batch_id):
    """Write to Delta MERGE with explicit batchId for auditability."""

    # Add batchId as a column for auditing
    batch_df_with_meta = batch_df.withColumn("_batch_id", F.lit(batch_id))

    target_path = "s3://bucket/delta/orders/"

    if DeltaTable.isDeltaTable(spark, target_path):
        target = DeltaTable.forPath(spark, target_path)
        (target.alias("t")
          .merge(batch_df_with_meta.alias("s"), "t.order_id = s.order_id")
          # Idempotent: same order written again just updates to same values
          .whenMatchedUpdateAll()
          .whenNotMatchedInsertAll()
          .execute()
        )
    else:
        # First run: create the table
        (batch_df_with_meta
          .write.format("delta").mode("overwrite").save(target_path))

    print(f"Batch {batch_id} merged to Delta ✓")

query = (df_stream
  .writeStream
  .foreachBatch(write_with_batch_id)
  .option("checkpointLocation", "s3://bucket/ckpt/orders-merge/")
  .start()
)

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, TimestampType

# ── Spark Session ──────────────────────────────────────────────────────────
spark = (SparkSession.builder
  .appName("ExactlyOncePipeline")
  .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
  .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog")
  .getOrCreate()
)

# ── Schema ────────────────────────────────────────────────────────────────
order_schema = StructType([
  StructField("order_id", StringType()),
  StructField("customer_id", StringType()),
  StructField("amount", DoubleType()),
  StructField("status", StringType()),
  StructField("order_time", TimestampType())
])

# ── LAYER 1: SOURCE — Kafka with checkpointed offsets ─────────────────────
# Kafka is replayable → at-least-once source
# Checkpoint saves exact offsets → prevents re-reads or skips
df_raw = (spark
  .readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092")
  .option("subscribe", "orders")
  .option("startingOffsets", "latest")     # only used on first start
  .option("failOnDataLoss", "true")         # alert on Kafka retention gap
  .option("maxOffsetsPerTrigger", "50000")  # control batch size
  .load()
)

# ── LAYER 2: PROCESSING — Spark deterministic transformation ──────────────
# Same input offsets always → same output rows (deterministic)
df_orders = (df_raw
  .selectExpr("CAST(value AS STRING) AS json")
  .select(F.from_json(F.col("json"), order_schema).alias("d"))
  .select("d.*")
  .filter(F.col("order_id").isNotNull())   # filter invalid records
  .withColumn("ingested_at", F.current_timestamp())
  .withColumn("amount_usd", F.round(F.col("amount"), 2))
)

# ── LAYER 3: SINK — Delta (idempotent via transaction log batchId) ─────────
# Delta embeds streaming query ID + batchId in every commit
# Re-running a batch → Delta detects duplicate batchId → skips write
query = (df_orders
  .writeStream
  .format("delta")                                # ACID transactional sink
  .outputMode("append")
  .option("checkpointLocation",                   # saves Kafka offsets
         "s3://prod/checkpoints/orders-pipeline/")
  .trigger(processingTime="30 seconds")
  .start("s3://prod/delta/bronze/orders/")
)

query.awaitTermination()

# ═══ EXACTLY-ONCE GUARANTEE SUMMARY ════════════════════════════════════════
# SOURCE:     Kafka offset checkpointing → no skips, no double reads
# PROCESSING: Deterministic Spark transforms → same input = same output
# SINK:       Delta batchId dedup → even if batch reruns, no duplicate rows
# RESULT:     End-to-end exactly-once ✅

# Verify no duplicates in Delta sink
df_orders_delta = spark.read.format("delta").load("s3://prod/delta/bronze/orders/")

total_rows = df_orders_delta.count()
unique_orders = df_orders_delta.select("order_id").distinct().count()

print(f"Total rows: {total_rows}")
print(f"Unique order_ids: {unique_orders}")
print(f"Duplicates: {total_rows - unique_orders}")
# Expected: Duplicates = 0  ✅

# Find actual duplicates if any
duplicates = (df_orders_delta
  .groupBy("order_id")
  .count()
  .filter(F.col("count") > 1)
)
print(f"Duplicate order IDs found: {duplicates.count()}")
duplicates.show(10)

# Check Delta transaction history for batchId tracking
from delta.tables import DeltaTable
dt = DeltaTable.forPath(spark, "s3://prod/delta/bronze/orders/")
dt.history().select(
  "version", "timestamp", "operation",
  "operationMetrics.numOutputRows"
).show(10)