"""
ornl_rebuild_bridges_into_databases_20260807.py
===========================

Created by: Andy Carter, PE
2026.08.10-- Assitance from Claude Opus 4.8

Rebuild the per-district TxDOT bridge FAST databases with freshly converted
data, then re-export each as a custom-format PostgreSQL dump.

For every district in `DISTRICTS`, the pipeline:

  1. Downloads the district's published `.dump` from TACC (skipped if present).
  2. Creates a local PostgreSQL 17 database and enables PostGIS.
  3. Restores the dump into that database.
  4. Verifies the parquet schema against the target table, then replaces
     `t_bridge_rating_curve` with the district's rating-curve parquet.
  5. Replaces `s_bridge_pnt` with the statewide bridge-point FGB.
  6. Dumps the revised database (custom format) to `OUT_FOLDER`.
  7. Drops the local database and deletes the downloaded source dump.

Progress is shown with tqdm; per-district results print as one summary line.

Requirements
------------
  - PostgreSQL 17 client tools (psql, createdb, pg_restore, pg_dump)
  - Python: requests, pandas, pyarrow, psycopg2, geopandas, tqdm
  - A running PostgreSQL 17 server reachable at DB_HOST:DB_PORT

Notes
-----
  - The v17 client tools are referenced by absolute path so a v14 install on
    PATH is never used by accident. Confirm DB_PORT points at the v17 server.
  - The full statewide FGB is loaded into EACH district database (matching the
    original behavior). It is read from disk only once and reused.
"""

from __future__ import annotations  # allow list[str]/tuple[...] hints on py<3.9

import io
import os
import subprocess
from pathlib import Path

import geopandas as gpd
import pandas as pd
import psycopg2
import pyarrow.parquet as pq
from psycopg2 import sql
from psycopg2.extras import execute_values
from tqdm import tqdm

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# --- PostgreSQL connection ---
DB_USER = "postgres"
DB_HOST = "localhost"
DB_PORT = "5432"          # <-- v17 server port (confirm; v14 may also use 5432)
os.environ["PGPASSWORD"] = "cornet37"

# --- PostgreSQL 17 client tools (absolute paths, so v14 on PATH is ignored) ---
PG_BIN = r"C:\Program Files\PostgreSQL\17\bin"
PSQL = os.path.join(PG_BIN, "psql.exe")
CREATEDB = os.path.join(PG_BIN, "createdb.exe")
PG_RESTORE = os.path.join(PG_BIN, "pg_restore.exe")
PG_DUMP = os.path.join(PG_BIN, "pg_dump.exe")

# --- Paths ---
BRIDGE_ROOT = Path(r"E:\bridge_conversions_20260805\output_20260807")
DOWNLOAD_DIR = Path(r"E:\bridge_conversions_20260805\localdownloads")
OUT_FOLDER = Path(r"E:\bridge_db_dump\bridge_db_revised_20260811")

# Statewide bridge-point FGB (loaded once, reused for every district).
#FGB_FILE = BRIDGE_ROOT / "merged_07_02_s_bridge_pnt_4326.fgb"

# Remote source dumps live here.
DUMP_URL_BASE = "https://web.corral.tacc.utexas.edu/ras2fim/fast_databases"

# --- Target tables ---
RATING_TABLE = "t_bridge_rating_curve"
POINT_TABLE = "s_bridge_pnt"

# Columns pandas adds when writing parquet -- not real data, ignore them.
IGNORE_COLS = {"__index_level_0__"}

# --- Behavior flags ---
DROP_AFTER = True          # drop the local DB once the revised dump is written
DELETE_DUMP_AFTER = True   # delete the downloaded source dump when finished

# Maps "<conversion folder name>" -> "<database / dump name>".
DISTRICTS = {
    "1_PAR":  "01_PAR_realtime_hand",
    "2_FTW":  "02_FTW_realtime_hand",
    "3_WFS":  "03_WFS_realtime_hand",
    "4_AMA":  "04_AMA_realtime_hand",
    "5_LBB":  "05_LBB_realtime_hand",
    "6_ODA":  "06_ODA_realtime_hand",
    "7_SJT":  "07_SJT_realtime_hand",
    "8_ABL":  "08_ABL_realtime_hand",
    "9_WAC":  "09_WAC_realtime_hand",
    "10_TYL": "10_TYL_realtime_hand",
    "11_LFK": "11_LFK_realtime_hand",
    "12_HOU": "12_HOU_realtime_hand",
    "13_YKM": "13_YKM_realtime_hand",
    "14_AUS": "14_AUS_realtime_hand",
    "15_SAT": "15_SAT_realtime_hand",
    "16_CRP": "16_CRP_realtime_hand",
    "17_BRY": "17_BRY_realtime_hand",
    "18_DAL": "18_DAL_realtime_hand",
    "19_ATL": "19_ATL_realtime_hand",
    "20_BMT": "20_BMT_realtime_hand",
    "21_PHR": "21_PHR_realtime_hand",
    "22_LRD": "22_LRD_realtime_hand",
    "23_BWD": "23_BWD_realtime_hand",
    "24_ELP": "24_ELP_realtime_hand",
    "25_CHS": "25_CHS_realtime_hand",
}

