from pyspark.sql.functions import count, sum as spark_sum, to_timestamp, col

# Each unique user_id → 1 state entry
# State entry stores: running count + running sum
agg_df = orders_stream \
    .withWatermark("event_time", "10 minutes") \
    .groupBy("user_id") \  # ← this becomes the state KEY
    .agg(
        count("*").alias("order_count"),       # state VALUE field
        spark_sum("amount").alias("total_spent")  # state VALUE field
    )

# Compound key example:
# state key = (user_id, product_category) tuple
agg_df2 = orders_stream \
    .withWatermark("event_time", "10 minutes") \
    .groupBy("user_id", "product_category") \
    .agg(count("*").alias("count"))

# Rule of thumb for streaming:
# numPartitions ≈ 2-3x number of executor cores
# (too many → small state files, overhead; too few → hotspots)

# For a cluster with 4 executors x 4 cores = 16 cores:
spark.conf.set("spark.sql.shuffle.partitions", "32")

# IMPORTANT: Once a streaming query starts, you CANNOT change
# spark.sql.shuffle.partitions without deleting the checkpoint.
# Changing it changes the partition layout → state keys move → broken!

# Verify your setting:
print(spark.conf.get("spark.sql.shuffle.partitions"))

import time

# Monitor state size after every few batches
while query.isActive:
    time.sleep(30)
    progress = query.lastProgress
    if progress and progress.get("stateOperators"):
        for op in progress["stateOperators"]:
            print(f"State rows: {op['numRowsTotal']}")         # TOTAL keys in store
            print(f"State memory: {op['memoryUsedBytes']} bytes")   # memory footprint
            print(f"Rows updated: {op['numRowsUpdated']}")        # keys touched this batch
            print(f"Rows removed: {op['numRowsRemoved']}")        # keys cleaned up this batch
            print("---")

# Healthy pattern: numRowsRemoved > 0 (cleanup happening)
# Danger signal: numRowsTotal keeps increasing, numRowsRemoved = 0

from pyspark.sql.streaming import GroupState, GroupStateTimeout
from pyspark.sql.types import StructType, StringType, LongType
from typing import Iterator, Tuple

# State schema: what we remember per user
StateType = StructType() \
    .add("order_count", LongType()) \
    .add("total_amount", LongType())

def update_user_state(
    user_id: str,
    new_orders,
    state: GroupState
) -> Tuple[str, int, int]:
    """
    Called per unique user_id per batch.
    state: the persisted state for this user (survives across batches)
    """
    
    # Check if this was a TIMEOUT call (no new data, state expired)
    if state.hasTimedOut:
        # Clean up expired state — user hasn't ordered in 30 minutes
        state.remove()
        return (user_id, 0, 0)  # or yield nothing for flatMap version
    
    # Get current state (or initialize if first time)
    if state.exists:
        curr_count = state.get["order_count"]
        curr_amount = state.get["total_amount"]
    else:
        curr_count, curr_amount = 0, 0
    
    # Process new rows for this user in this batch
    for order in new_orders:
        curr_count += 1
        curr_amount += order.amount
    
    # Update state
    state.update((curr_count, curr_amount))
    
    # Set TTL: if no new orders arrive for 30 min → evict this user's state
    state.setTimeoutDuration("30 minutes")  # ProcessingTimeTimeout
    
    return (user_id, curr_count, curr_amount)

# Apply the stateful function
output_schema = "user_id STRING, order_count LONG, total_amount LONG"

result = orders_stream \
    .groupBy("user_id") \
    .applyInPandas(update_user_state, schema=output_schema)
# Note: mapGroupsWithState uses .mapGroupsWithState() for typed Dataset API

from pyspark.sql.functions import window, count, sum as spark_sum

# Tumbling window: count orders per 1-hour window
# Watermark: tolerate up to 10 minutes of late data
# State cleanup: once watermark passes the window end → evict that window's state

windowed_agg = orders_stream \
    .withWatermark("event_time", "10 minutes") \  # ← cleanup trigger
    .groupBy(
        window(col("event_time"), "1 hour"),  # tumbling 1-hour window
        col("product_category")
    ) \
    .agg(
        count("*").alias("order_count"),
        spark_sum("amount").alias("revenue")
    )

# State lifecycle for window [10:00–11:00):
# Batch 1: watermark = 09:45 → window still open → state kept
# Batch 2: watermark = 10:50 → window still open → state kept
# Batch 3: watermark = 11:10 → window [10:00-11:00) is PAST watermark
#   → state for this window is EVICTED ✅
#   → result emitted in Append mode (window is now final)

query = windowed_agg.writeStream \
    .outputMode("append") \  # append: only emit when window is finalized
    .option("checkpointLocation", "/tmp/ckpt/windowed") \
    .format("delta") \
    .start("/delta/hourly_revenue")

