# gate: skip
"""CityFlow's real Connections + Variables demo (Module 4 Lesson 3).

Nothing here is a hardcoded connection string or a hardcoded table name: the
task reads the Airflow Variable 'cityflow_hook_demo_table' to learn which
table to write to, and reaches Postgres through the named 'cityflow_postgres'
Connection via PostgresHook -- both configured for real through the Airflow
UI / CLI, not passed as literals in this file."""
from datetime import datetime

from airflow.decorators import dag, task
from airflow.models import Variable

POSTGRES_CONN_ID = "cityflow_postgres"


@dag(
    dag_id="cityflow_connections_variables_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-4", "connections-variables"],
)
def cityflow_connections_variables_demo_dag():
    @task
    def add_nyc_total_row() -> int:
        """Read the target table's name from a real Airflow Variable, then
        use the real 'cityflow_postgres' Connection (via PostgresHook) to
        insert one more row and confirm the new total."""
        from airflow.providers.postgres.hooks.postgres import PostgresHook

        table_name = Variable.get("cityflow_hook_demo_table")
        print(f"read Variable cityflow_hook_demo_table -> {table_name!r}")

        hook = PostgresHook(postgres_conn_id=POSTGRES_CONN_ID)

        hook.run(
            f"DELETE FROM {table_name} WHERE zone_name = %s;",
            parameters=("ALL_NYC_TOTAL",),
        )
        hook.run(
            f"INSERT INTO {table_name} (zone_name, zone_count) VALUES (%s, %s);",
            parameters=("ALL_NYC_TOTAL", 265),
        )
        print(f"inserted ALL_NYC_TOTAL row into {table_name!r} via the named Connection")

        row_count = hook.get_first(f"SELECT COUNT(*) FROM {table_name};")[0]
        print(f"SELECT COUNT(*) via hook: {row_count} rows in {table_name!r}")
        return row_count

    add_nyc_total_row()


cityflow_connections_variables_demo_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


def trigger_and_wait(dag_id, run_id_prefix):
    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"{run_id_prefix}_{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() + 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}"
    return run_id


# Step 1: replay Lesson 2's DAG so this fresh Postgres instance has the real
# 8-borough table Lesson 3's DAG is about to add a 9th row to.
trigger_and_wait("cityflow_hooks_demo", "lesson3_seed")

# Step 2: this lesson's own Connections+Variables-driven DAG.
trigger_and_wait("cityflow_connections_variables_demo", "lesson3_m4_verify")

# 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 zone_name, zone_count FROM cityflow_hook_demo ORDER BY zone_name;")
rows = cur.fetchall()
for r in rows:
    print(" ", r)
print(f"direct psycopg2 check: {len(rows)} rows")
assert len(rows) == 9, f"expected 9 rows (8 boroughs + ALL_NYC_TOTAL), got {len(rows)}"
conn.close()