import json
import click
import pygeohash as pgh
import pyarrow.parquet as pq
import geopandas as gpd
from shapely.geometry import Point, LineString, MultiLineString
from pyproj import Geod
import numpy as np

@click.command()
@click.argument("latitude", type=float)
@click.argument("longitude", type=float)
@click.option("--distance", type=float, default=10, help="Distance threshold in meters for nearby lines.")
@click.option("--parquet-dir", type=str, default="geohash", help="Directory containing geohash Parquet files.")
def geohash_lookup(latitude, longitude, distance, parquet_dir):
    """
    Finds polygons containing the given point and lines within a threshold (in meters).
    """
    # Compute geohash (precision can be adjusted)
    geohash = pgh.encode(latitude, longitude, precision=4)  # Adjust precision as needed
    parquet_file = f"{parquet_dir}/{geohash}.parquet"

    try:
        # Load the Parquet file directly into a GeoDataFrame
        df = gpd.read_parquet(parquet_file)
    except FileNotFoundError:
        click.echo(json.dumps({"error": f"No Parquet file found for geohash {geohash}."}))
        return

    point = Point(longitude, latitude)
    geod = Geod(ellps="WGS84")  # Accurate geodetic distance calculations

    found_geometries = []

    # Vectorized distance calculation for LineStrings and MultiLineStrings
    def calculate_min_distance(geom):
        if isinstance(geom, LineString):
            return min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in geom.coords)
        elif isinstance(geom, MultiLineString):
            return min(min(geod.inv(point.x, point.y, p[0], p[1])[2] for p in line.coords) for line in geom.geoms)
        return np.inf

    for row in df.itertuples(index=False):
        geom = row.geometry  # Directly use the Shapely geometry object

        metadata = {}
        if isinstance(row.tags, dict):
            metadata = {k: v for k, v in row.tags.items() if v is not None}  # Remove empty values

        if geom.contains(point):
            found_geometries.append({"type": "polygon", **metadata})
        elif isinstance(geom, (LineString, MultiLineString)):
            min_distance = calculate_min_distance(geom)
            if min_distance <= distance:
                found_geometries.append({"type": "line" if isinstance(geom, LineString) else "multi_line", **metadata})

    # Output JSON result
    click.echo(json.dumps(found_geometries, indent=2, ensure_ascii=False))  # Ensure proper Unicode display

if __name__ == "__main__":
    geohash_lookup()