# ---------------------------------------------------------------------------
# Concat the merged_07_02_s_bridge_pnt_4326.fgb
# ---------------------------------------------------------------------------
TARGET_NAME = "07_02_s_bridge_pnt_4326.fgb"
FGB_FILE = BRIDGE_ROOT / "merged_07_02_s_bridge_pnt_4326.fgb"

# --- Find all matching files (recursive walk) ---
# rglob matches TARGET_NAME exactly, so the differently-named merged
# output won't be picked up even if this is re-run.
files = sorted(p for p in BRIDGE_ROOT.rglob(TARGET_NAME) if p.is_file())

if not files:
    raise FileNotFoundError(f"No files named {TARGET_NAME} found under {BRIDGE_ROOT}")

print(f"Found {len(files)} file(s) to merge:")
for f in files:
    print(f"  {f}")

# --- Read each file, tracking source for traceability ---
frames = []
target_crs = None

for f in files:
    gdf = gpd.read_file(f, engine="fiona")
    if gdf.empty:
        print(f"  (skipped, empty) {f}")
        continue

    # Lock CRS to the first non-empty file; reproject any that differ.
    if target_crs is None:
        target_crs = gdf.crs
    elif gdf.crs != target_crs:
        gdf = gdf.to_crs(target_crs)

    gdf["source_file"] = str(f)   # optional; drop if you don't want it
    frames.append(gdf)

if not frames:
    raise ValueError("All matching files were empty; nothing to merge.")

# --- Concatenate ---
merged = gpd.GeoDataFrame(
    pd.concat(frames, ignore_index=True),
    crs=target_crs,
)
print(f"\nMerged {len(merged)} bridge points.")

# --- Write out ---
FGB_FILE.parent.mkdir(parents=True, exist_ok=True)
merged.to_file(FGB_FILE, driver="FlatGeobuf", engine="fiona")
print(f"Wrote {FGB_FILE}")


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------

def log(msg: str) -> None:
    """Print without clobbering active tqdm progress bars."""
    tqdm.write(msg)


def _conn_args() -> list[str]:
    """Common host/port/user flags for the psql-family CLI tools."""
    return ["-h", DB_HOST, "-p", DB_PORT, "-U", DB_USER]


def sh(cmd: list[str], check: bool = True) -> subprocess.CompletedProcess:
    """Run a command quietly. Raise on failure (unless check=False)."""
    result = subprocess.run(cmd, capture_output=True, text=True)
    if check and result.returncode != 0:
        raise RuntimeError(
            f"Command failed (exit {result.returncode}): {' '.join(map(str, cmd))}\n"
            f"{(result.stderr or result.stdout).strip()}"
        )
    return result


def psql_scalar(db: str, query: str) -> str:
    """Run a query with psql -tAc and return the stripped stdout."""
    return sh([PSQL, *_conn_args(), "-d", db, "-tAc", query]).stdout.strip()


def connect(db: str) -> psycopg2.extensions.connection:
    """Open a psycopg2 connection to the given database."""
    return psycopg2.connect(
        host=DB_HOST, port=DB_PORT, user=DB_USER,
        password=os.environ["PGPASSWORD"], dbname=db,
    )


def normalize_type(t: str) -> str:
    """Collapse parquet/arrow and postgres type names to a common family."""
    t = t.lower()
    if any(k in t for k in ("string", "text", "char", "utf8", "varchar")):
        return "text"
    if any(k in t for k in ("int64", "int32", "bigint", "integer", "smallint")):
        return "integer"
    if any(k in t for k in ("double", "float", "real", "numeric", "decimal")):
        return "float"
    if "bool" in t:
        return "boolean"
    if any(k in t for k in ("timestamp", "date", "time")):
        return "datetime"
    return t  # fallback: compare raw


def db_columns(db: str, table: str) -> list[str]:
    """Return a table's columns in ordinal (definition) order."""
    out = psql_scalar(
        db,
        "SELECT column_name FROM information_schema.columns "
        f"WHERE table_name = '{table}' ORDER BY ordinal_position;",
    )
    return [c for c in out.splitlines() if c]


