from osgeo import gdal, ogr
from tqdm import tqdm
import os
from .converter import convert_dat_to_tiff  # Import conversion function

def clipper(input_file, shp_file, output_folder, prefix=None, field="OBJECTID"):
    """
    Splits an input raster (TIFF or DAT) into smaller plots using a shapefile.
    
    Args:
        input_file (str): Path to the input raster file (.tiff or .dat).
        shp_file (str): Path to the shapefile.
        output_folder (str): Directory to save the extracted plots.
        prefix (str, optional): Prefix for output file names. Defaults to None.
        field (str): Shapefile attribute field used for naming the plots.
    """
    
    # Convert .dat to .tiff if necessary
    if input_file.lower().endswith(".dat"):
        tiff_file = convert_dat_to_tiff(input_file)
        print(f"Converted {input_file} to {tiff_file}")
    else:
        tiff_file = input_file  # Use directly if already TIFF
    
    # Open the TIFF dataset
    tiff_dataset = gdal.Open(tiff_file, gdal.GA_ReadOnly)
    if tiff_dataset is None:
        raise Exception(f"Failed to open raster file: {tiff_file}")

    # Open the shapefile
    shp_dataset = ogr.Open(shp_file)
    if shp_dataset is None:
        raise Exception(f"Failed to open shapefile: {shp_file}")

    layer = shp_dataset.GetLayer()

    # Process each feature in the shapefile
    for i in tqdm(range(layer.GetFeatureCount()), desc="Processing plots"):
        feature = layer.GetFeature(i)
        geometry = feature.GetGeometryRef()
        if geometry and geometry.IsValid() and geometry.GetArea() > 0:
            xmin, xmax, ymin, ymax = geometry.GetEnvelope()

            # Define clipping options
            options = gdal.WarpOptions(
                outputBounds=(xmin, ymin, xmax, ymax),
                format="GTiff",
                cutlineDSName=shp_file,
                cropToCutline=True,
                cutlineWhere=f"{field}='{feature.GetField(field)}'",
            )

            # Construct filename
            plot_id = feature.GetField(field)
            output_filename = os.path.join(
                output_folder,
                f"Plotid_{plot_id}{'_' + prefix if prefix else ''}.tif"
            )

            # Save the clipped image
            output = gdal.Warp(output_filename, tiff_dataset, options=options)

            if output is None:
                print(f"Failed to clip feature {i} with field {plot_id}")
        else:
            print(f"Invalid or empty geometry for feature {i}")

    # Close datasets
    tiff_dataset = None
    shp_dataset = None
