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

# 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":
        logger.info(f"Converting CRS of {gdf} to EPSG:4326")
        gdf = gdf.to_crs("EPSG:4326")
    return gdf

def process_geohash_files(coastline_gdf, geohash_files):
    """Processes each geohash file: removes old coastlines and inserts updated ones."""
    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)

        # Remove existing coastline features
        filtered_gdf = existing_gdf[~((existing_gdf.geometry.type.isin(["LineString", "MultiLineString"])) &
                                      (existing_gdf["tags"].apply(lambda tags: isinstance(tags, dict) and tags.get("natural") == "coastline")))]

        # Get coastline geometries for this geohash
        bbox = geohash_bbox(geohash_code)
        new_coastline = coastline_gdf[coastline_gdf.intersects(bbox)].copy()

        # Clip coastline to the geohash boundary
        new_coastline["geometry"] = new_coastline["geometry"].apply(lambda geom: geom.intersection(bbox))

        if new_coastline.empty:
            logger.info(f"No coastline found for {geohash_code}. Skipping update.")
            continue

        # Assign "natural: coastline" tag
        new_coastline["tags"] = [{"natural": "coastline"}] * len(new_coastline)

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

        # Merge updated coastline into geohash file
        updated_gdf = gpd.GeoDataFrame(pd.concat([filtered_gdf, new_coastline], ignore_index=True), crs="EPSG:4326")

        # Save back to parquet
        updated_gdf.to_parquet(file_path, index=False)
        logger.info(f"Updated coastline in {file_path}")

def main():
    parser = argparse.ArgumentParser(description="Update coastline geometries in existing geohash parquet files.")
    parser.add_argument("coastline", help="Path to the coastline.parquet file.")
    parser.add_argument("geohash_folder", help="Path to the folder containing <geohash>.parquet files.")
    
    args = parser.parse_args()

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

    # 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(coastline_gdf, geohash_files)

    logger.info("Processing complete. All coastline updates applied.")

if __name__ == "__main__":
    main()