"""Utility to hold the file-level pipeline execution plan."""

from __future__ import annotations

import pickle
import re
from dataclasses import dataclass, field

import hats.pixel_math.healpix_shim as hp
import numpy as np
from hats import pixel_math
from hats.io import file_io
from hats.pixel_math.healpix_pixel import HealpixPixel
from hats.pixel_math.sparse_histogram import HistogramAggregator, SparseHistogram
from numpy import frombuffer
from upath import UPath

import hats_import.file_io as import_io
from hats_import.pipeline_resume_plan import PipelineResumePlan


@dataclass
class ResumePlan(PipelineResumePlan):
    """Container class for holding the state of each file in the pipeline plan."""

    input_paths: list[UPath] = field(default_factory=list)
    """Resolved list of all files that will be used in the importer"""
    map_files: list[tuple[str, str]] = field(default_factory=list)
    """List of files (and job keys) that have yet to be mapped"""
    split_keys: list[tuple[str, str]] = field(default_factory=list)
    """Set of files (and job keys) that have yet to be split"""
    destination_pixel_map: dict[HealpixPixel, int] | None = None
    """Destination pixels and their expected final count"""
    threshold_mode: str = "row_count"
    """Which mode to use for partitioning: 'row_count' or 'mem_size'.
    Determines whether to create additional mem_size histogram."""
    should_run_mapping: bool = True
    should_run_splitting: bool = True
    should_run_reducing: bool = True
    should_run_finishing: bool = True

    MAPPING_STAGE = "mapping"
    SPLITTING_STAGE = "splitting"
    REDUCING_STAGE = "reducing"

    ROW_COUNT_HISTOGRAM_BINARY_FILE = "row_count_mapping_histogram.npz"
    ROW_COUNT_HISTOGRAMS_DIR = "row_count_histograms"
    MEM_SIZE_HISTOGRAM_BINARY_FILE = "mem_size_mapping_histogram.npz"
    MEM_SIZE_HISTOGRAMS_DIR = "mem_size_histograms"

    ALIGNMENT_FILE = "alignment.pickle"

    # pylint: disable=too-many-arguments
    def __init__(
        self,
        resume: bool = True,
        progress_bar: bool = True,
        simple_progress_bar: bool = False,
        input_paths=None,
        tmp_path=None,
        tmp_base_path: UPath | None = None,
        delete_resume_log_files: bool = True,
        delete_intermediate_parquet_files: bool = True,
        run_stages: list[str] | None = None,
        import_args=None,
    ):
        if import_args:
            super().__init__(pipeline_name="catalog", **import_args.resume_kwargs_dict())
            if import_args.debug_stats_only:
                run_stages = ["mapping", "finishing"]
            self.input_paths = import_args.input_paths

            # Set threshold_mode based on byte_pixel_threshold
            if hasattr(import_args, "byte_pixel_threshold") and import_args.byte_pixel_threshold is not None:
                self.threshold_mode = "mem_size"
        else:
            super().__init__(
                resume=resume,
                progress_bar=progress_bar,
                simple_progress_bar=simple_progress_bar,
                tmp_path=tmp_path,
                tmp_base_path=tmp_base_path,
                delete_resume_log_files=delete_resume_log_files,
                delete_intermediate_parquet_files=delete_intermediate_parquet_files,
            )
            self.input_paths = input_paths
        self.gather_plan(run_stages)

    def gather_plan(self, run_stages: list[str] | None = None):
        """Initialize the plan."""
        with self.print_progress(total=4, stage_name="Planning") as step_progress:
            ## Make sure it's safe to use existing resume state.
            super().safe_to_resume()
            step_progress.update(1)

            ## Validate existing resume state.
            ## - if a later stage is complete, the earlier stages should be complete too.
            mapping_done = self.done_file_exists(self.MAPPING_STAGE)
            splitting_done = self.done_file_exists(self.SPLITTING_STAGE)
            reducing_done = self.done_file_exists(self.REDUCING_STAGE)

            if reducing_done and (not mapping_done or not splitting_done):
                raise ValueError("mapping and splitting must be complete before reducing")
            if splitting_done and not mapping_done:
                raise ValueError("mapping must be complete before splitting")
            step_progress.update(1)

            ## Figure out which stages we should run, based on requested `run_stages`
            self.should_run_mapping = not mapping_done
            self.should_run_splitting = not splitting_done
            self.should_run_reducing = not reducing_done
            self.should_run_finishing = True

            if run_stages:
                self.should_run_mapping &= self.MAPPING_STAGE in run_stages
                self.should_run_splitting &= self.SPLITTING_STAGE in run_stages
                self.should_run_reducing &= self.REDUCING_STAGE in run_stages
                self.should_run_finishing = "finishing" in run_stages

            ## Validate that we're operating on the same file set as the previous instance.
            self.input_paths = self.check_original_input_paths(self.input_paths)
            step_progress.update(1)

            ## Gather keys for execution.
            if self.should_run_mapping:
                self.map_files = self.get_remaining_map_keys()
                file_io.make_directory(
                    import_io.append_paths_to_pointer(self.tmp_path, self.ROW_COUNT_HISTOGRAMS_DIR),
                    exist_ok=True,
                )
                # If using mem_size thresholding, gather those keys too.
                if self.threshold_mode == "mem_size":
                    self.get_remaining_map_keys(which_histogram="mem_size")
                    file_io.make_directory(
                        import_io.append_paths_to_pointer(self.tmp_path, self.MEM_SIZE_HISTOGRAMS_DIR),
                        exist_ok=True,
                    )
            if self.should_run_splitting:
                if not (mapping_done or self.should_run_mapping):
                    raise ValueError("mapping must be complete before splitting")

                self.split_keys = self.get_remaining_split_keys()
                file_io.make_directory(
                    import_io.append_paths_to_pointer(self.tmp_path, self.SPLITTING_STAGE),
                    exist_ok=True,
                )
            if self.should_run_reducing:
                ## We don't pre-gather the plan for the reducing keys.
                ## It requires the full destination pixel map.
                if not (splitting_done or self.should_run_splitting):
                    raise ValueError("splitting must be complete before reducing")

                file_io.make_directory(
                    import_io.append_paths_to_pointer(self.tmp_path, self.REDUCING_STAGE),
                    exist_ok=True,
                )
            step_progress.update(1)

    def get_remaining_map_keys(self, which_histogram: str = "row_count"):
        """Gather remaining keys, dropping successful mapping tasks from histogram names.

        Args:
            which_histogram (str): Which histogram to check for completed tasks, either 'row_count'
                or 'mem_size'. Defaults to 'row_count'.

        Returns:
            list of tuple: The mapping keys *not* found in files like /resume/path/mapping_key.npz,
                along with their corresponding input file paths.

        Raises:
            ValueError: If `which_histogram` is not recognized, or if which_histogram is
                'mem_size' but the threshold_mode is not 'mem_size'.
        """
        if which_histogram == "row_count":
            prefix = file_io.get_upath(self.tmp_path) / self.ROW_COUNT_HISTOGRAMS_DIR
        elif which_histogram == "mem_size" and self.threshold_mode == "mem_size":
            prefix = file_io.get_upath(self.tmp_path) / self.MEM_SIZE_HISTOGRAMS_DIR
        elif which_histogram == "mem_size":
            raise ValueError("Cannot get remaining mem_size map keys when threshold_mode is not 'mem_size'.")
        else:
            raise ValueError(f"Unrecognized which_histogram value: {which_histogram}")

        map_file_pattern = re.compile(r"map_(\d+).npz")
        done_indexes = [
            int(match.group(1))
            for path in prefix.glob("*.npz")
            if (match := map_file_pattern.match(path.name))
        ]
        remaining_indexes = list(set(range(0, len(self.input_paths))) - (set(done_indexes)))
        return [(f"map_{key}", self.input_paths[key]) for key in remaining_indexes]

    def read_histogram(self, healpix_order, which_histogram: str = "row_count"):
        """Returns a histogram with the specified Healpix order's shape.

        This method attempts the following steps in order:
        1. Tries to locate and return a combined histogram.
        2. If a combined histogram is unavailable, combines partial histograms to create one.
        3. If no partial histograms are found, returns an empty histogram.

        Args:
            healpix_order (int): The desired Healpix order for the histogram.
            which_histogram (str): Which histogram to read, either "row_count" or "mem_size".
                Defaults to "row_count".


        Returns:
            numpy.ndarray: A one-dimensional array representing the histogram with the
                specified Healpix order.

        Raises:
            RuntimeError: If there are incomplete mapping stages.
            ValueError: If the histogram from the previous execution is incompatible with
                the highest Healpix order, or if `which_histogram` is invalid.
        """
        if which_histogram == "row_count":
            histogram_binary_file = self.ROW_COUNT_HISTOGRAM_BINARY_FILE
            histogram_directory = self.ROW_COUNT_HISTOGRAMS_DIR
        elif which_histogram == "mem_size" and self.threshold_mode == "mem_size":
            histogram_binary_file = self.MEM_SIZE_HISTOGRAM_BINARY_FILE
            histogram_directory = self.MEM_SIZE_HISTOGRAMS_DIR
        elif which_histogram == "mem_size":
            raise ValueError("Cannot read mem_size histogram when threshold_mode is not 'mem_size'.")
        else:
            raise ValueError(f"Unrecognized which_histogram value: {which_histogram}")

        file_name = import_io.append_paths_to_pointer(self.tmp_path, histogram_binary_file)

        # If no file, read the histogram from partial histograms and combine.
        if not file_name.exists():
            remaining_map_files = self.get_remaining_map_keys(which_histogram=which_histogram)
            if len(remaining_map_files) > 0:
                raise RuntimeError(f"{len(remaining_map_files)} map stages did not complete successfully.")
            histogram_files = import_io.find_files_matching_path(self.tmp_path, histogram_directory, "*.npz")
            aggregate_histogram = HistogramAggregator(healpix_order)
            for partial_file_name in histogram_files:
                partial = SparseHistogram.from_file(partial_file_name)
                aggregate_histogram.add(partial)

            file_name = import_io.append_paths_to_pointer(self.tmp_path, histogram_binary_file)
            with file_name.open("wb+") as file_handle:
                file_handle.write(aggregate_histogram.full_histogram)
            if self.delete_resume_log_files:
                file_io.remove_directory(
                    import_io.append_paths_to_pointer(self.tmp_path, histogram_directory),
                    ignore_errors=True,
                )

        with file_name.open("rb") as file_handle:
            full_histogram = frombuffer(file_handle.read(), dtype=np.int64)

        if len(full_histogram) != hp.order2npix(healpix_order):
            raise ValueError(
                "The histogram from the previous execution is incompatible with "
                + "the highest healpix order. To start the importing pipeline "
                + "from scratch with the current order set `resume` to False."
            )
        return full_histogram

    @classmethod
    def partial_histogram_file(cls, tmp_path, mapping_key: str, which_histogram: str = "row_count"):
        """File name for writing a histogram file to a special intermediate directory.

        As a side effect, this method may create the special intermediate directory.

        Args:
            tmp_path (str): where to write intermediate resume files.
            mapping_key (str): unique string for each mapping task (e.g. "map_57")
            which_histogram (str): which histogram to write, either "row_count" or "mem_size".
                Defaults to "row_count".

        Returns:
            str: Full path to the partial histogram file.
        """
        if which_histogram == "row_count":
            histograms_dir = cls.ROW_COUNT_HISTOGRAMS_DIR
        elif which_histogram == "mem_size":
            histograms_dir = cls.MEM_SIZE_HISTOGRAMS_DIR
        else:
            raise ValueError(f"Unrecognized which_histogram value: {which_histogram}")

        file_io.make_directory(
            import_io.append_paths_to_pointer(tmp_path, histograms_dir),
            exist_ok=True,
        )
        return import_io.append_paths_to_pointer(tmp_path, histograms_dir, f"{mapping_key}.npz")

    def get_remaining_split_keys(self):
        """Gather remaining keys, dropping successful split tasks from done file names.

        Returns:
            list of splitting keys *not* found in files like /resume/path/split_key.done
        """
        prefix = file_io.get_upath(self.tmp_path) / self.SPLITTING_STAGE
        split_file_pattern = re.compile(r"split_(\d+)_done")
        done_indexes = [int(split_file_pattern.match(path.name).group(1)) for path in prefix.glob("*_done")]
        remaining_indexes = list(set(range(0, len(self.input_paths))) - set(done_indexes))
        return [(f"split_{key}", self.input_paths[key]) for key in remaining_indexes]

    @classmethod
    def splitting_key_done(cls, tmp_path, splitting_key: str):
        """Mark a single splitting task as done

        Args:
            tmp_path (str): where to write intermediate resume files.
            splitting_key (str): unique string for each splitting task (e.g. "split_57")
        """
        cls.touch_key_done_file(tmp_path, cls.SPLITTING_STAGE, splitting_key)

    @classmethod
    def reducing_key_done(cls, tmp_path, reducing_key: str):
        """Mark a single reducing task as done

        Args:
            tmp_path (str): where to write intermediate resume files.
            reducing_key (str): unique string for each reducing task (e.g. "3_57")
        """
        cls.touch_key_done_file(tmp_path, cls.REDUCING_STAGE, reducing_key)

    def wait_for_mapping(self, futures):
        """Wait for mapping futures to complete."""
        self.wait_for_futures(futures, self.MAPPING_STAGE)
        remaining_map_items = self.get_remaining_map_keys()
        if len(remaining_map_items) > 0:
            raise RuntimeError("some map stages did not complete successfully.")
        self.touch_stage_done_file(self.MAPPING_STAGE)

    def get_alignment_file(
        self,
        raw_histogram,
        constant_healpix_order,
        highest_healpix_order,
        lowest_healpix_order,
        pixel_threshold,
        drop_empty_siblings,
        expected_total_rows,
        existing_pixels=None,
        raw_histogram_mem_size=None,
    ) -> UPath:
        """Get a pointer to the existing alignment file for the pipeline, or
        generate a new alignment using provided arguments.

        Args:
            raw_histogram (:obj:`np.array`): one-dimensional numpy array of long integers where the
                value at each index corresponds to the number of objects found at the healpix pixel.
            constant_healpix_order (int): if positive, use this as the order for
                all non-empty partitions. else, use remaining arguments.
            highest_healpix_order (int):  the highest healpix order (e.g. 5-10)
            lowest_healpix_order (int): the lowest healpix order (e.g. 1-5). specifying a lowest order
                constrains the partitioning to prevent spatially large pixels.
            pixel_threshold (int): the maximum number of objects allowed in a single pixel
            drop_empty_siblings (bool):  if 3 of 4 pixels are empty, keep only the non-empty pixel
            expected_total_rows (int): number of expected rows found in the dataset.
            existing_pixels (Sequence[tuple[int,int]]): the HEALPix pixels to include in the alignment
            raw_histogram_mem_size (:obj:`np.array`): one-dimensional numpy array of long integers
                where the value at each index corresponds to the memory size in bytes of objects
                found at the healpix pixel. Only required if threshold_mode is 'mem_size'.

        Returns:
            path to cached alignment file.
        """
        file_name = import_io.append_paths_to_pointer(self.tmp_path, self.ALIGNMENT_FILE)
        if not file_name.exists():
            # If existing_pixels, create an incremental alignment.
            if existing_pixels:
                alignment = pixel_math.generate_incremental_alignment(
                    raw_histogram,
                    existing_pixels=existing_pixels,
                    highest_order=highest_healpix_order,
                    lowest_order=lowest_healpix_order,
                    threshold=pixel_threshold,
                )
            # If constant_healpix_order is set, create a simple alignment.
            elif constant_healpix_order >= 0:
                alignment = self._generate_constant_healpix_order_alignment(
                    raw_histogram, constant_healpix_order
                )
            # Else, generate standard alignment based on thresholds.
            else:
                alignment = pixel_math.generate_alignment(
                    raw_histogram,
                    highest_order=highest_healpix_order,
                    lowest_order=lowest_healpix_order,
                    threshold=pixel_threshold,
                    drop_empty_siblings=drop_empty_siblings,
                    mem_size_histogram=raw_histogram_mem_size,
                )

            # Write alignment to file.
            with file_name.open("wb") as pickle_file:
                alignment = np.array([x if x is not None else [-1, -1, 0] for x in alignment], dtype=np.int64)
                pickle.dump(alignment, pickle_file)

        # Check that the destination pixel map (alignment file) matches expected total rows.
        if self.destination_pixel_map is None:
            with file_name.open("rb") as pickle_file:
                alignment = pickle.load(pickle_file)
            pixel_list = np.unique(alignment, axis=0)
            self.destination_pixel_map = {
                HealpixPixel(order, pix): row_count
                for (order, pix, row_count) in pixel_list
                if int(row_count) > 0
            }

        total_rows = sum(self.destination_pixel_map.values())
        if total_rows != expected_total_rows:
            raise ValueError(
                f"Number of rows ({total_rows}) does not match expectation ({expected_total_rows})"
            )

        return file_name

    def _generate_constant_healpix_order_alignment(self, raw_histogram_row_count, constant_healpix_order):
        """Generate alignment where all non-empty pixels are at the same healpix order.

        Args:
            raw_histogram_row_count (:obj:`np.array`): one-dimensional numpy array of long integers
                where the value at each index corresponds to the number of objects found at the
                healpix pixel.
            constant_healpix_order (int): the healpix order to assign to all non-empty pixels.

        Returns:
            np.array: alignment array where each entry is [order, pixel, row_count] or
                [order, pixel, row_count, mem_size] depending on threshold_mode.
        """
        alignment = np.full((len(raw_histogram_row_count), 3), [-1, -1, 0])
        for pixel_num, pixel_row_count_sum in enumerate(raw_histogram_row_count):
            alignment[pixel_num] = [
                constant_healpix_order,
                pixel_num,
                pixel_row_count_sum,
            ]
        return alignment

    def wait_for_splitting(self, futures):
        """Wait for splitting futures to complete."""
        self.wait_for_futures(futures, self.SPLITTING_STAGE)
        remaining_split_items = self.get_remaining_split_keys()
        if len(remaining_split_items) > 0:
            raise RuntimeError(f"{len(remaining_split_items)} split stages did not complete successfully.")
        self.touch_stage_done_file(self.SPLITTING_STAGE)

    def get_reduce_items(self):
        """Fetch a triple for each partition to reduce.

        Triple contains:

        - destination pixel (healpix pixel with both order and pixel)
        - number of rows expected for this pixel
        - reduce key (string of destination order+pixel)
        """
        if self.destination_pixel_map is None:
            raise RuntimeError("destination pixel map not provided for progress tracking.")

        reduced_pixels = self.read_done_pixels(self.REDUCING_STAGE)

        remaining_pixels = list(set(self.destination_pixel_map.keys()) - set(reduced_pixels))
        return [
            (hp_pixel, self.destination_pixel_map[hp_pixel], f"{hp_pixel.order}_{hp_pixel.pixel}")
            for hp_pixel in remaining_pixels
        ]

    def get_destination_pixels(self):
        """Create HealpixPixel list of all destination pixels."""
        if self.destination_pixel_map is None:
            raise RuntimeError("destination pixel map not known.")
        return list(self.destination_pixel_map.keys())

    def wait_for_reducing(self, futures):
        """Wait for reducing futures to complete."""
        self.wait_for_futures(futures, self.REDUCING_STAGE, fail_fast=True)
        remaining_reduce_items = self.get_reduce_items()
        if len(remaining_reduce_items) > 0:
            raise RuntimeError(f"{len(remaining_reduce_items)} reduce stages did not complete successfully.")
        self.touch_stage_done_file(self.REDUCING_STAGE)
