import arcpy
import geopandas as gpd
import os


def raster2polygon(in_raster, out_polygon, verbose=True):
    """
    Convert a raster to a polygon shapefile, keeping only the non-null parts.
    NOTE: This function is using ArcPy. For installing arcpy, use the following command:
    'from abja_utils.gis import install_arcpy'
    'install_arcpy()'
    Args:
        in_raster (str): Path to the input raster file.
        out_polygon (str): Path to the output polygon shapefile.
    Returns:
        None
    """

    arcpy.env.overwriteOutput = True

    in_raster = os.path.abspath(in_raster)
    if not os.path.exists(in_raster):
        raise FileNotFoundError(f"Input raster does not exist: {in_raster}")
    output_folder = os.path.abspath(os.path.dirname(out_polygon)) 
    os.makedirs(output_folder, exist_ok=True)
    out_polygon = os.path.abspath(out_polygon)
    tmp_polygon = out_polygon.replace(".shp", "_tmp.shp")

    if verbose:
        print(f"Input raster: {in_raster}")
        print(f"Output polygon: {out_polygon}")
        print("Converting raster to integer...")
    int_raster = arcpy.sa.Con(
        in_conditional_raster=in_raster,
        in_true_raster_or_constant=1,
        in_false_raster_or_constant=0,
        where_clause="VALUE IS NOT NULL"
    )

    # Convert raster to polygon with 0 and 1 in Value column
    if verbose:
        print("Converting raster to polygon") 
    arcpy.conversion.RasterToPolygon(
        in_raster=int_raster,
        out_polygon_features=tmp_polygon,
        simplify="NO_SIMPLIFY",
        raster_field="Value",
        create_multipart_features="SINGLE_OUTER_PART",
        max_vertices_per_feature=None
    )

    # Remove null parts
    if verbose:
        print("Cleaning up the polygon")
    gdf = gpd.read_file(tmp_polygon)
    mask = gdf["gridcode"] == 1
    gdf = gdf[mask]
    gdf.to_file(out_polygon)
    files_to_delete = [f for f in os.listdir(output_folder) if f.startswith(os.path.basename(tmp_polygon).replace(".shp", ""))]
    for f in files_to_delete:
        os.remove(os.path.join(output_folder, f))
    if verbose:
        print("Finished!")