import time

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

FILES = [
    "yellow_tripdata_2024-01.parquet",
    "yellow_tripdata_2024-02.parquet",
    "yellow_tripdata_2024-03.parquet",
]
KEEP = ["tpep_pickup_datetime", "PULocationID", "total_amount"]
WIN_LO = "2024-01-01"
WIN_HI = "2024-04-01"

start = time.time()

spark = (
    SparkSession.builder
    .appName("cityflow-capstone-local-proof")
    .master("local[*]")
    .config("spark.ui.showConsoleProgress", "false")
    .getOrCreate()
)
spark.sparkContext.setLogLevel("ERROR")

trips = spark.read.parquet(*FILES).select(*KEEP)
total_rows = trips.count()

in_window = trips.filter(
    (F.col("tpep_pickup_datetime") >= F.lit(WIN_LO))
    & (F.col("tpep_pickup_datetime") < F.lit(WIN_HI))
)
in_window_rows = in_window.count()
quarantined = total_rows - in_window_rows

bucketed = in_window.withColumn("hour", F.date_trunc("hour", F.col("tpep_pickup_datetime")))
report = (
    bucketed.groupBy(F.col("PULocationID").alias("zone"), "hour")
    .agg(
        F.count("*").alias("trips"),
        F.sum("total_amount").alias("revenue"),
    )
)
report.cache()

buckets = report.count()
trips_kept = report.agg(F.sum("trips")).first()[0]
total_revenue = report.agg(F.sum("revenue")).first()[0]
top = report.orderBy(F.col("trips").desc()).first()
hour_str = top["hour"].strftime("%Y-%m-%dT%H")

elapsed = time.time() - start

print(f"source rows across 3 files: {total_rows:,}")
print(f"zone-hour buckets:          {buckets:,}")
print(f"trips kept:                 {trips_kept:,}")
print(f"quarantined (outside qtr):  {quarantined}")
print(f"total revenue:              ${total_revenue:,.2f}")
print(f"busiest zone-hour:          zone {top['zone']} at {hour_str}, {top['trips']} trips")
print(f"elapsed seconds:            {elapsed:.2f}")

# Real assertions against the anchor every course in this path has verified --
# this is the same number, reproduced a fourth independent way.
assert total_rows == 9_554_778, total_rows
assert buckets == 240_917, buckets
assert trips_kept == 9_554_757, trips_kept
assert quarantined == 21, quarantined
assert round(total_revenue, 2) == 256_692_373.14, round(total_revenue, 2)
assert top["zone"] == 79, top["zone"]
assert hour_str == "2024-02-25T01", hour_str
assert top["trips"] == 846, top["trips"]
print("PROVED: PySpark local-mode aggregation reproduces the exact anchor:")
print("240,917 buckets, 9,554,757 trips, $256,692,373.14, zone 79 @ 2024-02-25T01, 846 trips")

spark.stop()

FROM python:3.13-slim

