# gate: skip
"""Module 3's guided project: one real DAG combining genuine branching
dependencies (ingest/validate feeding two independent summaries that merge
back together) with a deliberately retried final task (load_report fails
its first attempt, succeeds on the second) -- the two ideas Lessons 3 and 4
each demonstrated alone, now together in a single run."""
from datetime import datetime, timedelta

from airflow.decorators import dag, task

WORKDIR = "/opt/airflow/dags/data/work/m3_guided_project"
RAW_CSV = "/opt/airflow/dags/data/yellow_tripdata_sample.csv"
LOAD_MARKER = f"{WORKDIR}/load_report_marker.txt"


@dag(
    dag_id="cityflow_guided_project_authoring",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-3", "guided-project"],
)
def cityflow_guided_project_authoring_dag():
    @task
    def ingest_taxi_data() -> str:
        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

    @task
    def validate_trips(ingested_path: str) -> str:
        import pandas as pd

        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

    @task
    def compute_by_payment_type(validated_path: str) -> list:
        import pandas as pd

        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")

    @task
    def compute_by_passenger_count(validated_path: str) -> list:
        import pandas as pd

        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")

    @task
    def merge_summaries(by_payment: list, by_passengers: list) -> dict:
        return {"by_payment_type": by_payment, "by_passenger_count": by_passengers}

    @task(retries=1, retry_delay=timedelta(seconds=20))
    def load_report(merged: dict) -> None:
        """The deliberately-retried task: fails its first real attempt (no
        marker file present yet), succeeds on the second."""
        import json
        import os

        if not os.path.exists(LOAD_MARKER):
            with open(LOAD_MARKER, "w") as f:
                f.write("attempt-1-failed")
            raise RuntimeError(
                "simulated transient failure writing the report -- Airflow should retry"
            )
        os.remove(LOAD_MARKER)

        out_path = f"{WORKDIR}/guided_project_report.json"
        with open(out_path, "w") as f:
            json.dump(merged, f, indent=2)
        print(f"wrote {out_path} on retry (attempt 2)")
        print(json.dumps(merged, indent=2))

    ingested = ingest_taxi_data()
    validated = validate_trips(ingested)
    by_payment = compute_by_payment_type(validated)
    by_passengers = compute_by_passenger_count(validated)
    merged = merge_summaries(by_payment, by_passengers)
    load_report(merged)


cityflow_guided_project_authoring_dag()

import subprocess
import time
import json

dag_id = "cityflow_guided_project_authoring"

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"lesson5_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 on load_report means this run genuinely takes >20s
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_guided_project_authoring/dagRuns/"
    "lesson5_verify_1784885780/taskInstances/load_report",
    auth=("airflow", "airflow"),
)
info = resp.json()
print("load_report state:", info["state"])
print("load_report try_number:", info["try_number"])
assert info["try_number"] == 2, f"expected try_number 2, got {info['try_number']}"
print("confirmed: load_report needed 2 tries to succeed, same as Lesson 4's flaky_ingest")

{
  "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}
  ]
}