# gate: skip
"""CityFlow's real Airflow + Spark + Kubernetes guided project (Module 7
Lesson 5).

Lessons 2-4 combined into one clean, complete DAG: submit a real PySpark Job
to the real datatweets-docker-k8s Kind cluster from a real Airflow task,
poll the real Job's status to completion, and read back the real Spark
result from the real pod's logs -- the exact shape Lesson 3 proved, run
here as this module's capstone rather than a lesson demo. Still no
apache-airflow-providers-cncf-kubernetes -- every Kubernetes call is a real
subprocess to the same real kubectl binary + kubeconfig this course
host-mounted into the Airflow containers, talking to Kind's real API server
at https://host.docker.internal:<port> (insecure-skip-tls-verify: true,
because Kind's certificate has no host.docker.internal SAN -- a real,
stated local-dev trade-off, not a hidden one)."""
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-airflow-spark-k8s-guided-project"
IMAGE = "cityflow-pyspark-job: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-pyspark
          image: {IMAGE}
          imagePullPolicy: IfNotPresent
"""


@dag(
    dag_id="cityflow_airflow_spark_k8s_guided_project",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-7", "guided-project", "kubernetes", "pyspark"],
)
def cityflow_airflow_spark_k8s_guided_project_dag():
    @task
    def submit_pyspark_job() -> str:
        """Delete-then-apply, the same idempotent pattern Lesson 2 proved and
        Lesson 3 reused -- a second run of this exact DAG (or a second run of
        the code gate, proving idempotency) 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_pyspark_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:
        import subprocess
        import time

        deadline = time.time() + 180
        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(2)
        raise TimeoutError(f"real Job {job_name} did not succeed within the poll deadline")

    @task
    def verify_spark_result(poll_result: dict) -> dict:
        """Reads the real pod's real logs, parses the real
        CITYFLOW_PYSPARK_RESULT_JSON: line the Spark job printed, and checks
        it against the exact real numbers this module already established in
        Lesson 3 -- proof this guided project's end-to-end run reproduces
        the same real computation, not a different or approximate one."""
        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_PYSPARK_RESULT_JSON:(\{.*\})", logs)
        if not match:
            raise RuntimeError(f"no CITYFLOW_PYSPARK_RESULT_JSON line found in real pod logs:\n{logs}")
        result = json.loads(match.group(1))
        print(f"verify_spark_result: real pod {pod_name!r} computed: {result}")

        assert result["total_rows"] == 81013, result["total_rows"]
        assert result["top_borough"] == "Manhattan", result["top_borough"]
        assert result["top_borough_trip_count"] == 68416, result["top_borough_trip_count"]
        by_borough = {r["Borough"]: r for r in result["by_borough"]}
        assert by_borough["Manhattan"]["total_revenue"] == 1673556.42
        assert sum(r["trip_count"] for r in result["by_borough"]) == 81013
        print("verify_spark_result: every real number matches Lesson 3's already-verified result")
        return {"pod_name": pod_name, **result}

    verify_spark_result(poll_job_completion(submit_pyspark_job()))


cityflow_airflow_spark_k8s_guided_project_dag()

import json
import subprocess
import time

dag_id = "cityflow_airflow_spark_k8s_guided_project"

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

run_id = f"lesson5_m7_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() + 180
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)

# Confirm the real pod independently via kubectl, outside anything Airflow reported.
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-airflow-spark-k8s-guided-project"],
    capture_output=True, text=True, check=True,
)
print(pods.stdout)
assert "Completed" in pods.stdout