# gate: skip
"""CityFlow's real instrumented ETL DAG (Module 6 Guided Project).

A real variant of Module 5's cityflow_etl_dag.py -- the original file is
untouched, this is a separate DAG id and a separate file -- instrumented
with everything this module built:

- validate_trips no longer reimplements the five real data-quality checks
  inline. It imports and calls cityflow_validation_logic.validate_trips_df
  (Lesson 1), the exact function this module unit-tested in isolation.
- default_args wires cityflow_on_failure_alert (Lesson 4's real callback,
  itself a thin wrapper around the already-tested cityflow_alerting module)
  as this DAG's on_failure_callback.
- ingest_taxi_trips reads its parquet path from dag_run.conf (falling back
  to the real path), so this same file can be triggered against a
  genuinely bad path to prove the alert fires on a real failure, and
  against the real data to prove it does not fire on a real success --
  without editing the DAG file between the two runs."""
import sys
from datetime import datetime

from airflow.decorators import dag, task

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

PARQUET_PATH = "/opt/airflow/dags/data/yellow_tripdata_2024-01.parquet"
ZONE_LOOKUP_PATH = "/opt/airflow/dags/data/taxi_zone_lookup.csv"
WORK_DIR = "/opt/airflow/dags/data/work/etl_instrumented"
RAW_TRIPS_PATH = f"{WORK_DIR}/raw_trips.csv"
CLEAN_TRIPS_PATH = f"{WORK_DIR}/clean_trips.csv"
QUARANTINE_PATH = f"{WORK_DIR}/quarantine_trips.csv"
VALIDATION_REPORT_PATH = f"{WORK_DIR}/validation_report.json"
TRANSFORMED_PATH = f"{WORK_DIR}/daily_borough_summary.csv"
ALERT_PATH = "/opt/airflow/dags/data/work/etl_instrumented/cityflow_alerts.jsonl"
INGEST_START = "2024-01-01"
INGEST_END_EXCLUSIVE = "2024-01-04"  # same real three-day slice as Module 5's guided project
KEEP_COLUMNS = [
    "tpep_pickup_datetime", "tpep_dropoff_datetime", "passenger_count",
    "trip_distance", "PULocationID", "DOLocationID", "payment_type",
    "fare_amount", "tip_amount", "total_amount",
]
MIN_PASS_RATE = 0.5
POSTGRES_CONN_ID = "cityflow_postgres"
TABLE_NAME = "cityflow_etl_instrumented_daily_borough_summary"


