"""Helpers for the watermask."""

import duckdb
import geopandas as gpd

from entropice.utils.paths import watermask_file


def open():
    """Open the watermask shapefile.

    Returns:
        GeoDataFrame: The watermask polygons.

    """
    watermask = gpd.read_file(watermask_file)
    return watermask


def clip_grid(grid_gdf: gpd.GeoDataFrame, allow_duckdb: bool = False) -> gpd.GeoDataFrame:
    """Clip the input GeoDataFrame with the watermask.

    Args:
        gdf (gpd.GeoDataFrame): The input GeoDataFrame to clip.

    Returns:
        gpd.GeoDataFrame: The clipped GeoDataFrame.

    """
    watermask = open()
    watermask = watermask.to_crs(grid_gdf.crs)
    # ! Currently disabled - kernel crashes
    allow_duckdb = False
    if len(grid_gdf) >= 10000 and allow_duckdb:
        # Use duckdb for large datasets
        crs = grid_gdf.crs

        # Convert geometry columns to WKB format for DuckDB
        # No need to copy the watermask
        watermask["geometry"] = watermask.geometry.to_wkb()
        grid_gdf = grid_gdf.copy()
        grid_gdf["geometry"] = grid_gdf.geometry.to_wkb()

        # Connect to DuckDB
        con = duckdb.connect(":memory:")

        # Install and load spatial extension
        con.execute("INSTALL spatial;")
        con.execute("LOAD spatial;")

        # Register the DataFrames as tables in DuckDB
        con.register("watermask", watermask)
        con.register("grid", grid_gdf)

        query = """
            SELECT g.*
            FROM grid g
            WHERE NOT EXISTS (
                SELECT 1
                FROM watermask w
                WHERE ST_Intersects(ST_GeomFromWKB(g.geometry), ST_GeomFromWKB(w.geometry))
            )
        """

        query = """
            WITH clipped AS (
                SELECT
                    g.* EXCLUDE (geometry),
                    CASE
                        WHEN EXISTS (
                            SELECT 1
                            FROM watermask w
                            WHERE ST_Intersects(ST_GeomFromWKB(g.geometry), ST_GeomFromWKB(w.geometry))
                        )
                        THEN ST_Difference(
                            ST_GeomFromWKB(g.geometry),
                            (
                                SELECT ST_Union_Agg(ST_GeomFromWKB(w.geometry))
                                FROM watermask w
                                WHERE ST_Intersects(ST_GeomFromWKB(g.geometry), ST_GeomFromWKB(w.geometry))
                            )
                        )
                        ELSE ST_GeomFromWKB(g.geometry)
                    END AS geometry
                FROM grid g
            )
            SELECT * FROM clipped
            WHERE NOT ST_IsEmpty(geometry)
        """

        result = con.execute(query).df()
        # Convert back to GeoDataFrame
        result["geometry"] = gpd.GeoSeries.from_wkb(result["geometry"].apply(bytes))
        grid_gdf = gpd.GeoDataFrame(result, geometry="geometry", crs=crs)
        con.close()
    else:
        grid_gdf = grid_gdf.overlay(watermask, how="difference")
    return grid_gdf
