# gate: skip
"""CityFlow's real file-arrival sensor demo (Module 4 Lesson 1).

A FileSensor genuinely polls for a real incoming-trips file rather than
assuming a fixed delay before reading it. The file does not exist when this
DAG is triggered -- it is created moments later, from a second process, while
the sensor is actively polling."""
from datetime import datetime

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

INCOMING_FILE = "/opt/airflow/dags/data/work/sensor_demo/incoming_trips.csv"

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

    @task
    def read_incoming_file():
        import pandas as pd

        trips = pd.read_csv(INCOMING_FILE)
        print(f"read {len(trips)} real rows from {INCOMING_FILE}")
        print(trips.head(3).to_string(index=False))
        return len(trips)

    wait_for_incoming_file >> read_incoming_file()

import json
import os
import subprocess
import time

dag_id = "cityflow_sensor_demo"
host_incoming = ".airflow-local/dags/data/work/sensor_demo/incoming_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"lesson1_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:
    # brand-new DAG file the scheduler's DagModel hasn't synced yet -- force it
    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}")

# Poll BEFORE creating the file -- proves the sensor is genuinely still waiting.
for i in range(4):
    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_file" in l][0]
    print(f"poll {i}: {sensor_line.strip()}")
    assert "running" in sensor_line or "up_for_reschedule" in sensor_line, \
        f"expected the sensor still waiting, got: {sensor_line}"

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

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

import subprocess

log_path = ".airflow-local/logs/dag_id=cityflow_sensor_demo/run_id=lesson1_m4_verify_1784888176/task_id=wait_for_incoming_file/attempt=1.log"
with open(log_path) as f:
    log_text = f.read()

poke_lines = [l for l in log_text.splitlines() if "Poking for file" in l]
found_lines = [l for l in log_text.splitlines() if "Found File" in l]
print(f"{len(poke_lines)} real 'Poking for file' log lines")
for l in poke_lines:
    print(" ", l)
assert len(poke_lines) >= 5, f"expected several real poke attempts, got {len(poke_lines)}"
assert found_lines, "expected a 'Found File' line once the sensor detected the real file"
print("\nconfirmed: the sensor polled repeatedly, then detected the real file")

log_path = ".airflow-local/logs/dag_id=cityflow_sensor_demo/run_id=lesson1_m4_verify_1784888176/task_id=read_incoming_file/attempt=1.log"
with open(log_path) as f:
    read_log = f.read()

for line in read_log.splitlines():
    if "real rows" in line:
        print(line)
assert "read 15 real rows" in read_log