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, Polygon
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 generate_4digit_geohashes():
    """Generates all possible 4-digit geohashes covering the world."""
    base_geohashes = "0123456789bcdefghjkmnpqrstuvwxyz"
    return [a + b + c + d for a in base_geohashes for b in base_geohashes for c in base_geohashes for d in base_geohashes]

def intersect_geohash_tiles(water_polygons, coastline, geohashes):
    """Finds and slices water polygons into 4-digit 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")

        if not clipped_water.empty:
            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 save_data(geohash_tiles, output_dir, export_cbor, export_geojson):
    """Saves geohash-sliced water polygons in requested formats."""
    os.makedirs(output_dir, exist_ok=True)
    logger.info(f"Saving data to {output_dir}...")

    for ghash, gdf in tqdm(geohash_tiles.items(), desc="Saving geohashes"):
        parquet_path = os.path.join(output_dir, f"{ghash}.parquet")
        gdf.to_parquet(parquet_path)
        logger.debug(f"Saved {ghash}.parquet")

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

    logger.info("Data saving process completed.")

def main():
    parser = argparse.ArgumentParser(description="Slice water polygons along the coastline into 4-digit 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.")

    # Generate 4-digit geohashes
    logger.info("Generating all possible 4-digit geohashes...")
    geohashes = generate_4digit_geohashes()
    logger.info(f"Generated {len(geohashes)} geohash tiles.")

    # Intersect water polygons with geohash tiles
    geohash_tiles = intersect_geohash_tiles(water_polygons, coastline, geohashes)

    # Save results
    save_data(geohash_tiles, args.output_dir, args.cbor, args.geojson)

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

if __name__ == "__main__":
    main()