import os
import argparse
import pandas as pd
import geopandas as gpd
import geohash
import logging
from shapely.geometry import box

# Configure logging
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
logger = logging.getLogger(__name__)

def load_geohash_files(folder):
    """Finds all existing geohash parquet files in the folder."""
    geohash_files = {}
    for filename in os.listdir(folder):
        if filename.endswith(".parquet"):
            geohash_code = filename[:-8]  # Remove '.parquet' extension
            geohash_files[geohash_code] = os.path.join(folder, filename)
    logger.info(f"Found {len(geohash_files)} geohash parquet files in {folder}.")
    return geohash_files

def geohash_bbox(geohash_code):
    """Returns a bounding box polygon for a geohash."""
    bbox = geohash.bbox(geohash_code)
    return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"])

def ensure_crs_consistency(gdf):
    """Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches."""
    if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326":
        gdf = gdf.to_crs("EPSG:4326")
    return gdf

def process_geohash_files(water_gdf, geohash_files):
    """Processes each geohash file: adds water polygons from water.parquet."""
    # Precompute bounding boxes for all geohash codes
    geohash_bboxes = {geohash_code: geohash_bbox(geohash_code) for geohash_code in geohash_files.keys()}

    # Spatial index for faster spatial queries
    water_sindex = water_gdf.sindex

    for geohash_code, file_path in geohash_files.items():
        logger.info(f"Processing {file_path}...")

        # Load existing geohash file
        try:
            existing_gdf = gpd.read_parquet(file_path)
        except Exception as e:
            logger.warning(f"Failed to read {file_path}. Skipping. Error: {e}")
            continue

        existing_gdf = ensure_crs_consistency(existing_gdf)

        # Get the precomputed bounding box for this geohash
        bbox = geohash_bboxes[geohash_code]

        # Use spatial index to find intersecting water polygons
        possible_matches_index = list(water_sindex.intersection(bbox.bounds))
        possible_matches = water_gdf.iloc[possible_matches_index]
        water_in_tile = possible_matches[possible_matches.intersects(bbox)].copy()

        if water_in_tile.empty:
            logger.info(f"No water found for {geohash_code}. Skipping water addition.")
            continue

        # Clip water polygons to the geohash tile boundary
        water_in_tile["geometry"] = water_in_tile.intersection(bbox)

        # Assign "water" metadata
        water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile)

        # Ensure columns match before merging
        for col in ["feature_id", "tags"]:
            if col not in existing_gdf.columns:
                existing_gdf[col] = None
            if col not in water_in_tile.columns:
                water_in_tile[col] = None

        # Merge updated water polygons into geohash file
        updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326")

        # Save back to parquet
        updated_gdf.to_parquet(file_path, index=False)
        logger.info(f"Updated geohash file {file_path} with water polygons.")

def main():
    parser = argparse.ArgumentParser(description="Add water polygons to geohash parquet files from a global water dataset.")
    parser.add_argument("water", help="Path to the water.parquet file.")
    parser.add_argument("geohash_folder", help="Path to the folder containing <geohash>.parquet files.")
    
    args = parser.parse_args()

    # Load water dataset
    logger.info(f"Loading water dataset from {args.water}...")
    water_gdf = ensure_crs_consistency(gpd.read_parquet(args.water))

    # Get existing geohash parquet files
    geohash_files = load_geohash_files(args.geohash_folder)

    if not geohash_files:
        logger.warning("No geohash parquet files found. Exiting.")
        return

    # Process geohash files
    process_geohash_files(water_gdf, geohash_files)

    logger.info("Processing complete. All water polygons added.")

if __name__ == "__main__":
    main()