import os
import argparse
import pandas as pd
import geopandas as gpd
import geohash
import numpy as np
import logging
from shapely.geometry import box, LineString, MultiLineString, Point
from shapely.ops import unary_union, linemerge, split
import matplotlib.pyplot as plt

def split_polygon_by_coastline(tile_polygon, coastline):
    """
    Splits the geohash tile polygon into two by the coastline and keeps the polygon on the 'right' side
    based on the LineString's direction.
    
    Parameters:
        - tile_polygon: The full geohash bounding box as a Polygon.
        - coastline: The merged LineString coastline within the tile.
    
    Returns:
        - The remaining water polygon after removing the 'left' side.
    """
    if coastline.is_empty or coastline.geom_type not in ["LineString", "MultiLineString"]:
        return None  # No valid coastline

    # Convert MultiLineString to the longest single LineString if necessary
    if coastline.geom_type == "MultiLineString":
        coastline = max(coastline.geoms, key=lambda g: g.length)

    # Get the midpoint of the LineString
    midpoint = coastline.interpolate(0.5, normalized=True)

    # Compute the direction of the coastline (from first to last point)
    coords = np.array(coastline.coords)
    start, end = coords[0], coords[-1]
    dx, dy = end[0] - start[0], end[1] - start[1]

    # Compute normal vector (perpendicular to the coastline direction)
    normal = np.array([-dy, dx])  # Rotate 90° counterclockwise
    normal = normal / np.linalg.norm(normal)  # Normalize

    # Compute a test point on the 'right' side
    test_point = midpoint.x + normal[0], midpoint.y + normal[1]

    # Perform the split
    split_result = split(tile_polygon, coastline)

    if not split_result or len(split_result.geoms) < 2:
        return None  # No valid split

    # Choose the polygon that contains the test point
    for poly in split_result.geoms:
        if poly.contains(Point(test_point)):
            return poly  # Keep the 'right' side

    return None  # Fail-safe


def group_continuous_coastline(segments, gap_threshold=0.001):
    """
    Groups and merges only those LineStrings that are properly connected to each other.
    Separate coastlines remain unmerged (e.g., islands).
    
    Parameters:
        - segments: List of LineStrings.
        - gap_threshold: Max distance (degrees) between segment endpoints to be considered connected.
    
    Returns:
        - List of merged continuous LineStrings.
    """
    connected_groups = []
    remaining_segments = list(segments)

    while remaining_segments:
        # Start a new group with one segment
        group = [remaining_segments.pop(0)]

        added = True
        while added:
            added = False
            for seg in remaining_segments[:]:
                if is_connected_to_group(seg, group, gap_threshold):
                    group.append(seg)
                    remaining_segments.remove(seg)
                    added = True  # Continue expanding the group

        # Merge each connected group separately
        if len(group) > 1:
            merged = linemerge(MultiLineString(group))
        else:
            merged = group[0]

        connected_groups.append(merged)

    return connected_groups


def is_connected_to_group(line, group, gap_threshold=0.001):
    """
    Checks if a LineString is connected to any other in a group.
    A connection means they share an endpoint within a small gap threshold.
    """
    line_start, line_end = Point(line.coords[0]), Point(line.coords[-1])

    for other in group:
        other_start, other_end = Point(other.coords[0]), Point(other.coords[-1])
        if (
            line_start.distance(other_end) < gap_threshold or
            line_end.distance(other_start) < gap_threshold
        ):
            return True  # They are connected within the gap threshold

    return False

