import sys

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
from pyspark.sql.types import DoubleType, StringType, StructField, StructType

if len(sys.argv) != 3:
    raise SystemExit("Usage: s3_read_write.py <source-s3a-csv> <output-s3a-parquet>")

source_path = sys.argv[1]
output_path = sys.argv[2]

spark = (
    SparkSession.builder
    .appName("sg-s3-read-write")
    .config("spark.ui.showConsoleProgress", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

schema = StructType([
    StructField("order_id", StringType(), False),
    StructField("region", StringType(), False),
    StructField("amount", DoubleType(), False),
    StructField("status", StringType(), False),
    StructField("order_date", StringType(), False),
])

orders = (
    spark.read
    .option("header", "true")
    .schema(schema)
    .csv(source_path)
)

paid_orders = (
    orders
    .where(col("status") == "paid")
    .select("order_id", "region", "amount", "order_date")
)

print(f"source_rows={orders.count()}")
paid_orders.orderBy("order_id").show(truncate=False)

(
    paid_orders
    .coalesce(1)
    .write
    .mode("overwrite")
    .partitionBy("region")
    .parquet(output_path)
)

read_back = spark.read.parquet(output_path).orderBy("order_id")

print(f"output_rows={read_back.count()}")
read_back.show(truncate=False)

spark.stop()