import psycopg2

# From the host itself (this script, running in .venv-mlm-tutorials on the laptop):
conn = psycopg2.connect(host="localhost", port=55432, user="cityflow", password="cityflow", dbname="cityflow")
cur = conn.cursor()
cur.execute("SELECT version();")
print("host -> cityflow-lessons-db:", cur.fetchone()[0][:11])
conn.close()

import subprocess

# From INSIDE the Airflow scheduler container -- a different network entirely --
# host.docker.internal resolves to the host, and 55432 is the same published port.
result = subprocess.run(
    ["docker", "exec", "airflow-local-airflow-scheduler-1", "python", "-c",
     "import psycopg2; c = psycopg2.connect(host='host.docker.internal', port=55432, "
     "user='cityflow', password='cityflow', dbname='cityflow'); cur = c.cursor(); "
     "cur.execute('SELECT version();'); print('scheduler ->', cur.fetchone()[0][:11]); c.close()"],
    capture_output=True, text=True, check=True,
)
print(result.stdout.strip())
assert "scheduler ->" in result.stdout

# gate: skip
"""CityFlow's real PostgresHook demo (Module 4 Lesson 2).

A single Python task uses PostgresHook -- not raw psycopg2 -- to run a real
CREATE TABLE, a real set of INSERTs (CityFlow's real borough-zone counts from
the 265-zone lookup table), and a real SELECT that reads the rows back, all
against the real cityflow-lessons-db Postgres instance via the
'cityflow_postgres' Airflow Connection."""
from datetime import datetime

from airflow.decorators import dag, task

ZONE_LOOKUP = "/opt/airflow/dags/data/taxi_zone_lookup.csv"
POSTGRES_CONN_ID = "cityflow_postgres"
TABLE_NAME = "cityflow_hook_demo"


@dag(
    dag_id="cityflow_hooks_demo",
    schedule=None,
    start_date=datetime(2026, 1, 1),
    catchup=False,
    tags=["cityflow", "module-4", "hooks"],
)
def cityflow_hooks_demo_dag():
    @task
    def load_borough_counts_via_hook() -> int:
        """Compute CityFlow's real borough-zone counts, then create the
        table, insert every row, and read them back -- all through
        PostgresHook, never a raw psycopg2 connection."""
        import pandas as pd
        from airflow.providers.postgres.hooks.postgres import PostgresHook

        zones = pd.read_csv(ZONE_LOOKUP)
        counts = zones["Borough"].fillna("Outside of NYC").value_counts().to_dict()

        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,
                zone_name TEXT NOT NULL,
                zone_count INTEGER NOT NULL
            );
            """
        )
        print(f"created table {TABLE_NAME} via PostgresHook.run()")

        for borough, count in counts.items():
            hook.run(
                f"INSERT INTO {TABLE_NAME} (zone_name, zone_count) VALUES (%s, %s);",
                parameters=(borough, int(count)),
            )
        print(f"inserted {len(counts)} real borough rows via PostgresHook.run()")

        rows = hook.get_records(
            f"SELECT zone_name, zone_count FROM {TABLE_NAME} ORDER BY zone_name;"
        )
        for zone_name, zone_count in rows:
            print(f"  {zone_name}: {zone_count}")
        total = sum(r[1] for r in rows)
        assert total == 265, f"expected 265 real zones total, got {total}"
        print(f"SELECT via hook confirms {len(rows)} rows, {total} zones total")
        return len(rows)

    load_borough_counts_via_hook()


cityflow_hooks_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

dag_id = "cityflow_hooks_demo"

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

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

# Independent check #2: 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 zone_name, zone_count FROM cityflow_hook_demo ORDER BY zone_name;")
rows = cur.fetchall()
for zone_name, zone_count in rows:
    print(f"  {zone_name}: {zone_count}")
total = sum(r[1] for r in rows)
print(f"direct psycopg2 check: {len(rows)} rows, {total} zones total")
assert total == 265, f"expected 265, got {total}"
conn.close()