import json
from pathlib import Path

from pyspark.sql import SparkSession


BASE = Path("/tmp/sg-spark-streaming-checkpoint")
INPUT = BASE / "input"
OUTPUT = BASE / "output"
CHECKPOINT = BASE / "checkpoint"

if BASE.exists():
    raise RuntimeError(f"Use an empty test path: {BASE}")

INPUT.mkdir(parents=True)


def write_events(filename, rows):
    staging_file = BASE / f".{filename}.tmp"
    with staging_file.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, sort_keys=True) + "\n")
    staging_file.replace(INPUT / filename)


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

spark = (
    SparkSession.builder.appName("streaming-checkpoint-check")
    .config("spark.sql.shuffle.partitions", "2")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")


def start_query():
    events = spark.readStream.schema(
        "event STRING, region STRING, count INT"
    ).json(str(INPUT))
    return (
        events.writeStream.format("json")
        .option("path", str(OUTPUT))
        .option("checkpointLocation", str(CHECKPOINT))
        .trigger(availableNow=True)
        .start()
    )


first = start_query()
first.awaitTermination(60)
first_id = first.id
first_run_id = first.runId
first_input_rows = sum(item["numInputRows"] for item in first.recentProgress)

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

second = start_query()
second.awaitTermination(60)
second_input_rows = sum(item["numInputRows"] for item in second.recentProgress)

query_id_reused = first_id == second.id
run_id_changed = first_run_id != second.runId
rows = spark.read.json(str(OUTPUT)).orderBy("event", "count").collect()

if not query_id_reused or not run_id_changed:
    raise RuntimeError("The second query did not recover from the checkpoint")
if second_input_rows != 1 or len(rows) != 3:
    raise RuntimeError("The restarted query did not process only the new file")

checkpoint_entries = sorted(
    path.name for path in CHECKPOINT.iterdir() if not path.name.startswith(".")
)
print(f"first_input_rows={first_input_rows}")
print(f"second_input_rows={second_input_rows}")
print("checkpoint_entries=" + ",".join(checkpoint_entries))
print(f"query_id_reused={query_id_reused}")
print(f"run_id_changed={run_id_changed}")
print(f"output_rows={len(rows)}")
for row in rows:
    print(f"{row['event']},{row['region']},{row['count']}")

spark.stop()