# gate: skip
"""CityFlow's real retry demo -- flaky_ingest deliberately fails its first
attempt (a marker file it finds missing) and succeeds on the second,
configured with retries=1 / retry_delay so Airflow's own retry machinery,
not a mocked failure, is what gets observed. Module 3 Lesson 4."""
from datetime import datetime, timedelta

from airflow.decorators import dag, task

MARKER_PATH = "/opt/airflow/dags/data/work/retry_marker.txt"


@dag(
    dag_id="cityflow_retry_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-3", "retries"],
)
def cityflow_retry_demo_dag():
    @task(retries=1, retry_delay=timedelta(seconds=20))
    def flaky_ingest():
        """Every real run of this task fails its first attempt on purpose:
        no marker file means "this is attempt 1", so it drops one and
        raises. Airflow's own retry logic waits retry_delay, then tries
        again -- this time the marker exists, so it succeeds and cleans up,
        leaving the DAG ready to demonstrate the exact same failure on its
        next real run."""
        import os

        if not os.path.exists(MARKER_PATH):
            os.makedirs(os.path.dirname(MARKER_PATH), exist_ok=True)
            with open(MARKER_PATH, "w") as f:
                f.write("attempt-1-failed")
            raise RuntimeError(
                "simulated transient failure on attempt 1 -- Airflow should retry"
            )
        os.remove(MARKER_PATH)
        print("flaky_ingest succeeded on retry (attempt 2)")
        return "ok"

    @task
    def confirm_success(result: str) -> None:
        print(f"confirm_success received: {result!r} -- the retried task's real return value")

    confirm_success(flaky_ingest())


cityflow_retry_demo_dag()

import subprocess
import time
import json

dag_id = "cityflow_retry_demo"

deadline = time.time() + 320
found = False
while time.time() < deadline:
    result = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "details", dag_id],
        capture_output=True, text=True,
    )
    if result.returncode == 0:
        found = True
        break
    time.sleep(5)
assert found, f"{dag_id} was never picked up by the scheduler"

subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
    capture_output=True, text=True, check=True,
)
run_id = f"lesson4_retry_verify_{int(time.time())}"
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}")

# retry_delay=20s means this run genuinely takes >20s to reach a terminal state
deadline = time.time() + 90
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
    time.sleep(3)
    result = 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(result.stdout)
    state = next(r["state"] for r in runs if r["run_id"] == run_id)

print(f"run {run_id!r} reached state: {state}")
assert state == "success", f"expected success, got {state}"

states = 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(states.stdout)

import requests

resp = requests.get(
    "http://localhost:8080/api/v1/dags/cityflow_retry_demo/dagRuns/"
    "lesson4_retry_verify_1784885667/taskInstances/flaky_ingest",
    auth=("airflow", "airflow"),
)
info = resp.json()
print("state:", info["state"])
print("try_number:", info["try_number"])
assert info["try_number"] == 2, f"expected try_number 2, got {info['try_number']}"
print("confirmed: flaky_ingest needed 2 tries to succeed")

# gate: skip
"""CityFlow's real catchup/backfill demo -- a small daily-scheduled DAG
(catchup=False) that snapshots the real zone count for whatever logical
date it runs under. Module 3 Lesson 4 unpauses it to show catchup=False
creates no historical runs on its own, then uses `airflow dags backfill`
to create real, distinct historical runs on purpose."""
from datetime import datetime

from airflow.decorators import dag, task
from airflow.operators.python import get_current_context

SNAPSHOT_DIR = "/opt/airflow/dags/data/work/snapshots"


@dag(
    dag_id="cityflow_daily_zone_snapshot",
    schedule="@daily",
    start_date=datetime(2026, 7, 20),
    catchup=False,
    tags=["cityflow", "module-3", "catchup", "backfill"],
)
def cityflow_daily_zone_snapshot_dag():
    @task
    def snapshot_zone_count():
        import json
        import os

        import pandas as pd

        context = get_current_context()
        ds = context["ds"]

        zones = pd.read_csv("/opt/airflow/dags/data/taxi_zone_lookup.csv")
        os.makedirs(SNAPSHOT_DIR, exist_ok=True)
        out_path = f"{SNAPSHOT_DIR}/{ds}.json"
        with open(out_path, "w") as f:
            json.dump({"logical_date": ds, "total_zones": len(zones)}, f, indent=2)
        print(f"snapshot for {ds}: {len(zones)} zones -> {out_path}")

    snapshot_zone_count()


cityflow_daily_zone_snapshot_dag()

import subprocess
import time
import json

dag_id = "cityflow_daily_zone_snapshot"

deadline = time.time() + 320
found = False
while time.time() < deadline:
    result = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "details", dag_id],
        capture_output=True, text=True,
    )
    if result.returncode == 0:
        found = True
        break
    time.sleep(5)
assert found, f"{dag_id} was never picked up by the scheduler"

subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
    capture_output=True, text=True, check=True,
)
time.sleep(15)  # let the scheduler act on the unpause

result = 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_after_unpause = json.loads(result.stdout)
print(f"runs immediately after unpause (catchup=False): {len(runs_after_unpause)}")
for run in runs_after_unpause:
    print(" ", run["execution_date"], run["state"])

backfill = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1",
     "airflow", "dags", "backfill",
     "-s", "2026-07-20", "-e", "2026-07-22", "--reset-dagruns", "-y", dag_id],
    capture_output=True, text=True,
)
print("backfill exit code:", backfill.returncode)
print(backfill.stdout[-600:])
assert backfill.returncode == 0

result = 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_after_backfill = json.loads(result.stdout)
print(f"\ntotal runs after backfill: {len(runs_after_backfill)}")
for run in sorted(runs_after_backfill, key=lambda r: r["execution_date"]):
    print(" ", run["execution_date"], run["run_id"], run["state"])