import argparse
import duckdb
import geohash
import geopandas as gpd
import pandas as pd
import pyarrow.parquet as pq
import pyarrow as pa
import json
import cbor2
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:
            existing_data = gpd.read_parquet(file_path)
            combined_data = pd.concat([existing_data, new_data], ignore_index=True)

            if "feature_id" in combined_data.columns:
                combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last")

            combined_data = combined_data.reset_index(drop=True)
            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:
        new_data.to_parquet(file_path, index=False)
        print(f"✅ Created new file: {file_path}")

def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a GeoJSON file."""
    geojson_file = output_path / f"{geohash_code}.geojson"
    geojson_dict = json.loads(gdf.to_json())

    with open(geojson_file, "w", encoding="utf-8") as f:
        json.dump(geojson_dict, f)
    print(f"✅ Saved GeoJSON: {geojson_file}")

def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
    """Saves a GeoDataFrame as a CBOR file (compact binary format)."""
    cbor_file = output_path / f"{geohash_code}.cbor"
    geojson_dict = json.loads(gdf.to_json())

    with open(cbor_file, "wb") as f:
        cbor2.dump(geojson_dict, f)
    print(f"✅ Saved CBOR: {cbor_file}")

def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: bool, export_geojson: bool, export_cbor: bool):
    """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles."""
    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;")

    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'")

    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()

    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)

    geohash_data: Dict[str, gpd.GeoDataFrame] = {}

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

        filtered_data = con.execute(f"""
            SELECT feature_id, tags, 
                   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:
            filtered_data = filtered_data[filtered_data["clipped_geom"].apply(lambda x: isinstance(x, (bytes, bytearray)))]

            if not filtered_data.empty:
                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)

                filtered_data = filtered_data.dropna(subset=["geometry"])

                if not filtered_data.empty:
                    gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326")
                    geohash_data[geohash_code] = gdf

    for geohash_code, gdf in geohash_data.items():
        if export_parquet:
            geohash_file = output_path / f"{geohash_code}.parquet"
            append_to_parquet(geohash_file, gdf)
        
        if export_geojson:
            save_as_geojson(geohash_code, gdf, output_path)
        
        if export_cbor:
            save_as_cbor(geohash_code, gdf, output_path)

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

def main():
    parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles with export format options.")
    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 files.")
    parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet")
    parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON")
    parser.add_argument("--cbor", action="store_true", help="Export as CBOR")
    
    args = parser.parse_args()

    if not (args.parquet or args.geojson or args.cbor):
        args.parquet = True  # Default to GeoParquet if no options are given

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

if __name__ == "__main__":
    main()