def session_tracker(session_id, events, state: GroupState):
    # Explicitly remove state when session is "completed"
    for event in events:
        if event.event_type == "SESSION_END":
            # Session is done — no point keeping state
            state.remove()  # ← explicit cleanup
            return (session_id, "CLOSED", state.get["page_views"] if state.exists else 0)
    
    # Otherwise, update state and set TTL
    curr = state.get["page_views"] if state.exists else 0
    state.update((curr + len(list(events)),))
    state.setTimeoutDuration("30 minutes")  # auto-cleanup if no events for 30min
    return (session_id, "ACTIVE", state.get["page_views"])

spark = SparkSession.builder \
    .appName("RocksDBMemoryTuning") \
    # Enable RocksDB
    .config("spark.sql.streaming.stateStore.providerClass",
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider") \
    # Off-heap memory for RocksDB block cache per partition
    .config("spark.sql.streaming.stateStore.rocksdb.blockCacheSizeMB", "64") \
    # Enable off-heap memory (required for RocksDB)
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "2g") \   # per executor
    # Total executor memory = spark.executor.memory + offHeap
    .config("spark.executor.memory", "4g") \
    .getOrCreate()

# Memory layout per executor:
#   4g  on-heap  (for Spark execution, shuffle buffers, etc.)
#   2g  off-heap (for RocksDB block cache)
#   Total JVM memory request: ~6g + JVM overhead (~1g) = ~7g

# Signs you're running out of state memory:
# 1. Executor logs show: java.lang.OutOfMemoryError: Java heap space
# 2. GC time in Spark UI keeps increasing per batch
# 3. stateOperators[0].memoryUsedBytes approaches executor heap size

# Check via query progress:
progress = query.lastProgress
state_ops = progress.get("stateOperators", [])
for op in state_ops:
    mem_mb = op["memoryUsedBytes"] / (1024 * 1024)
    rows = op["numRowsTotal"]
    print(f"State: {rows:,} rows, {mem_mb:.1f} MB")
    
    # Alarm if state exceeds 80% of executor memory
    executor_heap_mb = 4096  # 4GB executor
    if mem_mb > executor_heap_mb * 0.8:
        print(f"⚠️  WARNING: State is {mem_mb/executor_heap_mb*100:.0f}% of executor heap!")
        print("Consider: (1) Enable watermark (2) Switch to RocksDB")

spark = SparkSession.builder \
    .config("spark.sql.streaming.stateStore.providerClass",
            "org.apache.spark.sql.execution.streaming.state.RocksDBStateStoreProvider") \
    # How often to write a full snapshot (default: 1000 batches)
    .config("spark.sql.streaming.stateStore.rocksdb.minDeltasForSnapshot", "1000") \
    # Checkpoint format: native (default), or zipped
    .config("spark.sql.streaming.stateStore.rocksdb.enableChangelogCheckpointing", "true") \
    .getOrCreate()

# With changelog checkpointing:
# batch 1: writes keys {A:1, B:2} → delta file
# batch 2: key A updated to {A:5} → delta file (only A written)
# batch 1000: full snapshot written
# Recovery: apply snapshot + all deltas since snapshot

# Reduce how many old checkpoint versions are kept
spark.conf.set("spark.sql.streaming.minBatchesToRetain", "5")
# default is 100 — this keeps 95 unnecessary old versions!

# What this controls:
# - Old offset log files (older than minBatchesToRetain batches) are cleaned
# - Old commit log files are cleaned
# - Old state snapshots are cleaned (RocksDB keeps what's needed for recovery)

# A value of 5 means: keep 5 recent batches of state history
# Good for recovery windows of a few minutes
# If your recovery window is longer, increase this value

# Cleaning old checkpoints manually (if needed):
# hdfs dfs -rm -r /ckpt/orders/state/0/0/1.delta  # don't do manually!
# Always let Spark manage this via minBatchesToRetain

from pyspark.sql.streaming import GroupState, GroupStateTimeout

def my_state_function(key, values, state: GroupState):
    
    # ── READ state ──────────────────────────────────────
    state.exists          # bool: does state exist for this key?
    state.get             # get current state (raises if not exists)
    state.getOption       # returns None if not exists
    
    # ── WRITE state ─────────────────────────────────────
    state.update(new_val) # replace state with new value
    state.remove()        # delete state for this key entirely
    
    # ── TIMEOUT ─────────────────────────────────────────
    state.hasTimedOut     # bool: was this call triggered by timeout?
    state.setTimeoutDuration("30 minutes")  # ProcessingTime timeout
    state.setTimeoutTimestamp(ts_ms)         # EventTime timeout (epoch ms)
    
    # ── METADATA ────────────────────────────────────────
    state.getCurrentProcessingTimeMs()  # current processing time
    state.getCurrentWatermarkMs()       # current watermark (for EventTime timeout)

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, from_json, to_timestamp
from pyspark.sql.types import (
    StructType, StringType, LongType, IntegerType
)
from pyspark.sql.streaming import GroupState, GroupStateTimeout

