from pyspark.sql import SparkSession
from pyspark.sql import functions as F


spark = (
    SparkSession.builder.master("local[1]")
    .appName("sg-spark-csv-read-write")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

schema = "order_id STRING, region STRING, item_count INT, order_total DOUBLE, notes STRING"

orders = (
    spark.read.schema(schema)
    .option("header", True)
    .option("mode", "FAILFAST")
    .csv("orders.csv")
)

priority_orders = (
    orders.where((F.col("region") == "APAC") & (F.col("item_count") >= 3))
    .select("order_id", "region", "item_count", "order_total", "notes")
    .orderBy("order_id")
)

print("Input rows:", orders.count())
orders.printSchema()
print("Priority rows before write:")
priority_orders.show(truncate=False)

(
    priority_orders.coalesce(1)
    .write.mode("overwrite")
    .option("header", True)
    .csv("output/csv-priority-orders")
)

read_back = (
    spark.read.schema(priority_orders.schema)
    .option("header", True)
    .csv("output/csv-priority-orders")
)

print("Output rows:", read_back.count())
read_back.orderBy("order_id").show(truncate=False)

spark.stop()