# gate: skip
"""CityFlow's real sensor-then-hook demo (Module 4 Lesson 4).

One DAG, two ideas composed: a FileSensor genuinely waits for a real batch of
CityFlow trip data to land, then a PostgresHook-based task loads that file's
real rows into the real cityflow-lessons-db Postgres instance -- wired end to
end, not two separate demos glued together in prose."""
from datetime import datetime

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

INCOMING_BATCH = "/opt/airflow/dags/data/work/sensor_hook_demo/incoming_batch.csv"
POSTGRES_CONN_ID = "cityflow_postgres"
TABLE_NAME = "cityflow_trip_batches"

with DAG(
    dag_id="cityflow_sensor_and_hook_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-4", "sensors-and-hooks"],
) as dag:
    wait_for_batch = FileSensor(
        task_id="wait_for_incoming_batch",
        filepath=INCOMING_BATCH,
        fs_conn_id="fs_default",
        poke_interval=5,
        timeout=120,
        mode="poke",
    )

    @task
    def load_batch_into_postgres() -> int:
        """Read the real trip rows the sensor just confirmed landed, and
        load them into Postgres through PostgresHook, not raw psycopg2."""
        import pandas as pd
        from airflow.providers.postgres.hooks.postgres import PostgresHook

        batch = pd.read_csv(INCOMING_BATCH)
        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,
                passenger_count REAL,
                trip_distance REAL,
                fare_amount REAL,
                payment_type INTEGER
            );
            """
        )

        rows = list(
            batch[["passenger_count", "trip_distance", "fare_amount", "payment_type"]]
            .itertuples(index=False, name=None)
        )
        hook.insert_rows(
            table=TABLE_NAME,
            rows=rows,
            target_fields=["passenger_count", "trip_distance", "fare_amount", "payment_type"],
        )
        print(f"loaded {len(rows)} real rows into {TABLE_NAME} via PostgresHook.insert_rows()")

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

    wait_for_batch >> load_batch_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_sensor_and_hook_demo"
host_incoming = ".airflow-local/dags/data/work/sensor_hook_demo/incoming_batch.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"lesson4_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_incoming_batch" in l][0]
    print(f"poll {i}: {sensor_line.strip()}")

print(">>> creating the real incoming batch file now (30 real trip rows) <<<")
import pandas as pd

raw = pd.read_csv(".airflow-local/dags/data/yellow_tripdata_sample.csv")
slice_ = raw.iloc[200:230]
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.
import psycopg2

conn = psycopg2.connect(host="localhost", port=55432, user="cityflow", password="cityflow", dbname="cityflow")
cur = conn.cursor()
cur.execute("SELECT COUNT(*) FROM cityflow_trip_batches;")
row_count = cur.fetchone()[0]
print(f"direct psycopg2 check: {row_count} rows in cityflow_trip_batches")
assert row_count == 30, f"expected 30, got {row_count}"
conn.close()