# gate: skip
"""CityFlow's real production-shaped ETL DAG (Module 5 Guided Project).

The same four real stages Lessons 2-4 built and verified individually --
ingest, validate, transform, load -- combined into one DAG and scaled up
from Lessons 2-4's one real day to three real days of CityFlow's actual
January 2024 taxi data (2024-01-01 through 2024-01-03, ~239K real rows,
easily CityFlow's largest single DAG run in this course so far)."""
from datetime import datetime

from airflow.decorators import dag, task

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"
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"
INGEST_START = "2024-01-01"
INGEST_END_EXCLUSIVE = "2024-01-04"  # three real days for the 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_daily_borough_summary"


@dag(
    dag_id="cityflow_etl",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-5", "guided-project"],
)
def cityflow_etl_dag():
    @task
    def ingest_taxi_trips() -> int:
        import os

        import pandas as pd

        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"ingested {len(sliced)} real CityFlow trips ({INGEST_START} to {INGEST_END_EXCLUSIVE}) -> {RAW_TRIPS_PATH}")
        return len(sliced)

    @task
    def validate_trips(ingested_rows: int) -> dict:
        import json

        import pandas as pd

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

        required = ["passenger_count", "trip_distance", "fare_amount", "PULocationID", "DOLocationID"]
        null_mask = trips[required].isna().any(axis=1)
        dist_mask = (trips["trip_distance"] < 0) | (trips["trip_distance"] > 100)
        fare_mask = (trips["fare_amount"] < 0) | (trips["fare_amount"] > 500)
        pax_mask = (trips["passenger_count"] < 0) | (trips["passenger_count"] > 6)
        time_mask = trips["tpep_pickup_datetime"] > trips["tpep_dropoff_datetime"]

        checks = {
            "null_required_fields": int(null_mask.sum()),
            "trip_distance_out_of_range": int(dist_mask.sum()),
            "fare_amount_out_of_range": int(fare_mask.sum()),
            "passenger_count_out_of_range": int(pax_mask.sum()),
            "pickup_after_dropoff": int(time_mask.sum()),
        }
        any_fail = null_mask | dist_mask | fare_mask | pax_mask | time_mask
        rows_failed = int(any_fail.sum())
        rows_passed = int((~any_fail).sum())
        pass_rate = rows_passed / len(trips)

        report = {
            "total_rows": len(trips), "checks": checks,
            "rows_failed_any_check": rows_failed, "rows_passed": rows_passed,
            "pass_rate": round(pass_rate, 4),
        }
        with open(VALIDATION_REPORT_PATH, "w") as f:
            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"validation report -> {VALIDATION_REPORT_PATH}: {report}")

        assert 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"{len(agg)} real day/borough summary rows -> {TRANSFORMED_PATH}, 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,
        )
        print(f"loaded {len(rows)} real summary rows into {TABLE_NAME} via PostgresHook.insert_rows()")

        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"SELECT COUNT(*) via hook: {row_count} rows; SELECT 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_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
# gate: teardown: docker rm -f cityflow-lessons-db >/dev/null 2>&1
import json
import subprocess
import time

dag_id = "cityflow_etl"

deadline = time.time() + 320
found = False
while time.time() < deadline:
    result = subprocess.run(
        ["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list", "--output", "json"],
        capture_output=True, text=True, check=True,
    )
    if any(d["dag_id"] == dag_id for d in json.loads(result.stdout)):
        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_m5_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() + 180
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)
    state = next(x["state"] for x in runs if x["run_id"] == run_id)
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/validation_report.json") as f:
    report = json.load(f)
print("validation_report.json:", json.dumps(report, indent=2))

# Independent check: query cityflow-lessons-db directly, bypassing Airflow entirely.
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_daily_borough_summary;")
row_count, total_trips = cur.fetchone()
print(f"direct psycopg2 check: {row_count} rows, SUM(trip_count)={total_trips}")
assert row_count == 21, f"expected 21, got {row_count}"
assert total_trips == 220898, f"expected 220898, got {total_trips}"
assert total_trips == report["rows_passed"], "loaded total doesn't match validate_trips's own reported pass count"
conn.close()