import argparse
import matplotlib
matplotlib.use("Agg")  # Use Agg backend for non-GUI environments

import geopandas as gpd
import matplotlib.pyplot as plt
from pathlib import Path

def get_color(tags, geom_type):
    """
    Determine the color based on the tags and geometry type.
    """
    if isinstance(tags, dict):
        if "natural" in tags:
            if tags["natural"] == "water":
                return "blue" if geom_type == "Polygon" else "darkblue"
            if tags["natural"] in ["brush", "forest"]:
                return "darkgreen"
            if tags["natural"] == "coastline":
                return "red"
        if "landuse" in tags:
            if tags["landuse"] == "farmland":
                return "green"
            if tags["landuse"] in ["residential", "commercial"]:
                return "orange"
            if tags["landuse"] == "industrial":
                return "darkgrey"
            return "yellow"
        if "military" in tags:
            return "lightred"
        if "waterway" in tags:
            return "darkblue"
        if "railway" in tags:
            return "black"
    return "grey"  # Default color

def plot_geometries(parquet_file: str, save_as: str):
    """
    Plots the geometries inside a GeoParquet file and saves the image instead of showing it.

    Args:
        parquet_file (str): Path to the Parquet file containing geometries.
        save_as (str): Path to save the output plot image.
    """
    parquet_path = Path(parquet_file)
    if not parquet_path.exists():
        raise FileNotFoundError(f"❌ File '{parquet_file}' not found.")

    # Load the Parquet file into a GeoDataFrame
    gdf = gpd.read_parquet(parquet_path)

    # Ensure it has a valid geometry column
    if "geometry" not in gdf.columns:
        raise ValueError("❌ No 'geometry' column found in the file.")
    
    # Drop rows with empty geometries
    gdf = gdf[gdf.geometry.notnull()]
    
    # Create a figure
    fig, ax = plt.subplots(figsize=(8, 8))
    
    # Classify colors efficiently
    gdf["color"] = gdf.apply(lambda row: get_color(row.tags, row.geometry.geom_type), axis=1)
    
    # Separate by geometry types for efficient plotting
    for geom_type, sub_gdf in gdf.groupby(gdf.geometry.geom_type):
        sub_gdf.plot(ax=ax, edgecolor="black" if geom_type in ["Polygon", "MultiPolygon"] else None,
                     facecolor=sub_gdf["color"] if geom_type in ["Polygon", "MultiPolygon"] else None,
                     color=sub_gdf["color"] if geom_type in ["LineString", "MultiLineString"] else None,
                     alpha=0.5, linewidth=2 if geom_type in ["LineString", "MultiLineString"] else None)
    
    # Add labels
    ax.set_title(f"Geometries in {parquet_file}")
    ax.set_xlabel("Longitude")
    ax.set_ylabel("Latitude")
    
    # Save the figure instead of showing it
    plt.savefig(save_as, dpi=300)
    print(f"✅ Plot saved as {save_as}")

def main():
    parser = argparse.ArgumentParser(description="Plot geometries from a GeoParquet file and save as an image.")
    parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.")
    parser.add_argument("-o", "--output", required=True, help="Path to save the output plot image.")
    
    args = parser.parse_args()

    try:
        plot_geometries(args.input, args.output)
    except Exception as e:
        print(f"❌ Error: {e}")

if __name__ == "__main__":
    main()