# gate: skip
"""CityFlow's real Airflow-submits-a-Kubernetes-Job demo (Module 7 Lesson 2).

No apache-airflow-providers-cncf-kubernetes is installed in this image (heavy,
and not needed) -- so this DAG's real Kubernetes work happens the honest way
this module explains: a @task that shells out, via subprocess, to the same
real kubectl binary + kubeconfig this course host-mounted at
/opt/airflow/dags/bin/kubectl and /opt/airflow/dags/data/kubeconfig-kind.yaml,
talking to the real datatweets-docker-k8s Kind cluster's API server at
https://host.docker.internal:<port>. That is a legitimate "equivalent" to
KubernetesPodOperator per this module's own spec -- not a shortcut pretending
to be the provider-based operator.

The Job's pod runs a trivial real command (python:3.13-slim printing a line
and exiting 0) on purpose -- this lesson proves the submission mechanism
(build a Job manifest, delete-then-apply for idempotency, poll
.status.succeeded, read real pod logs) before Lesson 3 replaces the trivial
pod with a real PySpark job.
"""
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-k8s-submit-demo"
IMAGE = "python:3.13-slim"
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-hello
          image: {IMAGE}
          imagePullPolicy: IfNotPresent
          command: ["python", "-c", "print('CityFlow: hello from a real Kubernetes Job, submitted by Airflow')"]
"""


@dag(
    dag_id="cityflow_k8s_submit_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-7", "kubernetes"],
)
def cityflow_k8s_submit_demo_dag():
    @task
    def submit_job() -> str:
        """Deletes any prior real Job of this name (idempotent -- a re-run of
        this DAG, or of the code gate, must not fail on 'already exists'),
        then applies the real manifest above via `kubectl apply -f -`."""
        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_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 -- not the Airflow
        task's own state, the real Kubernetes object's -- until it reports 1
        real completed pod, a real failure, or a real timeout."""
        import subprocess
        import time

        deadline = time.time() + 120
        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 report_job_logs(poll_result: dict) -> str:
        """Reads the real pod's real logs -- proof the trivial command
        actually ran on the real Kind cluster, not just that the Job object
        reached a green status."""
        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.strip()
        print(f"report_job_logs: real pod {pod_name!r} logged: {logs!r}")
        return logs

    report_job_logs(poll_job_completion(submit_job()))


cityflow_k8s_submit_demo_dag()

import json
import subprocess
import time

dag_id = "cityflow_k8s_submit_demo"

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

run_id = f"lesson2_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() + 90
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
    time.sleep(2)
    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 same real pod really ran, straight from kubectl -- independent
# of anything Airflow itself 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-k8s-submit-demo"],
    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 and confirm the
# delete-then-apply pattern doesn't error on "already exists".
run_id_2 = f"lesson2_m7_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() + 90
state2 = "queued"
while state2 not in ("success", "failed") and time.time() < deadline:
    time.sleep(2)
    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 Job name a second time is genuinely idempotent")