def cityflow_on_failure_alert(context):
    """This module's real on_failure_callback, wired at the DAG level via
    default_args -- fires once per real task failure, never on success."""
    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_etl_instrumented",
    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", "guided-project", "instrumented-etl"],
)
def cityflow_etl_instrumented_dag():
    @task
    def ingest_taxi_trips(**context) -> int:
        import os

        import pandas as pd

        # Reads the real parquet path by default -- but a caller can point
        # this at a genuinely bad path via dag_run.conf, the same
        # conf-driven pattern Module 5 Lesson 1 used to prove its own point.
        parquet_path = context["dag_run"].conf.get("parquet_path", PARQUET_PATH)
        raw = pd.read_parquet(parquet_path, columns=["tpep_pickup_datetime"] + [
            c for c in KEEP_COLUMNS if c != "tpep_pickup_datetime"
        ])
        raw["tpep_pickup_datetime"] = pd.to_datetime(raw["tpep_pickup_datetime"])
        raw["tpep_dropoff_datetime"] = pd.to_datetime(raw["tpep_dropoff_datetime"])
        start = pd.Timestamp(INGEST_START)
        end = pd.Timestamp(INGEST_END_EXCLUSIVE)
        sliced = raw[(raw["tpep_pickup_datetime"] >= start) & (raw["tpep_pickup_datetime"] < end)]
        sliced = sliced[KEEP_COLUMNS]
        os.makedirs(WORK_DIR, exist_ok=True)
        sliced.to_csv(RAW_TRIPS_PATH, index=False)
        print(f"ingest_taxi_trips: read {len(sliced)} real CityFlow trips from {parquet_path} -> {RAW_TRIPS_PATH}")
        return len(sliced)

    @task
    def validate_trips(ingested_rows: int) -> dict:
        """Unlike cityflow_etl_dag.py's original task, this one does not
        reimplement the five checks -- it imports and calls Lesson 1's
        already-unit-tested cityflow_validation_logic.validate_trips_df."""
        import sys

        import pandas as pd

        sys.path.insert(0, "/opt/airflow/dags")
        import cityflow_validation_logic as vlogic

        trips = pd.read_csv(RAW_TRIPS_PATH, parse_dates=["tpep_pickup_datetime", "tpep_dropoff_datetime"])
        assert len(trips) == ingested_rows

        report, any_fail = vlogic.validate_trips_df(trips)
        with open(VALIDATION_REPORT_PATH, "w") as f:
            import json
            json.dump(report, f, indent=2)
        trips[any_fail].to_csv(QUARANTINE_PATH, index=False)
        trips[~any_fail].to_csv(CLEAN_TRIPS_PATH, index=False)
        print(f"validate_trips: {report}")

        assert report["pass_rate"] >= MIN_PASS_RATE
        return report

    @task
    def transform_daily_borough_summary(validation_report: dict) -> dict:
        import pandas as pd

        clean = pd.read_csv(CLEAN_TRIPS_PATH, parse_dates=["tpep_pickup_datetime", "tpep_dropoff_datetime"])
        assert len(clean) == validation_report["rows_passed"]

        zones = pd.read_csv(ZONE_LOOKUP_PATH)[["LocationID", "Borough"]].rename(columns={"LocationID": "PULocationID"})
        clean = clean.merge(zones, on="PULocationID", how="left")
        clean["Borough"] = clean["Borough"].fillna("Unknown")
        clean["trip_duration_minutes"] = (
            clean["tpep_dropoff_datetime"] - clean["tpep_pickup_datetime"]
        ).dt.total_seconds() / 60.0
        clean["pickup_date"] = clean["tpep_pickup_datetime"].dt.date

        agg = clean.groupby(["pickup_date", "Borough"]).agg(
            trip_count=("fare_amount", "size"),
            avg_fare_amount=("fare_amount", "mean"),
            avg_trip_distance=("trip_distance", "mean"),
            avg_trip_duration_minutes=("trip_duration_minutes", "mean"),
            total_revenue=("total_amount", "sum"),
        ).reset_index()
        for col in ("avg_fare_amount", "avg_trip_distance", "avg_trip_duration_minutes", "total_revenue"):
            agg[col] = agg[col].round(2)
        agg.to_csv(TRANSFORMED_PATH, index=False)

        total_trip_count = int(agg["trip_count"].sum())
        assert total_trip_count == validation_report["rows_passed"]
        print(f"transform_daily_borough_summary: {len(agg)} rows, total_trip_count={total_trip_count}")
        return {"summary_rows": len(agg), "total_trip_count": total_trip_count}

    @task
    def load_daily_borough_summary_to_postgres(transform_result: dict) -> int:
        import pandas as pd
        from airflow.providers.postgres.hooks.postgres import PostgresHook

        summary = pd.read_csv(TRANSFORMED_PATH)
        hook = PostgresHook(postgres_conn_id=POSTGRES_CONN_ID)
        hook.run(f"DROP TABLE IF EXISTS {TABLE_NAME};")
        hook.run(f"""
            CREATE TABLE {TABLE_NAME} (
                id SERIAL PRIMARY KEY, pickup_date DATE NOT NULL, borough TEXT NOT NULL,
                trip_count INTEGER NOT NULL, avg_fare_amount REAL, avg_trip_distance REAL,
                avg_trip_duration_minutes REAL, total_revenue REAL
            );
        """)
        rows = list(summary[["pickup_date", "Borough", "trip_count", "avg_fare_amount", "avg_trip_distance",
                              "avg_trip_duration_minutes", "total_revenue"]].itertuples(index=False, name=None))
        hook.insert_rows(table=TABLE_NAME, rows=rows,
                          target_fields=["pickup_date", "borough", "trip_count", "avg_fare_amount",
                                         "avg_trip_distance", "avg_trip_duration_minutes", "total_revenue"],
                          commit_every=50)
        row_count = hook.get_first(f"SELECT COUNT(*) FROM {TABLE_NAME};")[0]
        total_trip_count = hook.get_first(f"SELECT SUM(trip_count) FROM {TABLE_NAME};")[0]
        assert row_count == len(summary)
        assert total_trip_count == transform_result["total_trip_count"]
        print(f"load_daily_borough_summary_to_postgres: {row_count} rows, SUM(trip_count)={total_trip_count}")
        return row_count

    load_daily_borough_summary_to_postgres(
        transform_daily_borough_summary(validate_trips(ingest_taxi_trips()))
    )


