import json
import shutil
from pathlib import Path

from pyspark.sql import SparkSession


base = Path("/tmp/sg-spark-streaming-checkpoint")
input_dir = base / "input"
output_dir = base / "output"
checkpoint_dir = base / "checkpoint"

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


def write_events(path, rows):
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")


write_events(
    input_dir / "events-001.json",
    [
        {"event": "checkout", "region": "MY", "count": 3},
        {"event": "search", "region": "SG", "count": 2},
    ],
)

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

schema = "event STRING, region STRING, count INT"


def run_query(label):
    events = spark.readStream.schema(schema).json(str(input_dir))
    query = (
        events.select("event", "region", "count")
        .writeStream.format("json")
        .queryName("sg_streaming_checkpoint")
        .option("path", str(output_dir))
        .option("checkpointLocation", str(checkpoint_dir))
        .trigger(availableNow=True)
        .start()
    )
    query.awaitTermination(60)
    progress = query.lastProgress or {}
    print(f"{label}_query_id={query.id}")
    print(f"{label}_run_id={query.runId}")
    print(f"{label}_input_rows={progress.get('numInputRows')}")
    return str(query.id), str(query.runId)


first_query_id, first_run_id = run_query("first")

write_events(
    input_dir / "events-002.json",
    [
        {"event": "checkout", "region": "MY", "count": 4},
    ],
)

second_query_id, second_run_id = run_query("second")

checkpoint_entries = sorted(
    path.name for path in checkpoint_dir.iterdir() if not path.name.startswith(".")
)
print("checkpoint_entries=" + ",".join(checkpoint_entries))
print(f"query_id_reused={first_query_id == second_query_id}")
print(f"run_id_changed={first_run_id != second_run_id}")

rows = spark.read.json(str(output_dir)).orderBy("event", "count").collect()
print(f"output_rows={len(rows)}")
for row in rows:
    print(f"{row['event']},{row['region']},{row['count']}")

spark.stop()