import os
import argparse
import geopandas as gpd
import pygeohash as pgh
import cbor2
import pandas as pd
import logging
from shapely.geometry import box
from shapely.ops import unary_union
from tqdm import tqdm

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

def geohash_bbox(geohash_code):
    """Returns a Shapely Polygon representing the bounding box of a geohash."""
    lat, lon, lat_err, lon_err = pgh.decode_exactly(geohash_code)
    lat_min, lon_min = lat - lat_err, lon - lon_err
    lat_max, lon_max = lat + lat_err, lon + lon_err
    return box(lon_min, lat_min, lon_max, lat_max)

def get_existing_geohashes(output_dir):
    """Returns a set of existing geohash parquet files in the output directory."""
    existing_geohashes = set()
    for filename in os.listdir(output_dir):
        if filename.endswith(".parquet"):
            geohash_code = filename[:-8]  # Remove '.parquet' extension
            existing_geohashes.add(geohash_code)
    logger.info(f"Found {len(existing_geohashes)} existing geohash files in {output_dir}.")
    return existing_geohashes

def subtract_existing_islands(water_gdf, existing_gdf):
    """Subtracts intersecting land polygons (islands) from water polygons."""
    if existing_gdf.empty:
        return water_gdf  # No islands to subtract

    # Collect all existing non-water geometries (islands or coastlines)
    land_geometries = existing_gdf[existing_gdf["tags"].apply(lambda tags: tags and tags.get("natural") != "water")]
    
    if land_geometries.empty:
        return water_gdf  # No islands to subtract

    # Merge islands into a single geometry
    islands_union = unary_union(land_geometries.geometry)

    # Subtract islands from the water polygons
    water_gdf["geometry"] = water_gdf.geometry.difference(islands_union)

    # Remove empty geometries resulting from subtraction
    water_gdf = water_gdf[~water_gdf.geometry.is_empty]

    return water_gdf

def intersect_geohash_tiles(water_polygons, coastline, geohashes):
    """Finds and slices water polygons into geohash tiles along the coastline."""
    results = {}

    logger.info("Processing geohashes for intersection with water polygons...")
    for ghash in tqdm(geohashes, desc="Processing geohashes"):
        ghash_poly = geohash_bbox(ghash)
        
        # Select only water polygons that intersect with the geohash tile
        water_subset = water_polygons[water_polygons.intersects(ghash_poly)]
        if water_subset.empty:
            continue
        
        # Clip with geohash bounding box
        clipped_water = gpd.clip(water_subset, ghash_poly)

        # Remove areas already covered by coastline
        coastline_subset = coastline[coastline.intersects(ghash_poly)]
        if not coastline_subset.empty:
            clipped_water = clipped_water.overlay(coastline_subset, how="difference")

        # Add water-specific tags
        if not clipped_water.empty:
            clipped_water["tags"] = [{"natural": "water", "water": "sea"}] * len(clipped_water)
            clipped_water = clipped_water[["tags", "geometry"]]  # Drop unnecessary columns
            results[ghash] = clipped_water

            logger.debug(f"Geohash {ghash} contains {len(clipped_water)} water polygons.")

    logger.info(f"Completed geohash processing. {len(results)} geohash tiles contain water polygons.")
    return results

def append_to_existing_parquet(geohash_tiles, output_dir):
    """Appends sliced water polygons to existing parquet files after subtracting islands."""
    for ghash, gdf in tqdm(geohash_tiles.items(), desc="Appending data"):
        parquet_path = os.path.join(output_dir, f"{ghash}.parquet")

        # Ensure input water polygons are in EPSG:4326
        gdf = gdf.to_crs("EPSG:4326")

        if os.path.exists(parquet_path):
            try:
                existing_gdf = gpd.read_parquet(parquet_path)

                # Ensure the existing file is in EPSG:4326 before merging
                if existing_gdf.crs != "EPSG:4326":
                    existing_gdf = existing_gdf.to_crs("EPSG:4326")

                # Subtract islands (land features) from the new water polygons
                gdf = subtract_existing_islands(gdf, existing_gdf)

                # Merge the existing and new data
                gdf = pd.concat([existing_gdf, gdf], ignore_index=True)
                logger.debug(f"Appended new water polygons to {ghash}.parquet")
            except Exception as e:
                logger.warning(f"Failed to read {parquet_path}. Overwriting instead. Error: {e}")

        # Save the updated file
        gdf.to_parquet(parquet_path)
        logger.debug(f"Updated {ghash}.parquet with new data.")

def save_optional_formats(geohash_tiles, output_dir, export_cbor, export_geojson):
    """Saves geohash-sliced water polygons in CBOR and GeoJSON formats."""
    for ghash, gdf in tqdm(geohash_tiles.items(), desc="Saving optional formats"):
        if export_cbor:
            cbor_path = os.path.join(output_dir, f"{ghash}.cbor")
            with open(cbor_path, "wb") as f:
                cbor2.dump(gdf.to_dict(), f)
            logger.debug(f"Saved {ghash}.cbor")

        if export_geojson:
            geojson_path = os.path.join(output_dir, f"{ghash}.geojson")
            gdf.to_file(geojson_path, driver="GeoJSON")
            logger.debug(f"Saved {ghash}.geojson")

def main():
    parser = argparse.ArgumentParser(description="Slice water polygons along the coastline into geohash tiles.")
    parser.add_argument("coastline", help="Path to the coastline Parquet file.")
    parser.add_argument("water", help="Path to the water polygons Parquet file.")
    parser.add_argument("output_dir", help="Directory to save the output files.")
    parser.add_argument("--cbor", action="store_true", help="Export as CBOR.")
    parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON.")
    parser.add_argument("--verbose", action="store_true", help="Enable detailed logging.")
    
    args = parser.parse_args()

    # Set logging level
    if args.verbose:
        logger.setLevel(logging.DEBUG)
        logger.debug("Verbose logging enabled.")

    logger.info("Loading input datasets...")
    coastline = gpd.read_parquet(args.coastline)
    water_polygons = gpd.read_parquet(args.water)

    # Ensure the geometries are valid
    coastline = coastline[coastline.geometry.is_valid]
    water_polygons = water_polygons[water_polygons.geometry.is_valid]

    logger.info(f"Loaded coastline dataset with {len(coastline)} features.")
    logger.info(f"Loaded water polygons dataset with {len(water_polygons)} features.")

    # Get existing geohashes from the output directory
    existing_geohashes = get_existing_geohashes(args.output_dir)

    if not existing_geohashes:
        logger.error("No existing geohash parquet files found in the output directory. Exiting.")
        return

    # Intersect water polygons with only the existing geohash tiles
    geohash_tiles = intersect_geohash_tiles(water_polygons, coastline, existing_geohashes)

    if not geohash_tiles:
        logger.info("No new data to append. Exiting.")
        return

    # Append results to existing parquet files
    append_to_existing_parquet(geohash_tiles, args.output_dir)

    # Save optional formats
    save_optional_formats(geohash_tiles, args.output_dir, args.cbor, args.geojson)

    logger.info("Processing complete. All files updated.")

if __name__ == "__main__":
    main()