cityflow_etl_instrumented_dag()

# gate: setup: docker rm -f cityflow-lessons-db >/dev/null 2>&1; docker run -d --name cityflow-lessons-db -e POSTGRES_USER=cityflow -e POSTGRES_PASSWORD=cityflow -e POSTGRES_DB=cityflow -p 55432:5432 postgres:13 >/dev/null && sleep 4
import json
import os
import subprocess
import time

dag_id = "cityflow_etl_instrumented"
alert_host_path = ".airflow-local/dags/data/work/etl_instrumented/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"lesson5_m6_fail_{int(time.time())}"
conf = json.dumps({"parquet_path": "/opt/airflow/dags/data/does_not_exist_january.parquet"})
trigger = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "trigger", dag_id, "--run-id", run_id, "--conf", conf],
    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, "--conf", conf],
        capture_output=True, text=True, check=True,
    )
print(f"triggered {dag_id!r} as {run_id!r} with conf={conf}")

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
assert lines[0]["dag_id"] == dag_id
assert lines[0]["task_id"] == "ingest_taxi_trips"
assert lines[0]["run_id"] == run_id

# gate: teardown: docker rm -f cityflow-lessons-db >/dev/null 2>&1
import json
import subprocess
import time

dag_id = "cityflow_etl_instrumented"
alert_host_path = ".airflow-local/dags/data/work/etl_instrumented/cityflow_alerts.jsonl"

with open(alert_host_path) as f:
    lines_before = [l for l in f if l.strip()]
print(f"alert file has {len(lines_before)} line(s) before this run")

run_id = f"lesson5_m6_success_{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} (no conf -- real parquet path, real 3-day slice)")

deadline = time.time() + 120
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
    time.sleep(3)
    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)

with open(".airflow-local/dags/data/work/etl_instrumented/validation_report.json") as f:
    report = json.load(f)
print("validation_report.json:", json.dumps(report, indent=2))

import psycopg2

conn = psycopg2.connect(host="localhost", port=55432, user="cityflow", password="cityflow", dbname="cityflow")
cur = conn.cursor()
cur.execute("SELECT COUNT(*), SUM(trip_count) FROM cityflow_etl_instrumented_daily_borough_summary;")
row_count, total_trips = cur.fetchone()
print(f"direct psycopg2 check: {row_count} rows, SUM(trip_count)={total_trips}")
assert total_trips == report["rows_passed"]
conn.close()

time.sleep(1)
with open(alert_host_path) as f:
    lines_after = [l for l in f if l.strip()]
print(f"alert file has {len(lines_after)} line(s) after this run")
assert len(lines_after) == len(lines_before), "the on_failure_callback must NOT have fired on a real success"
print("PROVED: callback did not fire on a genuine successful run")