# ---------------------------------------------------------------------------
# Pipeline steps
# ---------------------------------------------------------------------------

def download_dump(url: str, dest: Path) -> None:
    """Stream-download `url` to `dest`, skipping if already present."""
    if dest.exists() and dest.stat().st_size > 0:
        return
    import requests  # local import keeps the top-level import list lean
    with requests.get(url, stream=True, timeout=60) as r:
        r.raise_for_status()
        total = int(r.headers.get("content-length", 0))
        with open(dest, "wb") as f, tqdm(
            total=total, unit="B", unit_scale=True, unit_divisor=1024,
            desc="  download", leave=False,
        ) as bar:
            for chunk in r.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)
                    bar.update(len(chunk))


def ensure_database(db: str) -> None:
    """Create the database if it does not already exist."""
    if psql_scalar("postgres", f"SELECT 1 FROM pg_database WHERE datname = '{db}'") != "1":
        sh([CREATEDB, *_conn_args(), db])


def enable_postgis(db: str) -> None:
    """Enable PostGIS (raster/topology are best-effort)."""
    sh([PSQL, *_conn_args(), "-d", db, "-c", "CREATE EXTENSION IF NOT EXISTS postgis;"])
    for ext in ("postgis_raster", "postgis_topology"):
        sh([PSQL, *_conn_args(), "-d", db, "-c",
            f"CREATE EXTENSION IF NOT EXISTS {ext};"], check=False)


def restore_dump(db: str, dump_file: Path) -> None:
    """Restore a custom-format dump into `db` (benign warnings are ignored)."""
    sh([PG_RESTORE, *_conn_args(), "-d", db,
        "--no-owner", "--no-privileges", str(dump_file)], check=False)


def verify_rating_schema(db: str, parquet_file: Path) -> list[str]:
    """
    Compare the parquet schema against the DB table's schema.

    Raises if the column *sets* differ. Returns a list of columns whose type
    families differ (empty means a clean match).
    """
    pf = pq.ParquetFile(parquet_file)
    pq_map = {f.name: str(f.type) for f in pf.schema_arrow if f.name not in IGNORE_COLS}

    db_out = psql_scalar(
        db,
        "SELECT column_name, data_type FROM information_schema.columns "
        f"WHERE table_name = '{RATING_TABLE}' ORDER BY ordinal_position;",
    )
    db_map = dict(line.split("|") for line in db_out.splitlines() if line)

    only = (set(pq_map) - set(db_map)) | (set(db_map) - set(pq_map))
    if only:
        raise RuntimeError(f"{RATING_TABLE} column set mismatch: {sorted(only)}")

    return [n for n in pq_map if normalize_type(pq_map[n]) != normalize_type(db_map[n])]


def load_rating_curve(db: str, parquet_file: Path) -> tuple[int, int]:
    """TRUNCATE + COPY the rating-curve parquet into the DB table (one txn)."""
    df = pd.read_parquet(parquet_file).drop(
        columns=[c for c in IGNORE_COLS], errors="ignore",
    )
    col_order = db_columns(db, RATING_TABLE)
    assert set(col_order) == set(df.columns), (
        f"Column mismatch!\n db: {col_order}\n parquet: {list(df.columns)}"
    )
    df = df[col_order]  # match table order for COPY

    buf = io.StringIO()
    df.to_csv(buf, index=False, header=False, na_rep="\\N")
    buf.seek(0)

    conn = connect(db)
    try:
        with conn:  # commit on success, rollback on error
            with conn.cursor() as cur:
                cur.execute(f'SELECT count(*) FROM "{RATING_TABLE}";')
                before = cur.fetchone()[0]
                cur.execute(f'TRUNCATE TABLE "{RATING_TABLE}";')
                cur.copy_expert(
                    f'COPY "{RATING_TABLE}" ({", ".join(col_order)}) '
                    f"FROM STDIN WITH (FORMAT csv, NULL '\\N')",
                    buf,
                )
                cur.execute(f'SELECT count(*) FROM "{RATING_TABLE}";')
                after = cur.fetchone()[0]
    finally:
        conn.close()
    return before, after


