import json, re, time, urllib.request
from datetime import datetime
from pyspark.sql import functions as F

sc = spark.sparkContext
stg_orders.createOrReplaceTempView("stg_orders")
stg_users.createOrReplaceTempView("stg_users")

# ===== Exercise 1: birthdate_is_valid =====
print("\n===== EX1 =====")
first_orders = stg_orders.groupBy("user_id").agg(F.min("ordered_at").alias("first_order_at"))
users_ex1 = (
    stg_users.join(first_orders, "user_id", "left")
    .withColumn("birthdate_is_valid",
                F.coalesce(F.col("birthdate") <= F.to_date("first_order_at"), F.lit(True)))
    .drop("first_order_at")
)
users_ex1_sql = spark.sql('''
    SELECT u.*,
           coalesce(u.birthdate <= to_date(f.first_order_at), true) AS birthdate_is_valid
    FROM stg_users u
    LEFT JOIN (SELECT user_id, min(ordered_at) AS first_order_at
               FROM stg_orders GROUP BY user_id) f
      ON u.user_id = f.user_id
''')
users_ex1.groupBy("birthdate_is_valid").count().show()
users_ex1.filter(~F.col("birthdate_is_valid")).select("user_id", "username", "birthdate").show(truncate=False)
a = users_ex1.select("user_id", "birthdate_is_valid")
b = users_ex1_sql.select("user_id", "birthdate_is_valid")
print("EX1 differences DataFrame vs SQL :", a.exceptAll(b).count() + b.exceptAll(a).count())

# ===== Exercise 2: zip_code with Python UDF and pandas_udf =====
print("\n===== EX2 =====")
ZIP_RE = re.compile(r"([A-Z]{2}) ([0-9]{5})$")

def zip_py(address):
    lines = (address or "").split("\n")
    m = ZIP_RE.search(lines[1]) if len(lines) > 1 else None
    return m.group(2) if m else ""

zip_udf = F.udf(zip_py, "string")
zip_regexp = F.regexp_extract(F.element_at(F.split("address", "\n"), 2), "([A-Z]{2}) ([0-9]{5})$", 2)

def bench(label, expr, runs=5):
    times = []
    for _ in range(runs):
        t = time.perf_counter()
        users_bronze.select(expr.alias("zip")).write.format("noop").mode("overwrite").save()
        times.append(time.perf_counter() - t)
    print(f"{label:<14} best {min(times):.3f} s, mean {sum(times)/len(times):.3f} s")

bench("regexp_extract", zip_regexp)
bench("python udf", zip_udf("address"))
print("udf vs regexp mismatches :",
      users_bronze.filter(zip_udf("address") != zip_regexp).count())

try:
    import pandas as pd

    @F.pandas_udf("string")
    def zip_pandas(address: pd.Series) -> pd.Series:
        line2 = address.fillna("").str.split("\n").str[1]
        return line2.str.extract(r"([A-Z]{2}) ([0-9]{5})$")[1].fillna("")

    bench("pandas_udf", zip_pandas("address"))
    print("pandas_udf vs regexp mismatches :",
          users_bronze.filter(zip_pandas("address") != zip_regexp).count())
except Exception as e:
    print("pandas_udf not available :", type(e).__name__, e)

# ===== Exercise 3: cache the bronze DataFrames =====
print("\n===== EX3 =====")
def timed(df):
    t = time.perf_counter(); df.count(); return time.perf_counter() - t

print(f"without cache : users {timed(users_bronze):.3f} s, orders {timed(orders_bronze):.3f} s")
users_bronze.cache(); orders_bronze.cache()
users_bronze.count(); orders_bronze.count()  # fills the cache
print(f"with cache    : users {timed(users_bronze):.3f} s, orders {timed(orders_bronze):.3f} s")
url = sc.uiWebUrl + "/api/v1/applications/" + sc.applicationId + "/storage/rdd"
for r in json.load(urllib.request.urlopen(url)):
    print("  STORAGE", r["name"][:60], r["numCachedPartitions"], "/", r["numPartitions"],
          "partitions,", r["memoryUsed"], "bytes")

# ===== Exercise 4: append and time travel =====
print("\n===== EX4 =====")
snap_before = spark.sql(
    "SELECT snapshot_id FROM lab.silver.orders.snapshots ORDER BY committed_at DESC LIMIT 1"
).first()[0]
new_order = spark.createDataFrame(
    [("ex4-new-order", stg_orders.first()["user_id"], datetime(2026, 10, 28, 12, 0), 1, "bread")],
    schema=stg_orders.schema,
)
new_order.writeTo("lab.silver.orders").append()

spark.sql("SELECT snapshot_id, operation, committed_at FROM lab.silver.orders.snapshots").show(truncate=False)
print("Iceberg now    :", spark.table("lab.silver.orders").count())
print("Iceberg before :", spark.sql(f"SELECT * FROM lab.silver.orders VERSION AS OF {snap_before}").count())

# Same thing on a Parquet copy, so the silver Parquet used for the dbt comparison stays intact
ex4_path = f"s3a://{BUCKET}/silver/spark/orders_ex4.parquet"
stg_orders.write.mode("overwrite").parquet(ex4_path)
print("Parquet before append :", spark.read.parquet(ex4_path).count())
new_order.write.mode("append").parquet(ex4_path)
print("Parquet after append  :", spark.read.parquet(ex4_path).count())