# gate: skip
"""Module 4's guided project: CityFlow's real "wait for it to land, then
load it" DAG. A FileSensor waits for a real slice of CityFlow trip data to
land at a real path, then a PostgresHook-based task loads every real row
into the real cityflow-lessons-db Postgres instance via the same
'cityflow_postgres' Connection Lesson 2-4 used -- verified afterward with a
real SELECT COUNT(*)."""
from datetime import datetime

from airflow.decorators import task
from airflow.models.dag import DAG
from airflow.sensors.filesystem import FileSensor

LANDING_FILE = "/opt/airflow/dags/data/work/guided_project_m4/landing_trips.csv"
POSTGRES_CONN_ID = "cityflow_postgres"
TABLE_NAME = "cityflow_guided_project_trips"

with DAG(
    dag_id="cityflow_guided_project_wait_and_load",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-4", "guided-project"],
) as dag:
    wait_for_landing_file = FileSensor(
        task_id="wait_for_landing_file",
        filepath=LANDING_FILE,
        fs_conn_id="fs_default",
        poke_interval=5,
        timeout=120,
        mode="poke",
    )

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

        trips = pd.read_csv(LANDING_FILE)
        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_time TIMESTAMP,
                passenger_count REAL,
                trip_distance REAL,
                fare_amount REAL,
                payment_type INTEGER
            );
            """
        )

        cols = ["tpep_pickup_datetime", "passenger_count", "trip_distance", "fare_amount", "payment_type"]
        rows = list(trips[cols].itertuples(index=False, name=None))
        hook.insert_rows(
            table=TABLE_NAME,
            rows=rows,
            target_fields=["pickup_time", "passenger_count", "trip_distance", "fare_amount", "payment_type"],
            commit_every=100,
        )
        print(f"loaded {len(rows)} real CityFlow trip rows into {TABLE_NAME} via PostgresHook.insert_rows()")

        row_count = hook.get_first(f"SELECT COUNT(*) FROM {TABLE_NAME};")[0]
        assert row_count == len(trips), f"expected {len(trips)} rows, got {row_count}"
        print(f"SELECT COUNT(*) via hook confirms {row_count} real rows landed in Postgres")
        return row_count

    wait_for_landing_file >> load_landing_file_into_postgres()

# 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 os
import subprocess
import time

dag_id = "cityflow_guided_project_wait_and_load"
host_incoming = ".airflow-local/dags/data/work/guided_project_m4/landing_trips.csv"

os.makedirs(os.path.dirname(host_incoming), exist_ok=True)
if os.path.exists(host_incoming):
    os.remove(host_incoming)
print(f"confirmed {host_incoming} does not exist yet")

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

for i in range(3):
    time.sleep(4)
    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,
    )
    sensor_line = [l for l in r.stdout.splitlines() if "wait_for_landing_file" in l][0]
    print(f"poll {i}: {sensor_line.strip()}")

print(">>> creating the real landing file now (500 real trip rows) <<<")
import pandas as pd

raw = pd.read_csv(".airflow-local/dags/data/yellow_tripdata_sample.csv")
slice_ = raw.iloc[1000:1500]
slice_.to_csv(host_incoming, index=False)
print(f"wrote {len(slice_)} real rows to {host_incoming}")

deadline = time.time() + 90
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)

# 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(*) FROM cityflow_guided_project_trips;")
row_count = cur.fetchone()[0]
print(f"direct psycopg2 check: {row_count} rows in cityflow_guided_project_trips")
assert row_count == 500, f"expected 500, got {row_count}"
cur.execute("SELECT pickup_time, passenger_count, trip_distance, fare_amount, payment_type FROM cityflow_guided_project_trips ORDER BY id LIMIT 3;")
for row in cur.fetchall():
    print(" ", row)
conn.close()