"""Real, importable on_failure_callback logic for CityFlow DAGs (Module 6
Lesson 4, reused by Lesson 5's guided project).

record_failure_alert() has no Airflow import either -- it takes plain values
(dag id, task id, run id, execution date, exception text) and appends one
real JSON line to a real alert log file. Each cityflow_*_dag.py file that
wants alerting wires up a small on_failure_callback wrapper that pulls those
values out of Airflow's own `context` dict and calls this function -- the
sink itself stays pure and unit-testable without Airflow installed at all,
the same split Lesson 1 drew for validation logic."""
import json
from pathlib import Path


def record_failure_alert(alert_path: str, dag_id: str, task_id: str, run_id: str, execution_date: str, exception: str) -> dict:
    """Appends one real alert record to `alert_path` (creating its parent
    directory and the file itself if either is missing) and returns the
    record that was written, so a caller can log or assert on it directly."""
    Path(alert_path).parent.mkdir(parents=True, exist_ok=True)
    record = {
        "dag_id": dag_id, "task_id": task_id, "run_id": run_id,
        "execution_date": execution_date, "exception": exception,
    }
    with open(alert_path, "a") as f:
        f.write(json.dumps(record) + "\n")
    return record


def read_alerts(alert_path: str) -> list:
    """Reads every alert record back from `alert_path` as a list of dicts --
    an empty list if the file doesn't exist yet (no alert has ever fired)."""
    path = Path(alert_path)
    if not path.exists():
        return []
    return [json.loads(line) for line in path.read_text().splitlines() if line.strip()]

# gate: skip
"""CityFlow's real on_failure_callback demo (Module 6 Lesson 4).

Reuses Lesson 2's exact real-failure shape -- ingest_ok succeeds,
process_incoming_batch genuinely fails reading a batch file that was never
dropped -- and adds one thing: default_args wires cityflow_on_failure_alert
as this DAG's on_failure_callback, so every task's real failure (there's
only ever one candidate here) fires a real call to
cityflow_alerting.record_failure_alert, which appends one real JSON line to
a real local alert file with the real failure's dag_id, task_id, run_id, and
execution date."""
import sys
from datetime import datetime

from airflow.decorators import dag, task

sys.path.insert(0, "/opt/airflow/dags")
import cityflow_alerting  # noqa: E402

ZONE_LOOKUP_PATH = "/opt/airflow/dags/data/taxi_zone_lookup.csv"
MISSING_BATCH_PATH = "/opt/airflow/dags/data/work/alerting_demo/incoming_batch.csv"
ALERT_PATH = "/opt/airflow/dags/data/work/alerting_demo/cityflow_alerts.jsonl"


def cityflow_on_failure_alert(context):
    """The real on_failure_callback Airflow invokes with a real failure's own
    context -- pulls the real identity out of it and hands off to the pure,
    already-tested sink in cityflow_alerting.py."""
    ti = context["task_instance"]
    record = cityflow_alerting.record_failure_alert(
        alert_path=ALERT_PATH,
        dag_id=ti.dag_id,
        task_id=ti.task_id,
        run_id=context["run_id"],
        execution_date=str(context["logical_date"]),
        exception=repr(context.get("exception")),
    )
    print(f"cityflow_on_failure_alert: real alert recorded -> {record}")


@dag(
    dag_id="cityflow_alerting_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    default_args={"on_failure_callback": cityflow_on_failure_alert, "retries": 0},
    tags=["cityflow", "module-6", "alerting"],
)
def cityflow_alerting_demo_dag():
    @task
    def ingest_ok() -> str:
        import pandas as pd

        zones = pd.read_csv(ZONE_LOOKUP_PATH)
        print(f"ingest_ok: read {len(zones)} real CityFlow zone rows")
        return "zones_loaded"

    @task
    def process_incoming_batch(zones_ref: str) -> None:
        import pandas as pd

        print(f"process_incoming_batch: zones_ref={zones_ref!r}, reading {MISSING_BATCH_PATH}")
        pd.read_csv(MISSING_BATCH_PATH)  # genuinely missing -> real FileNotFoundError -> real on_failure_callback fires

    process_incoming_batch(ingest_ok())


cityflow_alerting_demo_dag()

import json
import os
import subprocess
import time

dag_id = "cityflow_alerting_demo"
alert_host_path = ".airflow-local/dags/data/work/alerting_demo/cityflow_alerts.jsonl"

if os.path.exists(alert_host_path):
    os.remove(alert_host_path)
    print(f"removed pre-existing {alert_host_path} to start from a clean slate")

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_m6_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 == "failed", f"expected failed, 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)

time.sleep(1)
with open(alert_host_path) as f:
    lines = [json.loads(l) for l in f if l.strip()]
print(f"real alert file now has {len(lines)} line(s):")
for l in lines:
    print(" ", l)
assert len(lines) == 1, f"expected exactly 1 alert record, got {len(lines)}"
assert lines[0]["dag_id"] == dag_id
assert lines[0]["task_id"] == "process_incoming_batch"
assert lines[0]["run_id"] == run_id