import urllib.request
import json

import pyarrow as pa
from pyarrow import parquet
import geoarrow.pyarrow as ga

manifest_url = (
    "https://raw.githubusercontent.com/geoarrow/geoarrow-data/v0.2.0-rc4/manifest.json"
)
files = {}
with urllib.request.urlopen(manifest_url) as f:
    manifest = json.load(f)
    for group in manifest["groups"]:
        for file in group["files"]:
            if file["format"] == "arrows/wkb":
                files[group["name"] + "_" + file["name"]] = file["url"]

out_dir = "/Users/dewey/gh/parquet-testing/data/geospatial"
ones_that_didnt_work = []
for name, url in files.items():
    # Skip big files + one CRS example that includes a non-PROJJSON value
    # on purpose (allowed in GeoArrow), which is rightly rejected
    # by Parquet
    if (
        "microsoft-buildings" in name
        or ("ns-water" in name and name != "ns-water_water-point")
        or "wkt2" in name
    ):
        print(f"Skipping {name}")
        continue

    # Maintain chunking from IPC into Parquet
    out = f"{out_dir}/{name}.parquet"
    with (
        urllib.request.urlopen(url) as f,
        pa.ipc.open_stream(f) as reader,
        parquet.ParquetWriter(
            out,
            reader.schema,
            store_schema=False,
            compression="none",
            write_geospatial_logical_types=True,
        ) as writer,
    ):
        original_schema = reader.schema
        print(f"Reading {url}")
        for batch in reader:
            writer.write_batch(batch)
        print(f"Wrote {out}")
    
    # Read in original table for comparison
    with (
        urllib.request.urlopen(url) as f,
        pa.ipc.open_stream(f) as reader
    ):
        original_table = reader.read_all()

    print(f"Checking {out}")
    with parquet.ParquetFile(out, arrow_extensions_enabled=True) as f:
        if f.schema_arrow != original_table.schema:
            print(f"Schema mismatch:\n{f.schema_arrow}\nvs\n{original_schema}")
            continue

        reread = f.read()
        if reread != original_table:
            print("Table mismatch")