import fused
from fused.core._fs_utils import is_non_empty_dir


@fused.udf(cache_max_age=0)
def udf_divide(
    src_path: str,
    chunk_id: int,
    x_chunks: int,
    y_chunks: int,
    tmp_path: str,
    res: int,
    k_ring: int,
    parent_offset: int,
    file_res: int,
):
    # define UDF that imports the helper function inside the UDF
    from job2.partition.raster_to_h3 import udf_divide as run_udf_divide

    run_udf_divide(
        src_path,
        chunk_id,
        x_chunks,
        y_chunks,
        tmp_path,
        res=res,
        k_ring=k_ring,
        parent_offset=parent_offset,
        file_res=file_res,
    )


def run_divide(
    src_path: str,
    tmp_path: str,
    x_chunks: int,
    y_chunks: int,
    chunk_ids: list[int] = None,
    res: int = 12,
    k_ring: int = 1,
    parent_offset: int = 1,
    file_res: int = 2,
    debug_mode: bool = False,
    **kwargs,
):
    """
    Chunk up the input raster file, extract pixel values and assign to H3 cells.

    Parameters
    ----------
    src_path : str
        Path to the input raster file.
    tmp_path : str
        Path to store intermediate files.
    x_chunks : int
        Number of chunks in the x direction.
    y_chunks : int
        Number of chunks in the y direction.
    chunk_ids : list of int, default None
        List of chunk IDs to process, default to all.
    res : int
        The resolution at which to assign the pixel values to H3 cells.
    k_ring : int
        The k-ring distance at resolution `res` to which the pixel value
        is assigned (in addition to the center cell).
    parent_offset : int
        Offset to parent resolution to which to assign the pixel values
        and counts
    file_res : int
        The H3 resolution to chunk the resulting files of the Parquet dataset

    """

    params = {
        "src_path": src_path,
        "tmp_path": tmp_path,
        "x_chunks": x_chunks,
        "y_chunks": y_chunks,
        #
        "res": res,
        "k_ring": k_ring,
        "parent_offset": parent_offset,
        "file_res": file_res,
    }

    if chunk_ids is None:
        chunk_ids = range(x_chunks * y_chunks)
        if debug_mode:
            # avoid creating huge submit_params list in case of tiny target_chunk_size
            chunk_ids = range(2)

    submit_params = [{"chunk_id": i} for i in chunk_ids]
    if debug_mode:
        submit_params = submit_params[:2]

    pool = fused.submit(
        udf_divide,
        submit_params,
        **params,
        collect=False,
        **kwargs,
    )
    return pool


@fused.udf(cache_max_age=0)
def udf_combine(
    file_id: int,
    tmp_path: str,
    output_path: str,
    chunk_res: int = 3,
):
    # define UDF that imports the helper function inside the UDF
    from job2.partition.raster_to_h3 import udf_combine as run_udf_combine

    run_udf_combine(
        file_id,
        tmp_path,
        output_path,
        chunk_res=chunk_res,
    )


def run_combine(
    tmp_path: str,
    file_ids: list[str],
    output_path: str,
    chunk_res: int = 3,
    debug_mode: bool = False,
    **kwargs,
):
    """
    Combine chunks per file_id H3 cell and recalculate data.

    Parameters
    ----------
    tmp_path : str
        Path where the intermediate files to combine are stored.
    output_path : str
        Path for the resulting Parquet dataset.
    chunk_res : int
        The H3 resolution to chunk the row groups within each file of the Parquet dataset

    """
    params = {
        "tmp_path": tmp_path,
        "output_path": output_path,
        #
        "chunk_res": chunk_res,
    }
    submit_params = [{"file_id": i} for i in file_ids]

    if debug_mode:
        submit_params = submit_params[:2]

    pool = fused.submit(
        udf_combine,
        submit_params,
        **params,
        collect=False,
        **kwargs,
    )
    return pool


def _cleanup_tmp_files(tmp_path: str, remove_tmp_files: bool):
    if remove_tmp_files:
        if tmp_path.startswith("file://"):
            import shutil

            shutil.rmtree(tmp_path.replace("file://", ""), ignore_errors=True)
        else:
            orig = fused.options.request_timeout
            fused.options.request_timeout = 10
            fused.api.delete(tmp_path)
            fused.options.request_timeout = orig


def _list_tmp_file_ids(tmp_path: str):
    if tmp_path.startswith("file://"):
        from pathlib import Path

        path = Path(tmp_path.replace("file://", ""))
        file_ids = [p.name for p in path.iterdir() if p.is_dir()]
    else:
        orig = fused.options.request_timeout
        fused.options.request_timeout = 10
        file_ids = [path.strip("/").split("/")[-1] for path in fused.api.list(tmp_path)]
        fused.options.request_timeout = orig
    return file_ids


