import argparse
import duckdb
import geohash
import geopandas as gpd
import pandas as pd
import pyarrow.parquet as pq
import pyarrow as pa
from pathlib import Path
from shapely import wkb
from typing import List, Dict

def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
    """Generates all 4-character Geohashes that intersect a bounding box."""
    geohashes = set()
    step_size = 0.25  # 4-character Geohash grid size

    lat = min_y
    while lat <= max_y:
        lon = min_x
        while lon <= max_x:
            gh = geohash.encode(lat, lon, precision=4)
            geohashes.add(gh)
            lon += step_size
        lat += step_size

    return sorted(geohashes)

def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame):
    """Appends new data to an existing Parquet file while ensuring correct merging."""
    if file_path.exists():
        try:
            # ✅ Load existing data
            existing_data = gpd.read_parquet(file_path)
            
            # ✅ Merge with new data
            combined_data = pd.concat([existing_data, new_data], ignore_index=True)

            # ✅ Drop duplicates *only if feature_id exists*
            if "feature_id" in combined_data.columns:
                combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last")

            # ✅ Reset index before saving
            combined_data = combined_data.reset_index(drop=True)

            # ✅ Save back to the same file
            combined_data.to_parquet(file_path, index=False)
            print(f"✅ Appended new data to {file_path}")

        except Exception as e:
            print(f"❌ Error while merging {file_path}: {e}")

    else:
        # If the file doesn't exist, create it
        new_data.to_parquet(file_path, index=False)
        print(f"✅ Created new file: {file_path}")

def slice_and_split_geoparquet(input_file: str, output_dir: str):
    """
    Splits the GeoParquet file into multiple files based on 4-character Geohash boxes.
    Appends new data to existing Geohash Parquet files while preventing duplicate geometries.
    """
    input_path, output_path = Path(input_file), Path(output_dir)
    if not input_path.exists():
        raise FileNotFoundError(f"File '{input_file}' not found.")
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")

    con = duckdb.connect()
    con.execute("INSTALL spatial; LOAD spatial;")

    # Create a temporary table of all non POINT geometries
    con.execute(f"CREATE TEMP TABLE geoparquet AS SELECT * FROM read_parquet('{input_path}') WHERE ST_GeometryType(geometry) IS NOT NULL AND ST_GeometryType(geometry) != 'POINT'")

    # Compute bounding box for each geometry
    bbox_results = con.execute("""
        SELECT feature_id, rowid, ST_XMin(ST_Envelope(geometry)), ST_YMin(ST_Envelope(geometry)),
                     ST_XMax(ST_Envelope(geometry)), ST_YMax(ST_Envelope(geometry))
        FROM geoparquet;
    """).fetchall()

    # Compute all intersecting geohashes
    geohash_mapping: Dict[str, List[int]] = {}

    for feature_id, rowid, min_x, min_y, max_x, max_y in bbox_results:
        intersecting_geohashes = get_geohashes_from_bbox(min_x, min_y, max_x, max_y)
        for geohash_code in intersecting_geohashes:
            if geohash_code not in geohash_mapping:
                geohash_mapping[geohash_code] = []
            geohash_mapping[geohash_code].append(rowid)

    # Dictionary to store in-memory results before writing
    geohash_data: Dict[str, gpd.GeoDataFrame] = {}

    for geohash_code, rowids in geohash_mapping.items():
        geohash_bbox = geohash.bbox(geohash_code)

        # Clip geometries inside this geohash box
        filtered_data = con.execute(f"""
            SELECT feature_id, rowid, 
                   ST_AsWKB(ST_Intersection(geometry, 
                                            ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                                            {geohash_bbox["e"]}, {geohash_bbox["n"]}))) 
                   AS clipped_geom 
            FROM geoparquet 
            WHERE rowid IN ({','.join(map(str, rowids))})
            AND ST_Intersects(geometry, 
                              ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, 
                                              {geohash_bbox["e"]}, {geohash_bbox["n"]}));
        """).fetchdf()

        if not filtered_data.empty:
            # ✅ Remove non-WKB values
            filtered_data = filtered_data[filtered_data["clipped_geom"].apply(lambda x: isinstance(x, (bytes, bytearray)))]

            if not filtered_data.empty:
                # ✅ Convert WKB to Shapely geometries safely
                filtered_data["geometry"] = filtered_data["clipped_geom"].apply(
                    lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None
                )
                filtered_data.drop(columns=["clipped_geom"], inplace=True)

                # Remove rows where geometry conversion failed
                filtered_data = filtered_data.dropna(subset=["geometry"])

                if not filtered_data.empty:
                    # Convert to GeoPandas GeoDataFrame
                    gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")

                    # Store in memory
                    geohash_data[geohash_code] = gdf

    # Step 4: Write each geohash's data **only once** and prevent duplicates
    for geohash_code, gdf in geohash_data.items():
        geohash_file = output_path / f"{geohash_code}.parquet"
        append_to_parquet(geohash_file, gdf)
        print(f"✅ Updated: {geohash_file}")

    con.close()
    print("🎉 Processing complete!")

def main():
    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles.")
    parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.")
    parser.add_argument("-o", "--output", required=True, help="Directory to save output Geohash tiles.")
    
    args = parser.parse_args()

    try:
        slice_and_split_geoparquet(args.input, args.output)
    except Exception as e:
        print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()