"""Fetch NCES EDGE public school locations (2023-24) and load Ohio into PostGIS.

Bulk shapefile from NCES; we filter to STATE='OH' after read.
"""
from __future__ import annotations

import io
import zipfile
from pathlib import Path

import geopandas as gpd
import requests
from sqlalchemy import create_engine, text

DATA_DIR = Path("/home/m3ac/genoa-entwuerfe.com")
CACHE = DATA_DIR / "gis_scripts" / ".cache"
CACHE.mkdir(exist_ok=True)
YEAR = "2324"
URL = f"https://nces.ed.gov/programs/edge/data/EDGE_GEOCODE_PUBLICSCH_{YEAR}.zip"


def _read_gispass() -> str:
    for line in (DATA_DIR / ".env.gis").read_text().splitlines():
        if line.startswith("GISPASS="):
            return line.split("=", 1)[1].strip()
    raise SystemExit("GISPASS missing")


def download() -> Path:
    zip_path = CACHE / f"EDGE_GEOCODE_PUBLICSCH_{YEAR}.zip"
    if zip_path.exists() and zip_path.stat().st_size > 100000:
        print(f"[cache] {zip_path.name}")
        return zip_path
    print(f"[fetch] {URL}")
    r = requests.get(URL, timeout=180)
    r.raise_for_status()
    zip_path.write_bytes(r.content)
    print(f"  {len(r.content)/1e6:.1f} MB")
    return zip_path


def main() -> None:
    zpath = download()
    out = CACHE / f"edge_pubsch_{YEAR}"
    out.mkdir(exist_ok=True)
    with zipfile.ZipFile(zpath) as z:
        z.extractall(out)
    # find the .shp
    shps = list(out.rglob("*.shp"))
    if not shps:
        raise SystemExit("no .shp found in the zip")
    shp = shps[0]
    print(f"[read] {shp.name}")
    gdf = gpd.read_file(shp)
    print(f"[shp] {len(gdf)} US schools, cols={list(gdf.columns)}")
    # STFIP == '39' for Ohio, or STATE == 'OH'
    if "STATE" in gdf.columns:
        oh = gdf[gdf["STATE"] == "OH"].copy()
    elif "STFIP" in gdf.columns:
        oh = gdf[gdf["STFIP"] == "39"].copy()
    else:
        oh = gdf.copy()
    print(f"[oh] {len(oh)} Ohio schools")

    # normalize column names, lower case
    oh.columns = [c.lower() for c in oh.columns]
    if oh.crs and oh.crs.to_epsg() != 4326:
        oh = oh.to_crs("EPSG:4326")

    engine = create_engine(
        f"postgresql+psycopg2://gisapp:{_read_gispass()}@127.0.0.1:5432/westerville"
    )
    oh.to_postgis("schools", engine, if_exists="replace", index=False)
    with engine.begin() as conn:
        conn.execute(text("CREATE INDEX schools_geom_idx ON schools USING GIST (geometry);"))
        # try to add PK on NCES school id if present
        cols = [r[0] for r in conn.execute(text(
            "SELECT column_name FROM information_schema.columns WHERE table_name='schools'"
        ))]
        for cand in ("ncessch", "nces_sch", "school_id", "schoolid"):
            if cand in cols:
                conn.execute(text(f"ALTER TABLE schools ADD PRIMARY KEY ({cand});"))
                print(f"[pk] on {cand}")
                break
        (n,) = conn.execute(text("SELECT count(*) FROM schools")).one()
    print(f"[postgis] schools loaded: {n} rows")


if __name__ == "__main__":
    main()