def plot_coastline_segments(geohash_code, fragmented_coastline, output_folder="debug_plots"):
    """Saves a PNG plot of coastline segments for debugging."""
    os.makedirs(output_folder, exist_ok=True)  # Ensure the output folder exists
    output_file = os.path.join(output_folder, f"{geohash_code}.png")

    fig, ax = plt.subplots(figsize=(6, 6))
    colors = ['red', 'blue', 'green', 'purple', 'orange', 'cyan']

    # Convert to list if MultiLineString
    if fragmented_coastline.geom_type == "MultiLineString":
        segments = list(fragmented_coastline.geoms)
    else:
        segments = [fragmented_coastline]

    for idx, segment in enumerate(segments):
        x, y = segment.xy
        ax.plot(x, y, color=colors[idx % len(colors)], linewidth=2, label=f"Segment {idx+1}")

    ax.set_title(f"Coastline Segments in {geohash_code}")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    ax.legend()
    
    plt.savefig(output_file, dpi=300, bbox_inches="tight")  # Save as PNG
    plt.close(fig)  # Close the plot to free memory

    print(f"✅ Saved coastline debug plot: {output_file}")

# 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 check_coastline_continuity(coastline_gdf, geohash_code):
    """Checks if the coastline forms a continuous LineString and attempts to merge fragmented parts."""
    if coastline_gdf.empty:
        logger.warning(f"Geohash {geohash_code}: No coastline to check.")
        return False

    # Convert to a list of valid LineStrings
    coastline_list = [geom for geom in coastline_gdf.geometry if geom and geom.geom_type in ["LineString", "MultiLineString"]]

    if not coastline_list:
        logger.error(f"Geohash {geohash_code}: No valid LineString geometries found.")
        return False

    # Merge the coastline segments (ensuring not empty)
    merged_coastlines = group_continuous_coastline(coastline_list)

    if not merged_coastlines:
        logger.error(f"Geohash {geohash_code}: Merging failed. No valid coastlines found.")
        merged_coastline = MultiLineString([])  # Return an empty MultiLineString
    else:
        merged_coastline = max(merged_coastlines, key=lambda ls: ls.length)

    # Ensure the result is valid
    if merged_coastline.is_empty:
        logger.error(f"Geohash {geohash_code}: Merged coastline is empty.")
        return False

    if isinstance(merged_coastline, LineString):
        logger.info(f"Geohash {geohash_code}: Coastline is continuous.")
        plot_coastline_segments(geohash_code, merged_coastline)
        return True

    if isinstance(merged_coastline, MultiLineString):
        logger.warning(f"Geohash {geohash_code}: Coastline remains fragmented ({len(merged_coastline.geoms)} parts).")
        plot_coastline_segments(geohash_code, merged_coastline)
        return False

    logger.error(f"Geohash {geohash_code}: Unexpected coastline geometry ({merged_coastline.geom_type}).")
    return False

def process_geohash_files(coastline_gdf, geohash_files):
    """Processes each geohash file: updates coastline and adds water polygons."""
    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()

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

        # Merge continuous coastline segments
        merged_coastlines = group_continuous_coastline(new_coastline.geometry)
        if merged_coastlines:
            merged_coastline = max(merged_coastlines, key=lambda ls: ls.length)  # Pick the longest coastline
        else:
            merged_coastline = None

        if merged_coastline is None:
            logger.warning(f"Geohash {geohash_code}: Coastline merging failed.")
            continue

        logger.info(f"Adding water polygon for {geohash_code}...")

        # Create a full water polygon covering the entire geohash tile
        tile_polygon = bbox

        # Split the tile polygon using the coastline
        water_polygon = split_polygon_by_coastline(tile_polygon, merged_coastline)

        if water_polygon:
            # Assign water tags
            water_gdf = gpd.GeoDataFrame(
                {"geometry": [water_polygon], "tags": [{"natural": "water", "water": "sea"}]},
                crs="EPSG:4326"
            )

            # Append the water polygon to the updated geohash file
            updated_gdf = gpd.GeoDataFrame(pd.concat([filtered_gdf, new_coastline, water_gdf], ignore_index=True), crs="EPSG:4326")
            logger.info(f"Water polygon added for {geohash_code}.")
        else:
            logger.warning(f"Could not split water polygon for {geohash_code}. Skipping water addition.")
            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 geohash file {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()