def run_ingest_raster_to_h3(
    src_path: str,
    output_path: str,
    target_chunk_size: int = 10_000_000,
    res: int = 12,
    k_ring: int = 1,
    parent_offset: int = 1,
    chunk_res: int = 3,
    file_res: int = 2,
    # ops=["divide", "combine", "sample"],
    # n_jobs_divide=20,
    # n_jobs_combine=5,
    debug_mode: bool = False,
    remove_tmp_files: bool = True,
    cache_id: int = 0,
    divide_kwargs={},
    combine_kwargs={},
    **kwargs,
):
    """
    Run the raster to H3 ingestion process.

    This process involves multiple steps:
    - extract pixels values and assign to H3 cells in chunks (divide step)
    - combine the chunks per file and prepare metadata (combine step)
    - create the metadata `_sample` file and overviews files

    Parameters
    ----------
    src_path : str
        Path to the input raster file.
    output_path : str
        Path for the resulting Parquet dataset.
    target_chunk_size : int
        The approximate number of pixel values to process per chunk in
        the first "divide" step.
    res : int
        The resolution at which to assign the pixel values to H3 cells.
    k_ring : int
        The k-ring distance at resolution `res` to which the pixel value
        is assigned (in addition to the center cell).
    parent_offset : int
        Offset to parent resolution (relative to `res`) to which to assign
        the pixel values and counts.
    file_res : int
        The H3 resolution to chunk the resulting files of the Parquet dataset
    chunk_res : int
        The H3 resolution to chunk the row groups within each file of the
        Parquet dataset
    debug_mode : bool, default False
        If True, run only the first two chunks for debugging purposes.
    remove_tmp_files : bool, default True
        If True, remove the temporary files after ingestion is complete.
    cache_id : int, default 0
        An identifier to use for the temporary path to avoid collisions
         when running multiple ingestions simultaneously.
    divide_kwargs : dict
        Additional keyword arguments to pass to `fused.submit` for the divide step.
    combine_kwargs : dict
        Additional keyword arguments to pass to `fused.submit` for the combine step.
    **kwargs
        Additional keyword arguments to pass to `fused.submit` for
        both the divide and combine steps. Keys specified here are further
        overridden by those in `divide_kwargs` and `combine_kwargs` respectively.

    The divide and combine steps are run in parallel using fused.submit. By
    default, the function will first attempt to run this using "realtime"
    instances, and retry any failed runs using "large" instances.

    You can override this behavior by specifying the `engine`, `instance_type`,
    `max_workers`, `n_processes_per_worker`, etc parameters as additional
    keyword arguments to this function (`**kwargs`). If you want to specify
    those per step, use `divide_kwargs` and `combine_kwargs`.
    For example, to run everything locally on the same machine where this
    function runs, use:

        run_ingest_raster_to_h3(..., engine="local")

    To run the divide step on realtime and the combine step on medium
    instance, you could do:

        run_ingest_raster_to_h3(...,
            divide_kwargs={"instance_type": "realtime", "max_workers": 256, "max_retry": 1},
            combine_kwargs={"instance_type": "medium", "max_workers": 5, n_processes_per_worker": 2},
        )
    """
    import datetime

    import numpy as np
    import rasterio

    try:
        from job2.partition.raster_to_h3 import (
            udf_overview,
            udf_sample,
        )
    except ImportError:
        raise RuntimeError(
            "The ingestion functionality can only be run using the remote engine"
        )

    result_divide = None
    result_combine = None

    print(f"Starting ingestion process for {src_path}\n")
    start_time = datetime.datetime.now()

    # Verify that the output path is empty
    if is_non_empty_dir(output_path):
        raise ValueError(f"Output path {output_path} is not empty.")

    # Construct path for intermediate results
    src_path_part = src_path.replace("/", "_").replace(":", "_").replace(".", "_")
    if output_path.startswith("file:///"):
        # create local tmp path
        import tempfile

        local_tmp_dir = tempfile.gettempdir()
        tmp_path = f"file://{local_tmp_dir}/fused-tmp/tmp/{src_path_part}-{cache_id}/"
    else:
        api = fused.api.FusedAPI()
        tmp_path = api._resolve(f"fd://fused-tmp/tmp/{src_path_part}-{cache_id}/")
    if is_non_empty_dir(tmp_path):
        raise ValueError(
            f"Temporary path {tmp_path} is not empty."  # Specify cleanup_tmp_files=True to force remove old files, or clean up manually."
        )

    print(f"-- Using {tmp_path=}")

    ###########################################################################
    # Step one: extracting pixel values and converting to hex divided in chunks

    print("\nRunning divide step")
    start_divide_time = datetime.datetime.now()

    # determine number of chunks based on target chunk size
    with rasterio.open(src_path) as src:
        meta = src.meta

    x_chunks = max(round(meta["width"] / np.sqrt(target_chunk_size)), 1)
    y_chunks = max(round(meta["height"] / np.sqrt(target_chunk_size)), 1)

    divide_run_kwargs = dict(
        src_path=src_path,
        tmp_path=tmp_path,
        x_chunks=x_chunks,
        y_chunks=y_chunks,
        res=res,
        k_ring=k_ring,
        parent_offset=parent_offset,
        file_res=file_res,
        debug_mode=debug_mode,
    )
    divide_submit_kwargs = {"max_retry": 0} | kwargs | divide_kwargs

    # Run the actual divide step
    if not (
        divide_submit_kwargs.get("engine", None) == "local"
        or "instance_type" in divide_submit_kwargs
    ):
        # default logic: try realtime and fallback to large instance
        divide_submit_kwargs["instance_type"] = "realtime"
        result_divide = run_divide(
            **divide_run_kwargs,
            **divide_submit_kwargs,
        )
        result_divide.wait()
        if not result_divide.all_succeeded():
            errors = result_divide.errors()
            print(
                f"-- Divide step failed on realtime instances ({len(errors)} out of "
                f"{result_divide.n_jobs} runs, first error: {next(iter(errors.values()))}), "
                "retrying failed runs on large instances"
            )
            divide_run_kwargs["chunk_ids"] = list(errors.keys())
            divide_submit_kwargs["instance_type"] = "large"
            result_divide = run_divide(
                **divide_run_kwargs,
                **divide_submit_kwargs,
            )
            result_divide.wait()
    else:
        # otherwise run once with user provided configuration
        result_divide = run_divide(
            **divide_run_kwargs,
            **divide_submit_kwargs,
        )
        result_divide.wait()

    end_divide_time = datetime.datetime.now()
    if not result_divide.all_succeeded():
        print("\nDivide step failed!")
        try:
            _cleanup_tmp_files(tmp_path, remove_tmp_files)
        except Exception:
            pass
        return result_divide, result_combine
    print(f"-- Done divide! (took {end_divide_time - start_divide_time})")

    ###########################################################################
    # Step two: combining the chunks per file (resolution 2) and preparing
    # metadata and overviews

    print("\nRunning combine step")

    # list available file_ids from the previous step
    file_ids = _list_tmp_file_ids(tmp_path)
    print(f"-- processing {len(file_ids)} file_ids")

    combine_run_kwargs = dict(
        tmp_path=tmp_path,
        file_ids=file_ids,
        output_path=output_path,
        chunk_res=chunk_res,
    )
    combine_submit_kwargs = {"max_retry": 0} | kwargs | combine_kwargs

    # Run the actual combine step
    if not (
        combine_submit_kwargs.get("engine", None) == "local"
        or "instance_type" in combine_submit_kwargs
    ):
        # default logic: try realtime and fallback to large instance
        combine_submit_kwargs["instance_type"] = "realtime"
        result_combine = run_combine(
            **combine_run_kwargs,
            **combine_submit_kwargs,
        )
        result_combine.wait()
        if not result_combine.all_succeeded():
            errors = result_combine.errors()
            print(
                f"-- Combine step failed on realtime instances ({len(errors)} out of "
                f"{result_combine.n_jobs} runs, first error: {next(iter(errors.values()))}), "
                "retrying failed runs on large instances"
            )
            combine_run_kwargs["file_ids"] = [file_ids[i] for i in errors.keys()]
            combine_submit_kwargs["instance_type"] = "large"
            result_combine = run_combine(
                **combine_run_kwargs,
                **combine_submit_kwargs,
            )
            result_combine.wait()
    else:
        # otherwise run once with user provided configuration
        result_combine = run_combine(
            **combine_run_kwargs,
            **combine_submit_kwargs,
        )
        result_combine.wait()

    end_combine_time = datetime.datetime.now()
    if not result_combine.all_succeeded():
        print("\nCombine step failed!")
        _cleanup_tmp_files(tmp_path, remove_tmp_files)
        return result_divide, result_combine
    print(f"-- Done combine! (took {end_combine_time - end_divide_time})")

    ###########################################################################
    # Step 3: combining the metadata and overview tmp files

    print("\nRunning sample step")

    @fused.udf(cache_max_age=0)
    def udf_sample(tmp_path: str, output_path: str):
        from job2.partition.raster_to_h3 import udf_sample as run_udf_sample

        return run_udf_sample(tmp_path, output_path)

    sample_file = fused.run(
        udf_sample,
        tmp_path=tmp_path,
        output_path=output_path,
        verbose=False,
        engine=kwargs.get("engine", None),
    )
    end_sample_time = datetime.datetime.now()
    print(f"-- Written: {sample_file}")
    print(f"-- Done sample! (took {end_sample_time - end_combine_time})")

    print("\nRunning overview step")

    @fused.udf(cache_max_age=0)
    def udf_overview(tmp_path: str, output_path: str, res: int):
        from job2.partition.raster_to_h3 import udf_overview as run_udf_overview

        return run_udf_overview(tmp_path, output_path, res=res)

    for res in [3, 4, 5, 6]:
        overview_file = fused.run(
            udf_overview,
            tmp_path=tmp_path,
            output_path=output_path,
            res=res,
            verbose=False,
            engine=kwargs.get("engine", None),
        )
        print(f"-- Written: {overview_file}")

    end_overview_time = datetime.datetime.now()
    print(f"-- Done overview! (took {end_overview_time - end_sample_time})")

    # remove tmp files
    _cleanup_tmp_files(tmp_path, remove_tmp_files)

    print(f"\nIngestion process done! (took {datetime.datetime.now() - start_time})")

    return result_divide, result_combine