spark = SparkSession.builder.appName("SessionTracker").getOrCreate()

# ── State Schema: what we store per user ──────────────
# (page_views, total_time_sec, session_start_ms)
SessionState = StructType() \
    .add("page_views", IntegerType()) \
    .add("total_time_sec", LongType()) \
    .add("session_start_ms", LongType())

# ── Input Schema ─────────────────────────────────────
EventSchema = StructType() \
    .add("user_id", StringType()) \
    .add("page", StringType()) \
    .add("duration_sec", IntegerType()) \
    .add("event_time", StringType())

# ── State Update Function ────────────────────────────
def track_session(user_id: str, events, state: GroupState):
    """
    Returns one summary row per user per batch.
    State = running session totals.
    """
    
    # Handle timeout: session expired (no activity for 30 min)
    if state.hasTimedOut:
        if state.exists:
            s = state.get
            state.remove()
            return (user_id, s["page_views"], s["total_time_sec"], "EXPIRED")
        return (user_id, 0, 0, "EXPIRED_EMPTY")
    
    # Initialize or restore state
    if state.exists:
        page_views = state.get["page_views"]
        total_time = state.get["total_time_sec"]
        start_ms   = state.get["session_start_ms"]
    else:
        page_views, total_time, start_ms = 0, 0, state.getCurrentProcessingTimeMs()
    
    # Process new events for this user in this batch
    for event in events:
        page_views += 1
        total_time += event.duration_sec
    
    # Persist updated state
    state.update((page_views, total_time, start_ms))
    
    # Reset the 30-minute inactivity timeout
    state.setTimeoutDuration("30 minutes")
    
    return (user_id, page_views, total_time, "ACTIVE")

# ── Output Schema ─────────────────────────────────────
output_schema = "user_id STRING, page_views INT, total_time_sec LONG, status STRING"

# ── Pipeline ──────────────────────────────────────────
events_stream = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "clickstream") \
    .load() \
    .select(from_json(col("value").cast("string"), EventSchema).alias("e")) \
    .select("e.*")

session_df = events_stream \
    .groupBy("user_id") \
    .mapGroupsWithState(
        track_session,
        outputSchema=output_schema,
        stateSchema=SessionState,
        timeoutConf=GroupStateTimeout.ProcessingTimeTimeout
    )

query = session_df.writeStream \
    .outputMode("update") \  # mapGroupsWithState supports update or append
    .format("delta") \
    .option("checkpointLocation", "/tmp/ckpt/sessions") \
    .trigger(processingTime="10 seconds") \
    .start("/delta/sessions")

query.awaitTermination()

from pyspark.sql.streaming import GroupState, GroupStateTimeout
from pyspark.sql.types import StructType, StringType, IntegerType
from typing import Iterator, Tuple

# State: count of failed logins per user
AlertState = StructType().add("failed_attempts", IntegerType())

def detect_fraud(
    user_id: str,
    events,
    state: GroupState
) -> Iterator[Tuple[str, str, int]]:
    """
    flatMapGroupsWithState: can yield 0 or many rows.
    Only emits when fraud is detected.
    """
    
    # Timeout: reset counter after 10 minutes of no activity
    if state.hasTimedOut:
        state.remove()
        return  # yield nothing — no output for timeout cleanup
    
    curr_attempts = state.get["failed_attempts"] if state.exists else 0
    
    for event in events:
        if event.event_type == "LOGIN_FAILED":
            curr_attempts += 1
            
            if curr_attempts >= 5:
                # ALERT! Emit a fraud alert row
                yield (user_id, "FRAUD_ALERT", curr_attempts)  # ← yield, not return
                state.remove()  # reset after alerting
                curr_attempts = 0
            # else: accumulate silently — NO output emitted
        
        elif event.event_type == "LOGIN_SUCCESS":
            curr_attempts = 0  # reset on successful login
    
    # Update state and reset timeout
    if curr_attempts > 0:
        state.update((curr_attempts,))
        state.setTimeoutDuration("10 minutes")
    else:
        if state.exists:
            state.remove()
    # If no alert triggered → function returns without yielding → 0 output rows

# Apply flatMapGroupsWithState
output_schema = "user_id STRING, alert_type STRING, attempt_count INT"

fraud_alerts = login_stream \
    .groupBy("user_id") \
    .flatMapGroupsWithState(
        detect_fraud,
        outputMode="append",      # ← append required for flatMap
        stateSchema=AlertState,
        outputSchema=output_schema,
        timeoutConf=GroupStateTimeout.ProcessingTimeTimeout
    )

query = fraud_alerts.writeStream \
    .outputMode("append") \
    .format("delta") \
    .option("checkpointLocation", "/tmp/ckpt/fraud") \
    .start("/delta/fraud_alerts")