def load_bridge_points(db: str, gdf: gpd.GeoDataFrame) -> int:
    """TRUNCATE + batch-INSERT the statewide bridge points into the DB table."""
    geo_meta = psql_scalar(
        db,
        "SELECT f_geometry_column, type, srid FROM geometry_columns "
        f"WHERE f_table_name = '{POINT_TABLE}';",
    )
    geom_col, _, srid_str = geo_meta.split("|")
    srid = int(srid_str)

    # Align CRS with the table (gdf is reused across districts).
    if gdf.crs is None:
        gdf = gdf.set_crs(epsg=srid)
    elif gdf.crs.to_epsg() != srid:
        gdf = gdf.to_crs(epsg=srid)

    attr_cols = [c for c in db_columns(db, POINT_TABLE) if c != geom_col]
    attrs = gdf[attr_cols].astype(object).where(pd.notnull(gdf[attr_cols]), None)
    wkb_hex = gdf.geometry.apply(lambda g: g.wkb_hex if g is not None else None)
    rows = [tuple(attrs.iloc[i]) + (wkb_hex.iloc[i],) for i in range(len(gdf))]

    col_list = ", ".join(f'"{c}"' for c in attr_cols) + f', "{geom_col}"'
    template = (
        "(" + ", ".join(["%s"] * len(attr_cols))
        + f", ST_GeomFromWKB(decode(%s, 'hex'), {srid}))"
    )
    insert_sql = f'INSERT INTO "{POINT_TABLE}" ({col_list}) VALUES %s'

    conn = connect(db)
    try:
        with conn:  # commit on success, rollback on error
            with conn.cursor() as cur:
                cur.execute(f'TRUNCATE TABLE "{POINT_TABLE}";')
                batch = 5000
                for start in tqdm(range(0, len(rows), batch),
                                  desc="  points", leave=False):
                    execute_values(cur, insert_sql, rows[start:start + batch],
                                   template=template, page_size=1000)
                cur.execute(f'SELECT count(*) FROM "{POINT_TABLE}";')
                after = cur.fetchone()[0]
    finally:
        conn.close()
    return after


def dump_database(db: str, out_file: Path) -> int:
    """Export `db` as a custom-format dump. Returns the file size in bytes."""
    sh([PG_DUMP, *_conn_args(), "-d", db, "-Fc", "-f", str(out_file)])
    return out_file.stat().st_size


def drop_database(db: str) -> None:
    """Drop `db`, forcibly terminating any other sessions (Postgres 13+)."""
    conn = connect("postgres")  # cannot drop the DB we are connected to
    conn.autocommit = True      # DROP DATABASE cannot run inside a transaction
    try:
        with conn.cursor() as cur:
            cur.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(
                sql.Identifier(db)))
    finally:
        conn.close()


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------

def process_district(folder: str, db: str, gdf: gpd.GeoDataFrame) -> str:
    """Run the full rebuild pipeline for one district; return a summary line."""
    parquet_file = (BRIDGE_ROOT / folder / "07_for_FAST_database"
                    / "07_01_t_bridge_rating_curve.parquet")
    dump_file = DOWNLOAD_DIR / f"{db}.dump"
    out_file = OUT_FOLDER / f"{db}.dump"

    download_dump(f"{DUMP_URL_BASE}/{db}.dump", dump_file)
    ensure_database(db)
    enable_postgis(db)
    restore_dump(db, dump_file)

    type_diffs = verify_rating_schema(db, parquet_file)
    if type_diffs:
        log(f"  ! {db}: type differences in {RATING_TABLE}: {type_diffs}")

    before, after = load_rating_curve(db, parquet_file)
    n_points = load_bridge_points(db, gdf)
    size_gb = dump_database(db, out_file) / 1e9

    if DROP_AFTER:
        drop_database(db)
    if DELETE_DUMP_AFTER:
        dump_file.unlink(missing_ok=True)

    return (f"{db}: rating {before:,}->{after:,} | "
            f"points {n_points:,} | dump {size_gb:.2f} GB")


def main() -> None:
    """Prepare output/download dirs, load the FGB once, process all districts."""
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
    OUT_FOLDER.mkdir(parents=True, exist_ok=True)

    # Confirm the client tooling is v17 before touching any data.
    log("pg_restore " + sh([PG_RESTORE, "--version"]).stdout.strip())

    # Read the statewide bridge-point FGB once and reuse it for every district.
    log(f"Loading FGB: {FGB_FILE}")
    gdf = gpd.read_file(FGB_FILE)
    log(f"  {len(gdf):,} features, CRS={gdf.crs}, "
        f"geom={sorted(gdf.geom_type.unique())}")

    failures = []
    bar = tqdm(DISTRICTS.items(), desc="districts", unit="db")
    for folder, db in bar:
        bar.set_postfix_str(db)
        try:
            log("  " + process_district(folder, db, gdf))
        except Exception as exc:  # keep going; report at the end
            failures.append(db)
            log(f"  X {db} FAILED: {exc}")

    log("\nDone." + (f" Failures: {failures}" if failures else " All districts OK."))


if __name__ == "__main__":
    main()