# gate: skip
"""CityFlow's real branching DAG -- ingest/validate feed two independent
summary tasks in parallel, which both feed a merge step before the final
report. Deliberately wires its five edges with all four dependency
mechanisms Module 3 Lesson 3 covers: `set_downstream`, `>>` (including to a
list, the actual branch), `set_upstream`, and `<<`."""
from datetime import datetime

from airflow import DAG
from airflow.operators.python import PythonOperator

WORKDIR = "/opt/airflow/dags/data/work/m3l3"
RAW_CSV = "/opt/airflow/dags/data/yellow_tripdata_sample.csv"


def ingest_taxi_data(**context):
    import os
    import pandas as pd

    os.makedirs(WORKDIR, exist_ok=True)
    raw = pd.read_csv(RAW_CSV)
    out_path = f"{WORKDIR}/ingested.csv"
    raw.to_csv(out_path, index=False)
    print(f"ingested {len(raw):,} raw trips -> {out_path}")
    return out_path


def validate_trips(**context):
    import pandas as pd

    ti = context["ti"]
    ingested_path = ti.xcom_pull(task_ids="ingest_taxi_data")
    raw = pd.read_csv(ingested_path)
    valid = raw[
        (raw["passenger_count"] > 0)
        & (raw["trip_distance"] > 0)
        & (raw["fare_amount"] > 0)
    ]
    out_path = f"{WORKDIR}/validated.csv"
    valid.to_csv(out_path, index=False)
    print(f"kept {len(valid):,} of {len(raw):,} rows -> {out_path}")
    return out_path


def compute_by_payment_type(**context):
    import pandas as pd

    ti = context["ti"]
    validated_path = ti.xcom_pull(task_ids="validate_trips")
    valid = pd.read_csv(validated_path)
    summary = (
        valid.groupby("payment_type")
        .agg(trips=("fare_amount", "size"), avg_fare=("fare_amount", "mean"))
        .round(2)
        .reset_index()
    )
    print(summary.to_string(index=False))
    return summary.to_dict(orient="records")


def compute_by_passenger_count(**context):
    import pandas as pd

    ti = context["ti"]
    validated_path = ti.xcom_pull(task_ids="validate_trips")
    valid = pd.read_csv(validated_path)
    summary = (
        valid.groupby("passenger_count")
        .agg(trips=("fare_amount", "size"), avg_distance=("trip_distance", "mean"))
        .round(2)
        .reset_index()
    )
    print(summary.to_string(index=False))
    return summary.to_dict(orient="records")


def merge_summaries(**context):
    ti = context["ti"]
    by_payment = ti.xcom_pull(task_ids="compute_by_payment_type")
    by_passengers = ti.xcom_pull(task_ids="compute_by_passenger_count")
    return {"by_payment_type": by_payment, "by_passenger_count": by_passengers}


def load_report(**context):
    import json

    ti = context["ti"]
    merged = ti.xcom_pull(task_ids="merge_summaries")
    out_path = f"{WORKDIR}/branch_report.json"
    with open(out_path, "w") as f:
        json.dump(merged, f, indent=2)
    print(f"wrote {out_path}")
    print(json.dumps(merged, indent=2))


with DAG(
    dag_id="cityflow_branching_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-3", "branching"],
) as dag:
    ingest = PythonOperator(task_id="ingest_taxi_data", python_callable=ingest_taxi_data)
    validate = PythonOperator(task_id="validate_trips", python_callable=validate_trips)
    by_payment = PythonOperator(task_id="compute_by_payment_type", python_callable=compute_by_payment_type)
    by_passengers = PythonOperator(task_id="compute_by_passenger_count", python_callable=compute_by_passenger_count)
    merge = PythonOperator(task_id="merge_summaries", python_callable=merge_summaries)
    report = PythonOperator(task_id="load_report", python_callable=load_report)

    # 1. set_downstream(): the linear start of the chain.
    ingest.set_downstream(validate)
    # 2. `>>` to a list: this is the actual branch -- both tasks depend on
    #    validate, neither depends on the other, so Airflow can run them
    #    in parallel.
    validate >> [by_payment, by_passengers]
    # 3. set_upstream() with a list: the branch closes back up here -- merge
    #    only runs once BOTH parallel tasks have succeeded.
    merge.set_upstream([by_payment, by_passengers])
    # 4. `<<`: the same relationship as `>>`, written with the arrow reversed.
    report << merge

import subprocess
import time
import json

dag_id = "cityflow_branching_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"lesson3_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}")

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)

{
  "by_payment_type": [
    {"payment_type": 1, "trips": 153384, "avg_fare": 18.39},
    {"payment_type": 2, "trips": 28056,  "avg_fare": 18.46},
    {"payment_type": 3, "trips": 676,    "avg_fare": 16.55},
    {"payment_type": 4, "trips": 1517,   "avg_fare": 19.27}
  ],
  "by_passenger_count": [
    {"passenger_count": 1.0, "trips": 143638, "avg_distance": 3.15},
    {"passenger_count": 2.0, "trips": 27038,  "avg_distance": 3.84},
    {"passenger_count": 3.0, "trips": 5912,   "avg_distance": 3.67},
    {"passenger_count": 4.0, "trips": 3304,   "avg_distance": 3.99},
    {"passenger_count": 5.0, "trips": 2263,   "avg_distance": 3.00},
    {"passenger_count": 6.0, "trips": 1475,   "avg_distance": 2.87},
    {"passenger_count": 8.0, "trips": 3,      "avg_distance": 0.05}
  ]
}