"""Spark4.0.1 real file stream; bounded demo, not external-sink EOS."""
import argparse
import json
import hashlib
import os
from pathlib import Path
from pyspark.sql import SparkSession,functions as F,types as T

def run(input_dir,checkpoint,output_dir,generation):
    checkpoint=Path(checkpoint).resolve();output_dir=Path(output_dir).resolve()
    output_dir.mkdir(parents=True,exist_ok=True)
    identity=output_dir/"lineage.json"
    contract={"generation":generation,"checkpoint":str(checkpoint),
              "input_uri":Path(input_dir).resolve().as_uri(),
              "code_sha256":hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
    if identity.exists() and json.loads(identity.read_text())!=contract:
        raise ValueError("Use a new output directory for a new checkpoint lineage")
    identity.write_text(json.dumps(contract),encoding="utf-8")
    spark=SparkSession.builder.master("local[2]").appName("handbook-stream").config(
      "spark.sql.session.timeZone","UTC").config("spark.sql.shuffle.partitions","2").getOrCreate()
    schema=T.StructType([
      T.StructField("event_id",T.StringType()),T.StructField("tenant_id",T.StringType()),
      T.StructField("currency",T.StringType()),T.StructField("amount_minor",T.LongType()),
      T.StructField("occurred_at",T.StringType()),T.StructField("entity_version",T.LongType()),
      T.StructField("schema_version",T.IntegerType())])
    try:
        source=spark.readStream.schema(schema).json(Path(input_dir).resolve().as_uri())
        timed=source.withColumn("event_time",F.to_timestamp("occurred_at"))
        aggregates=timed.withWatermark("event_time","5 seconds").groupBy(
          "tenant_id","currency",F.window("event_time","10 seconds")).agg(
          F.sum("amount_minor").alias("amount_minor"))
        def publish(frame,batch_id):
            rows=frame.orderBy("tenant_id","currency","window.start").limit(10001).collect()
            if len(rows)>10000:raise RuntimeError("bounded demo output exceeded")
            values=[{"tenant_id":r.tenant_id,"currency":r.currency,
                     "window_start":r.window.start.isoformat(),
                     "window_end":r.window.end.isoformat(),
                     "amount_minor":r.amount_minor} for r in rows]
            payload={"generation":generation,"batch_id":batch_id,"rows":values}
            temporary=output_dir/"snapshot.json.tmp"
            temporary.write_text(json.dumps(payload,sort_keys=True),encoding="utf-8")
            os.replace(temporary,output_dir/"snapshot.json")
        query=aggregates.writeStream.outputMode("complete").foreachBatch(publish).option(
          "checkpointLocation",checkpoint.as_uri()).trigger(availableNow=True).start()
        query.awaitTermination(120)
        if query.isActive:
            query.stop()
            raise TimeoutError("query did not finish bounded input")
        if query.exception():raise RuntimeError(str(query.exception()))
        return json.loads((output_dir/"snapshot.json").read_text(encoding="utf-8"))
    finally:spark.stop()

if __name__=="__main__":
    ap=argparse.ArgumentParser()
    ap.add_argument("--input",type=Path,required=True)
    ap.add_argument("--checkpoint",type=Path,required=True)
    ap.add_argument("--output",type=Path,required=True)
    ap.add_argument("--generation",required=True)
    a=ap.parse_args()
    print(json.dumps(run(a.input,a.checkpoint,a.output,a.generation),indent=2))
# COMPLETE mode retains all aggregate state: suitable only for this bounded
# fixture. Production needs bounded state/correction and real distributed sink.
# One writer per output/checkpoint; local os.replace is NOT S3 multi-object commit.