"""One-time prep: load GeoJSONs, reproject to TN State Plane US Feet, cache to parquet."""
from __future__ import annotations
import sys
import time
from pathlib import Path

import geopandas as gpd

ROOT = Path(__file__).resolve().parent.parent
GIS = ROOT / "gis"
CACHE = ROOT / "cache"
CACHE.mkdir(exist_ok=True)

TN_FT = "EPSG:2274"  # NAD83 / Tennessee State Plane, US Survey Feet


def prep(name: str, src: Path, keep_cols: list[str]) -> None:
    out = CACHE / f"{name}.parquet"
    if out.exists():
        print(f"[skip] {name} already cached -> {out}")
        return
    t0 = time.time()
    print(f"[read] {src.name} ...", flush=True)
    gdf = gpd.read_file(src)
    print(f"  loaded {len(gdf):,} features in {time.time()-t0:.1f}s; crs={gdf.crs}")
    keep_cols = [c for c in keep_cols if c in gdf.columns]
    gdf = gdf[keep_cols + ["geometry"]]
    t1 = time.time()
    print(f"[reproject] -> {TN_FT}", flush=True)
    gdf = gdf.to_crs(TN_FT)
    # Drop Z dimension if present (force XY)
    try:
        gdf["geometry"] = gpd.GeoSeries(
            [g if g is None or not g.has_z else _drop_z(g) for g in gdf.geometry],
            crs=gdf.crs,
        )
    except Exception:
        pass
    print(f"  reprojected in {time.time()-t1:.1f}s")
    gdf.to_parquet(out)
    print(f"[wrote] {out} ({out.stat().st_size/1e6:.1f} MB)")


def _drop_z(geom):
    from shapely.ops import transform
    return transform(lambda x, y, z=None: (x, y), geom)


def main() -> None:
    prep(
        "zoning",
        GIS / "Zoning.geojson",
        keep_cols=["OBJECTID", "ZONECLASS", "ZONEDESC", "ACRES_CALC"],
    )
    prep(
        "parcels",
        GIS / "Parcels.geojson",
        keep_cols=[
            "OBJECTID", "PARID", "GISLINK", "ADDRESS", "STREET",
            "OWNER", "OWNER2", "ZONECLASS", "ZONEDESC",
            "LANDUSE", "PROPTYPE", "CALC_ACRE",
        ],
    )
    prep(
        "addresses",
        GIS / "Addresses.geojson",
        keep_cols=[
            "OBJECTID", "GISLINK", "FULLADDR", "PLACETYPE",
            "ADDRCLASS", "ZONING", "OWNER_CALC",
        ],
    )
    prep(
        "roads",
        GIS / "Road_Centerlines.geojson",
        keep_cols=[
            "OBJECTID", "CENTERLINEID", "FULLNAME", "ROADCLASS",
            "ROADCLASS_PW", "MAINTBY", "PRIVDRIVETYPE",
        ],
    )


if __name__ == "__main__":
    sys.exit(main())