# Spark 4.x runs on the JVM and needs Java 17 or 21 -- a headless JRE is
# enough (no compiler, no javadoc, no javafx). Same recipe Module 7 Lesson 3
# used for cityflow-pyspark-job:1.0: python:3.13-slim's current Debian base
# (trixie) only ships JDK 21 in its repos, not a standalone
# openjdk-17-jre-headless package.
RUN apt-get update \
    && apt-get install --no-install-recommends -y openjdk-21-jre-headless \
    && rm -rf /var/lib/apt/lists/*
ENV JAVA_HOME=/usr/lib/jvm/java-21-openjdk-arm64

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY cityflow_capstone_job.py .
COPY yellow_tripdata_2024-01.parquet /opt/cityflow/data/yellow_tripdata_2024-01.parquet
COPY yellow_tripdata_2024-02.parquet /opt/cityflow/data/yellow_tripdata_2024-02.parquet
COPY yellow_tripdata_2024-03.parquet /opt/cityflow/data/yellow_tripdata_2024-03.parquet

CMD ["python", "cityflow_capstone_job.py"]

# gate: skip
"""CityFlow's real quarterly zone-hour ETL, as a real PySpark-on-Kubernetes
job (Module 8 - Capstone). Runs inside the cityflow-capstone-pyspark:1.0
image, on a real Kind Kubernetes Job pod, submitted by a real Airflow DAG --
the exact submit/poll/read-logs mechanism Module 7 built, with the image
carrying the full real Jan-Mar 2024 quarter instead of one day's slice."""
import json
import time

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

FILES = [
    "/opt/cityflow/data/yellow_tripdata_2024-01.parquet",
    "/opt/cityflow/data/yellow_tripdata_2024-02.parquet",
    "/opt/cityflow/data/yellow_tripdata_2024-03.parquet",
]
KEEP = ["tpep_pickup_datetime", "PULocationID", "total_amount"]
WIN_LO = "2024-01-01"
WIN_HI = "2024-04-01"


def main():
    start = time.time()
    spark = (
        SparkSession.builder
        .appName("cityflow-capstone-quarterly-etl")
        .master("local[*]")
        .config("spark.ui.showConsoleProgress", "false")
        .getOrCreate()
    )
    spark.sparkContext.setLogLevel("ERROR")

    trips = spark.read.parquet(*FILES).select(*KEEP)
    total_rows = trips.count()

    in_window = trips.filter(
        (F.col("tpep_pickup_datetime") >= F.lit(WIN_LO))
        & (F.col("tpep_pickup_datetime") < F.lit(WIN_HI))
    )
    quarantined = total_rows - in_window.count()

    bucketed = in_window.withColumn("hour", F.date_trunc("hour", F.col("tpep_pickup_datetime")))
    report = (
        bucketed.groupBy(F.col("PULocationID").alias("zone"), "hour")
        .agg(F.count("*").alias("trips"), F.sum("total_amount").alias("revenue"))
    )
    report.cache()

    buckets = report.count()
    trips_kept = report.agg(F.sum("trips")).first()[0]
    total_revenue = round(float(report.agg(F.sum("revenue")).first()[0]), 2)
    top = report.orderBy(F.col("trips").desc()).first()
    hour_str = top["hour"].strftime("%Y-%m-%dT%H")
    elapsed = round(time.time() - start, 2)

    result = {
        "source_files": [f.rsplit("/", 1)[-1] for f in FILES],
        "window_start": WIN_LO, "window_end": WIN_HI,
        "source_rows": total_rows, "zone_hour_buckets": buckets,
        "trips_kept": int(trips_kept), "quarantined": int(quarantined),
        "total_revenue": total_revenue, "busiest_zone": int(top["zone"]),
        "busiest_hour": hour_str, "busiest_trips": int(top["trips"]),
        "elapsed_seconds": elapsed,
    }

    print("CityFlow capstone quarterly zone-hour report -- real PySpark job on Kubernetes:")
    print(f"  source rows (3 real months):  {total_rows:,}")
    print(f"  zone-hour buckets:            {buckets:,}")
    print(f"  trips kept:                   {int(trips_kept):,}")
    print(f"  quarantined (outside qtr):    {int(quarantined)}")
    print(f"  total revenue:                ${total_revenue:,.2f}")
    print(f"  busiest zone-hour:            zone {int(top['zone'])} at {hour_str}, {int(top['trips'])} trips")
    print(f"  elapsed seconds:              {elapsed}")
    print("CITYFLOW_CAPSTONE_RESULT_JSON:" + json.dumps(result))

    spark.stop()


if __name__ == "__main__":
    main()

# gate: skip
"""CityFlow's real Airflow + PySpark + Kubernetes capstone DAG (Module 8 -
Capstone, the final lesson of the final course of the 4-course Data
Engineering path).

Same submit -> poll -> read-and-verify-logs shape Module 7 built and proved:
still no apache-airflow-providers-cncf-kubernetes -- every Kubernetes call is
a real subprocess to the real kubectl binary + kubeconfig this course
host-mounted into the Airflow containers, talking to the real
datatweets-docker-k8s Kind cluster's real API server at
https://host.docker.internal:<port> (insecure-skip-tls-verify: true, the
same real, stated local-dev trade-off Module 7 Lesson 1 explained).

What's new here is the payload, not the mechanism: the Job's pod runs
cityflow-capstone-pyspark:1.0 -- a real PySpark image with all three real
Jan-Mar 2024 monthly parquet files baked in (9,554,778 real trips) -- and
verify_capstone_result asserts the full anchor this entire 4-course path has
verified three times already: 240,917 zone-hour buckets, 9,554,757 trips,
$256,692,373.14 in revenue, busiest zone-hour zone 79 at 2024-02-25T01 with
846 trips. This DAG's own green run is the FOURTH independent reproduction
of that exact result."""
import re
from datetime import datetime

from airflow.decorators import dag, task

KUBECTL_BIN = "/opt/airflow/dags/bin/kubectl"
KUBECONFIG_PATH = "/opt/airflow/dags/data/kubeconfig-kind.yaml"
NAMESPACE = "default"
JOB_NAME = "cityflow-capstone-quarterly-report"
IMAGE = "cityflow-capstone-pyspark:1.0"
JOB_MANIFEST = f"""\
apiVersion: batch/v1
kind: Job
metadata:
  name: {JOB_NAME}
  namespace: {NAMESPACE}
  labels:
    app: {JOB_NAME}
spec:
  backoffLimit: 0
  ttlSecondsAfterFinished: 3600
  template:
    metadata:
      labels:
        app: {JOB_NAME}
    spec:
      restartPolicy: Never
      containers:
        - name: cityflow-capstone
          image: {IMAGE}
          imagePullPolicy: IfNotPresent
          resources:
            requests:
              cpu: "1"
              memory: "1Gi"
            limits:
              cpu: "2"
              memory: "2Gi"
"""


@dag(
    dag_id="cityflow_capstone",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-8", "capstone", "kubernetes", "pyspark"],
)
def cityflow_capstone_dag():
    @task
    def submit_capstone_job() -> str:
        """Delete-then-apply, the same idempotent pattern Module 7 proved --
        a second run of this exact DAG (or the code gate) creates the same
        real Job again cleanly, never fails on 'already exists'."""
        import subprocess

        subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "delete", "job", JOB_NAME,
             "-n", NAMESPACE, "--ignore-not-found"],
            capture_output=True, text=True, check=True,
        )
        result = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "apply", "-f", "-"],
            input=JOB_MANIFEST, capture_output=True, text=True,
        )
        print(f"submit_capstone_job: kubectl apply -> {result.stdout.strip()}")
        if result.returncode != 0:
            raise RuntimeError(f"kubectl apply failed: {result.stderr}")
        return JOB_NAME

    @task
    def poll_job_completion(job_name: str) -> dict:
        """Polls the real Job's .status.succeeded field. A real Spark job
        over the full real quarter needs more headroom than Module 7's
        single-day slice, hence the longer deadline."""
        import subprocess
        import time

        deadline = time.time() + 300
        polls = 0
        while time.time() < deadline:
            polls += 1
            succeeded = subprocess.run(
                [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "job", job_name,
                 "-n", NAMESPACE, "-o", "jsonpath={.status.succeeded}"],
                capture_output=True, text=True, check=True,
            ).stdout.strip()
            failed = subprocess.run(
                [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "job", job_name,
                 "-n", NAMESPACE, "-o", "jsonpath={.status.failed}"],
                capture_output=True, text=True, check=True,
            ).stdout.strip()
            print(f"poll_job_completion: poll #{polls} succeeded={succeeded!r} failed={failed!r}")
            if succeeded == "1":
                return {"job_name": job_name, "polls": polls, "outcome": "succeeded"}
            if failed and int(failed) > 0:
                raise RuntimeError(f"real Job {job_name} reported {failed} failed pod(s)")
            time.sleep(3)
        raise TimeoutError(f"real Job {job_name} did not succeed within the poll deadline")

    @task
    def verify_capstone_result(poll_result: dict) -> dict:
        """Reads the real pod's real logs, parses CITYFLOW_CAPSTONE_RESULT_JSON:,
        and asserts it against the exact anchor this whole 4-course path has
        already verified three separate times -- proof of a fourth
        independent reproduction, not just that a Spark job ran."""
        import json
        import subprocess

        job_name = poll_result["job_name"]
        pod_name = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "get", "pods", "-n", NAMESPACE,
             "-l", f"app={job_name}", "-o", "jsonpath={.items[0].metadata.name}"],
            capture_output=True, text=True, check=True,
        ).stdout.strip()
        logs = subprocess.run(
            [KUBECTL_BIN, "--kubeconfig", KUBECONFIG_PATH, "logs", pod_name, "-n", NAMESPACE],
            capture_output=True, text=True, check=True,
        ).stdout

        match = re.search(r"CITYFLOW_CAPSTONE_RESULT_JSON:(\{.*\})", logs)
        if not match:
            raise RuntimeError(f"no CITYFLOW_CAPSTONE_RESULT_JSON line found in real pod logs:\n{logs}")
        result = json.loads(match.group(1))
        print(f"verify_capstone_result: real pod {pod_name!r} computed: {result}")

        assert result["source_rows"] == 9_554_778, result["source_rows"]
        assert result["zone_hour_buckets"] == 240_917, result["zone_hour_buckets"]
        assert result["trips_kept"] == 9_554_757, result["trips_kept"]
        assert result["quarantined"] == 21, result["quarantined"]
        assert result["total_revenue"] == 256_692_373.14, result["total_revenue"]
        assert result["busiest_zone"] == 79, result["busiest_zone"]
        assert result["busiest_hour"] == "2024-02-25T01", result["busiest_hour"]
        assert result["busiest_trips"] == 846, result["busiest_trips"]
        print("verify_capstone_result: every real number matches the whole path's own anchor -- "
              "240,917 buckets, 9,554,757 trips, $256,692,373.14, zone 79 @ 2024-02-25T01, 846 trips")
        return {"pod_name": pod_name, **result}

    verify_capstone_result(poll_job_completion(submit_capstone_job()))


cityflow_capstone_dag()

import json
import subprocess
import time

dag_id = "cityflow_capstone"

subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
                capture_output=True, text=True, check=True)

run_id = f"capstone_verify_{int(time.time())}"
trigger = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
    capture_output=True, text=True,
)
if trigger.returncode != 0:
    subprocess.run(["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "reserialize"],
                    capture_output=True, text=True, check=True)
    subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id],
        capture_output=True, text=True, check=True,
    )
print(f"triggered {dag_id!r} as {run_id!r}")

deadline = time.time() + 300
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
    time.sleep(3)
    r = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
        capture_output=True, text=True, check=True,
    )
    runs = json.loads(r.stdout)
    match = [x for x in runs if x["run_id"] == run_id]
    state = match[0]["state"] if match else "queued"
print(f"run {run_id!r} reached state: {state}")
assert state == "success", f"expected success, got {state}"

r = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "tasks", "states-for-dag-run", dag_id, run_id],
    capture_output=True, text=True, check=True,
)
print(r.stdout)

pods = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "/opt/airflow/dags/bin/kubectl",
     "--kubeconfig", "/opt/airflow/dags/data/kubeconfig-kind.yaml", "get", "pods", "-n", "default",
     "-l", "app=cityflow-capstone-quarterly-report"],
    capture_output=True, text=True, check=True,
)
print(pods.stdout)
assert "Completed" in pods.stdout

# Real idempotency check: submit the same Job a second time.
run_id_2 = f"capstone_verify_idem_{int(time.time())}"
subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id_2],
    capture_output=True, text=True, check=True,
)
deadline = time.time() + 300
state2 = "queued"
while state2 not in ("success", "failed") and time.time() < deadline:
    time.sleep(3)
    r = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
        capture_output=True, text=True, check=True,
    )
    runs = json.loads(r.stdout)
    match = [x for x in runs if x["run_id"] == run_id_2]
    state2 = match[0]["state"] if match else "queued"
print(f"second run (idempotency check) {run_id_2!r} reached state: {state2}")
assert state2 == "success", f"expected success on re-submission, got {state2}"
print("PROVED: submitting the same capstone Job a second time is genuinely idempotent")