from pathlib import Path
import json
import shutil

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


base_dir = Path("/tmp/sg-structured-streaming")
input_dir = base_dir / "input"
checkpoint_dir = base_dir / "checkpoint"

shutil.rmtree(base_dir, ignore_errors=True)
input_dir.mkdir(parents=True)

spark = (
    SparkSession.builder
    .master("local[2]")
    .appName("sg-structured-streaming-run")
    .config("spark.ui.showConsoleProgress", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

schema = T.StructType(
    [
        T.StructField("event", T.StringType(), False),
        T.StructField("amount", T.IntegerType(), False),
    ]
)

events = spark.readStream.schema(schema).json(str(input_dir))
totals = events.groupBy("event").agg(F.sum("amount").alias("total"))

query = (
    totals.writeStream
    .format("memory")
    .queryName("event_totals")
    .outputMode("complete")
    .option("checkpointLocation", str(checkpoint_dir))
    .start()
)


def add_batch(name, rows):
    tmp_file = input_dir / f".{name}.tmp"
    final_file = input_dir / name
    with tmp_file.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row) + "\n")
    tmp_file.replace(final_file)


def show_totals(label):
    print(label)
    spark.sql(
        "SELECT event, total FROM event_totals ORDER BY event"
    ).show(truncate=False)


try:
    print(f"query_name={query.name}")
    print(f"query_active={query.isActive}")

    add_batch(
        "batch-001.json",
        [
            {"event": "checkout", "amount": 4},
            {"event": "search", "amount": 2},
            {"event": "checkout", "amount": 3},
        ],
    )
    query.processAllAvailable()
    show_totals("after_batch_1")

    add_batch(
        "batch-002.json",
        [
            {"event": "search", "amount": 5},
            {"event": "view", "amount": 8},
        ],
    )
    query.processAllAvailable()
    show_totals("after_batch_2")

    print(f"last_progress_batch={query.lastProgress['batchId']}")
    print(f"last_progress_input_rows={query.lastProgress['numInputRows']}")
finally:
    query.stop()
    print(f"query_active_after_stop={query.isActive}")
    print(f"active_streams_after_stop={len(spark.streams.active)}")
    spark.stop()
    shutil.rmtree(base_dir, ignore_errors=True)
    print("cleanup_done=true")