from upath import UPath

from olmoearth_shared.api.run.prediction_geometry import PredictionRequestFeature, PredictionResultCollection


class PostprocessInterfaceVector:
    """
    Interface describing how we post-process the vector results of either a Task or an entire Workflow.
    """

    def process_window(self, window_request: PredictionRequestFeature, window_result: PredictionResultCollection) -> PredictionResultCollection:
        """
        In the context of a single Window, combine the request geojson used to initialize the Window with the FeatureCollection generated by inference
        """
        raise NotImplementedError

    def process_partition(self, partition_window_results: list[PredictionResultCollection]) -> PredictionResultCollection:
        """
        Combine the results of process_window() for all windows in the partition, resulting in a single FeatureCollection for the partition
        """
        raise NotImplementedError

    def process_dataset(self, all_partitions_results: list[PredictionResultCollection]) -> PredictionResultCollection:
        """
        Combine the results of all partitions into a single FeatureCollection for the entire prediction request
        """
        raise NotImplementedError


class PostprocessInterfaceRaster:
    """
    Interface describing how we post-process the raster results of either a Task or an entire Workflow.
    """

    def process_window(self, window_request: PredictionRequestFeature, window_result_path: UPath) -> None:
        """
        In the context of a single Window, combine the request geojson used to initialize the Window with the FeatureCollection generated by inference.
        This might involve setting values in the GeoTIFF properties, applying a vector mask to the raster, and so on.
        """
        raise NotImplementedError

    def process_partition(self, partition_window_results: list[UPath], output_dir: UPath) -> list[UPath]:
        """
        Combine the GeoTIFFs from each window into a single geotiff for the whole partition.
        """
        raise NotImplementedError

    def process_dataset(self, all_partitions_result_paths: list[UPath], output_dir: UPath) -> list[UPath]:
        """
        Combine the GeoTIFFs from each window into a single geotiff for the whole partition.
        """
        raise NotImplementedError
