import stackstac
from rasterio import RasterioIOError
from time import sleep

class AutoParallelRioReaderWithRetry(stackstac.rio_reader.AutoParallelRioReader):
    def _get_ds(self) -> stackstac.rio_reader.SelfCleaningDatasetReader:
        """
        Retry
        """
        retries = 10
        retries_delay = 10

        for _ in range(retries):
            try:
                return stackstac.rio_reader.SelfCleaningDatasetReader(self.url, sharing=False)
            except RasterioIOError as ex:
                exception = ex
                dns_problem_condition = ("Could not resolve host" in str(ex))
                read_problem_condition = ("not recognized as a supported dataset name" in str(ex))
                if dns_problem_condition or read_problem_condition:
                    print("retrying")
                    sleep(retries_delay)
                    continue
                print(f"Failed to open {self.url} with exception {ex}")
                raise ex
        raise Exception(f"Failed to open {self.url} after {retries} retries with error {exception}")

import rasterio as rio
import gc
import time

with rio.Env():
        # Time the process 10 times
        for i in range(10):
                start = time.time()
                a = rio.DatasetReader("<path to to remote file>", sharing=False)
                print("DS", time.time() - start)
                a.close()
                del a
                gc.collect()

gdal_env = stackstac.rio_env.LayeredEnv(
    always=dict(
        GDAL_HTTP_MULTIRANGE="YES",  # unclear if this actually works
        GDAL_HTTP_MERGE_CONSECUTIVE_RANGES="YES",
        # ^ unclear if this works either. won't do much when our dask chunks are aligned to the dataset's chunks.
        CPL_VSIL_CURL_USE_HEAD="NO",
        CPL_VSIL_CURL_NON_CACHED="/vsicurl",
    ),
    open=dict(
        GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR",
        # ^ stop GDAL from requesting `.aux` and `.msk` files from the bucket (speeds up `open` time a lot)
        VSI_CACHE=True
        # ^ cache HTTP requests for opening datasets. This is critical for `ThreadLocalRioDataset`,
        # which re-opens the same URL many times---having the request cached makes subsequent `open`s
        # in different threads snappy.
    ),
    read=dict(
        VSI_CACHE=False
        # ^ *don't* cache HTTP requests for actual data. We don't expect to re-request data,
        # so this would just blow out the HTTP cache that we rely on to make repeated `open`s fast
        # (see above)
    ),
)


from rasterio import RasterioIOError
from time import sleep, time


class AutoParallelRioReaderWithRetry(stackstac.rio_reader.AutoParallelRioReader):

    def _sefefy_rasterio_operation(self, fn, *args, **kwargs):
        """
        Retry
        """
        retries = 10
        retries_delay = 10

        for i in range(retries):
            try:
                time_start = time()
                return fn(*args, **kwargs)
            except RasterioIOError as ex:
                exception = ex
                dns_problem_condition = ("Could not resolve host" in str(ex))
                read_problem_condition = ("not recognized as a supported" in str(ex))
                if dns_problem_condition or read_problem_condition:
                    print("retrying")
                    print(f"Time for attempt {i+1}: {time() - time_start}")
                    sleep(retries_delay)
                    continue
                print(f"Failed to open {self.url} with exception {ex}")
                raise ex
        raise Exception(f"Failed to open {self.url} after {retries} retries with error {exception}")

    def _get_ds(self) -> stackstac.rio_reader.SelfCleaningDatasetReader:
        """
        Retry
        """
        return self._sefefy_rasterio_operation(stackstac.rio_reader.SelfCleaningDatasetReader, self.url, sharing=False)

    def _reader_read(self, reader, window, **kwargs):
            return self._sefefy_rasterio_operation(
                reader.read,
                window=window,
                out_dtype=self.dtype,
                masked=True,
                # ^ NOTE: we always do a masked array, so we can safely apply scales and offsets
                # without potentially altering pixels that should have been the ``fill_value``
                **kwargs,
            )
....
items = ...
stack = stackstac.stack(items, gdal_env=gdal_env, reader=AutoParallelRioReaderWithRetry)