"""For iterative merging of pre-sorted files.
"""
from merge_sort import common
from pyaddons import utils
from pyaddons.flat_files import header as pyhead
import gzip
import heapq
import os
import time
# import pprint as pp


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class IterativeHeapqMerge(object):
    """Performs an interative merge sort over many files.

    Parameters
    ----------
    files : `list` of `str`
        One or more file names. If header is ``True`` then the first row of
        each file must be the header.
    key : `function`, optional, default: `NoneType`
        The function to provide the sort keys to ``heapq.merge``.
    max_files : `int`, optional, default: `16`
        The maximum number of files that can be open (and merged) at any
        one time. The default (16) is the same as unix sort.
    tmpdir : `str`, optional, default: `NoneType`
        The directory for temp merge files to be placed if the merge can't be
        performed in a single pass. The default is the system temp location.
    read_method : `function`, optional, default: `gzip.open`
        The method to use to open the input files for merging.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write any output files if the merge can't be
        completed in a single pass.
    read_mode : `str`, optional, default: `rt`
        The mode to use to open the input files for merging.
    write_mode : `str`, optional, default: `wt`
        The mode to use to write and intermediate files if the merge cannot be
        completed in a single pass.
    header : `bool`, optional, default: `True`
        Do the input files have a header, if so, the header of the first file
        is stored in the ``header`` attribute.
    check_header : `bool`, optional, default: `True`
        If ``header`` is ``True``, do you want to make sure that all the
        headers are the same in each of the input files?
    delete : `bool`, optional, default: `False`
        Delete the source files as they are merged?
    reverse : `bool`, optional, default: `False`
        Perform the sort in reverse.
    comment : `str`, optional, default: `NoneType`
        A comment character that might appear before the header/data row, these
        are skipped. These are not checked beyond the first header/data row.
    read_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``read_method``.
    write_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``write_method``.
    """

    TEMP_PREFIX = "heapq_merge_"
    """A prefix that will be added to any intermediate merge files should the
    merge not be completed in a single pass (`str`).
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, files, key=None, max_files=16, tmpdir=None,
                 read_method=gzip.open, write_method=gzip.open,
                 read_mode='rt', write_mode='wt', header=True,
                 check_header=True, delete=False, reverse=False,
                 comment=None, read_kwargs=None, write_kwargs=None):
        self._files = files
        self._key = key
        self._tmpdir = tmpdir
        self._read_method = read_method
        self._write_method = write_method
        self._read_mode = read_mode
        self._write_mode = write_mode
        self._read_kwargs = read_kwargs or dict()
        self._write_kwargs = write_kwargs or dict()
        self._has_header = header
        self._reverse = reverse
        self._delete = delete
        self._check_header = check_header
        self._comment = comment

        self._max_files = 0
        self.max_files = max_files

        # This will hold intermediate file names created during the merge
        # thee will be deleted at the end of the run.
        self._tempfiles = []
        self._remaining_files = []
        self._cur_time = time.time()
        self._temp_prefix = self.TEMP_PREFIX
        self.header = None
        self._final_merge = []

        # indicator to make sure we have opened the merger
        self._opened = False

        # Initialised to None so will be set up properly on a call to next()
        self._heapq = None

        # Will hold the iterations, number of files merged in each iteration
        # and time taken for each iteration.
        self._iterations = []

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __enter__(self):
        """Entry point for the context manager
        """
        self.open()
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __exit__(self, *args):
        """Exit point for the context manager

        Parameters
        ----------
        *args
            Any errors that are raised in the context manager, or  a tuple of
            3 `NoneType` if none were raised.
        """
        self.close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __iter__(self):
        """Initialise the iterator.
        """
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __next__(self):
        """Grab the next row, rows are returned in sort order.

        Returns
        -------
        row : `list`
            All the entries are assumed to be bytes apart.
        """
        try:
            row = next(self._heapq)
        except TypeError:
            # The first call so we initialise, although we make sure that the
            # heapq is NoneType, if not then something else may have happened.
            if self._heapq is not None:
                raise

            if self._opened is False:
                raise IOError("did you forget to open the merge object?")

            # Otherwise setup the sort and return the first row
            self._merge()
            self._final_merge_start = time.time()
            try:
                row = next(self._heapq)
            except StopIteration:
                self._log_final_iteration()
                raise
        except StopIteration:
            self._log_final_iteration()
            raise

        return row

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _log_final_iteration(self):
        """Log the timings for the final iteration.
        """
        try:
            iteration_no = self._iterations[-1][0]+1
        except IndexError:
            # no previous iterations it is doing it in one
            iteration_no = 1

        self._iterations.append(
            (
                iteration_no,
                len(self._final_merge),
                time.time() - self._final_merge_start,
                [i.name for i in self._final_merge]
            )
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def max_files(self):
        """Get the max open files files at one time (`int`)
        """
        return self._max_files

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @max_files.setter
    def max_files(self, max_files):
        """Set the max open files at one time but be > 1 (`int`)
        """
        max_files = int(max_files)
        if max_files < 2:
            raise ValueError(
                "must allow at least two files to be open at one time"
            )
        self._max_files = max_files

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def iterations(self):
        """Get the iterations used to merge the files (`list` of `tuple`).

        Each tuple is the iteration number, number of files, time for the
        iteration and files merged.
        """
        return self._iterations

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Open all the files that need merging and perform header checks.

        Notes
        -----
        Note that this will open and close all the files needed for merging to
        make sure they exist before the merge takes place. This is done as it
        can be time consuming, so we can fail due to header issues or
        FileNotFoundError at the start rather then during the process.
        """
        # If the files have headers
        if self._has_header is True:
            # Get the reference header row
            self.header = self.get_header(self._files[0])

            if self._check_header is True:
                # Make sure all the files can be opened and have the same
                # header
                for i in self._files[1:]:
                    comp_header = self.get_header(i)
                    if self.header != self.get_header(i):
                        raise ValueError(
                            f"headers are different: {self._files[0]} vs. {i}"
                            f"[{self.header}] vs. [{comp_header}]"
                        )
            else:
                # If we are not checking the header just make sure we can open
                for i in self._files:
                    open(i).close()
        else:
            # We have no header so just make sure we can open
            for i in self._files:
                open(i).close()

        # Indicate that open has been called
        self._opened = True

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close all the files, and ensure any files that need deleting are
        removed.
        """
        self.close_file_batch(self._final_merge)

        for i in self._tempfiles:
            try:
                os.unlink(i)
            except Exception:
                pass

        if self._delete is True:
            for i in self._remaining_files:
                try:
                    os.unlink(i)
                except Exception:
                    pass

        self._opened = False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_header(self, filename):
        """Grab a header row from a input file. This will open read the first
        row and close the file.

        Parameters
        ----------
        filename : `str`
            The path to the file.

        Returns
        -------
        header : `str`
            The first row from the file.
        """
        with self.open_file(filename) as infile:
            return next(infile)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_file(self, filename):
        """Open a file for reading. This will also skip any fixed numbers of
        leading lines and comment characters.

        Parameters
        ----------
        filename : `str`
            The path to the file.

        Returns
        -------
        file_obj : `File`
            A file object for reading.
        """
        file_obj = self._read_method(
            filename, self._read_mode, **self._read_kwargs
        )

        # If we have leading comment characters or fixed numbers of lines to
        # skip then do so before returning the file object.
        if self._comment is not None:
            pyhead.move_to_header(
                file_obj, comment=self._comment
            )
        return file_obj

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_file_batch(self, files):
        """Open a batch of files for reading and merging.

        Parameters
        ----------
        files : `list` of `str`
            A list of files to open.

        Returns
        -------
        file_objs : `list` or `File`
            A list of open file objects.
        """
        # The final merge
        file_objs = [
            self.open_file(i) for i in files
        ]
        return file_objs

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close_file_batch(self, files):
        """Open a batch of files for reading and merging.

        Parameters
        ----------
        files : `list` of `File`
            A list of files objects to close.
        """
        for i in files:
            i.close()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def read_file_batch_header(self, files):
        """Read a single row from a batch of files.

        Parameters
        ----------
        files : `list` of `File`
            A list of files objects to read from.
        """
        for i in files:
            next(i)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_merge_file(self):
        """Open a merge file file for writing.

        Returns
        -------
        merge_file : `File`
            A file object to write an intermediate merge file to.
        """
        # This will hold the sorted output for the current loop
        tmp_merge = utils.get_tmp_file(
            dir=self._tmpdir, prefix=self._temp_prefix
        )
        self._tempfiles.append(tmp_merge)

        return self._write_method(
            tmp_merge, self._write_mode, **self._write_kwargs
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_header(self, file):
        """Write a header row to a file.

        Returns
        -------
        file : `File`
            A file object to write the header to.
        """
        file.write(self.header)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _merge(self):
        """Start the merge process until we have a small enough remaining
        files to merge in a single pass.
        """
        # This will hold all the files that need merging over the various
        # iterations
        self._remaining_files = list(self._files)

        # Holds the file names on intermediate (so files that are created
        # because we can't merge all the files in one go). These will be
        # deleted irrespective of the overall delete parameter
        self._tempfiles = []

        iterations = 0

        # We merge until we get below are our file limit and then hopefully
        # we will have a one final merge
        while len(self._remaining_files) > self._max_files:
            start_at = time.time()
            to_merge = []
            for i in range(self._max_files):
                try:
                    # Fill up from the start of the file list
                    to_merge.append(self._remaining_files.pop(0))
                except IndexError:
                    # No more files remaining
                    break

            merge_file = self.open_merge_file()

            try:
                merge_fobjs = self.open_file_batch(to_merge)

                # Now we merge into the temp file
                if self._has_header is True:
                    self.read_file_batch_header(merge_fobjs)
                    self.write_header(merge_file)

                # Now merge the rest
                for row in heapq.merge(*merge_fobjs, key=self._key,
                                       reverse=self._reverse):
                    merge_file.write(row)
            except Exception:
                merge_file.close()
                self.close_file_batch(merge_fobjs)
                raise

            merge_file_name = merge_file.name
            merge_file.close()
            self.close_file_batch(merge_fobjs)

            # Also the tmp_file we just created is now remaining
            self._remaining_files.append(merge_file_name)

            # Delete the batch
            if self._delete is True:
                for i in to_merge:
                    os.unlink(i)

            iterations += 1
            end_at = time.time()
            self._iterations.append(
                (iterations, len(merge_fobjs), end_at - start_at,
                 [i.name for i in merge_fobjs])
            )

        # The final merge
        self._final_merge = self.open_file_batch(self._remaining_files)

        # Now we merge into the temp file
        if self._has_header is True:
            self.read_file_batch_header(self._final_merge)

        self._heapq = heapq.merge(
            *self._final_merge,
            key=self._key,
            reverse=self._reverse
        )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvIterativeHeapqMerge(IterativeHeapqMerge):
    """Performs an interative merge sort over many files.

    Parameters
    ----------
    files : `list` of `str`
        One or more file names. If header is ``True`` then the first row of
        each file must be the header.
    key : `function`, optional, default: `NoneType`
        The function to provide the sort keys to ``heapq.merge``.
    max_files : `int`, optional, default: `16`
        The maximum number of files that can be open (and merged) at any
        one time. The default (16) is the same as unix sort.
    tmpdir : `str`, optional, default: `NoneType`
        The directory for temp merge files to be placed if the merge can't be
        performed in a single pass. The default is the system temp location.
    read_method : `function`, optional, default: `gzip.open`
        The method to use to open the input files for merging.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write any output files if the merge can't be
        completed in a single pass.
    read_mode : `str`, optional, default: `rt`
        The mode to use to open the input files for merging.
    write_mode : `str`, optional, default: `wt`
        The mode to use to write and intermediate files if the merge cannot be
        completed in a single pass.
    header : `bool`, optional, default: `True`
        Do the input files have a header, if so, the header of the first file
        is stored in the ``header`` attribute.
    check_header : `bool`, optional, default: `True`
        If ``header`` is ``True``, do you want to make sure that all the
        headers are the same in each of the input files?
    delete : `bool`, optional, default: `False`
        Delete the source files as they are merged?
    reverse : `bool`, optional, default: `False`
        Perform the sort in reverse.
    comment : `str`, optional, default: `NoneType`
        A comment character that might appear before the header/data row, these
        are skipped. These are not checked beyond the first header/data row.
    read_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``read_method``.
    write_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``write_method``.
    csv_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the
        ``csv.reader``/``csv.writer``.
    """

    TEMP_PREFIX = "csv_heapq_merge_"
    """A prefix that will be added to any intermediate merge files should the
    merge not be completed in a single pass (`str`).
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, csv_kwargs=None, **kwargs):
        super().__init__(*args, **kwargs)
        csv_kwargs = csv_kwargs or dict()
        self._csv_kwargs = csv_kwargs

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_file(self, filename):
        """Open a file for reading.

        Parameters
        ----------
        filename : `str`
            The path to the file.

        Returns
        -------
        csv_wrapper : `merge_sort.common.CsvWrapper`
            A file object for reading.
        """
        file_obj = super().open_file(filename)
        csv_reader = common.CsvWrapper(file_obj, mode='r', **self._csv_kwargs)
        return csv_reader.open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_merge_file(self):
        """Open a merge file file for writing.

        Returns
        -------
        csv_wrapper : `merge_sort.common.CsvWrapper`
            A file object for writing.
        """
        file_obj = super().open_merge_file()
        csv_reader = common.CsvWrapper(file_obj, mode='w', **self._csv_kwargs)
        return csv_reader.open()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvDictIterativeHeapqMerge(IterativeHeapqMerge):
    """Performs an interative merge sort over many files.

    Parameters
    ----------
    files : `list` of `str`
        One or more file names. If header is ``True`` then the first row of
        each file must be the header.
    key : `function`, optional, default: `NoneType`
        The function to provide the sort keys to ``heapq.merge``.
    max_files : `int`, optional, default: `16`
        The maximum number of files that can be open (and merged) at any
        one time. The default (16) is the same as unix sort.
    tmpdir : `str`, optional, default: `NoneType`
        The directory for temp merge files to be placed if the merge can't be
        performed in a single pass. The default is the system temp location.
    read_method : `function`, optional, default: `gzip.open`
        The method to use to open the input files for merging.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write any output files if the merge can't be
        completed in a single pass.
    read_mode : `str`, optional, default: `rt`
        The mode to use to open the input files for merging.
    write_mode : `str`, optional, default: `wt`
        The mode to use to write and intermediate files if the merge cannot be
        completed in a single pass.
    header : `bool`, optional, default: `True`
        Do the input files have a header, if so, the header of the first file
        is stored in the ``header`` attribute.
    check_header : `bool`, optional, default: `True`
        If ``header`` is ``True``, do you want to make sure that all the
        headers are the same in each of the input files?
    delete : `bool`, optional, default: `False`
        Delete the source files as they are merged?
    reverse : `bool`, optional, default: `False`
        Perform the sort in reverse.
    comment : `str`, optional, default: `NoneType`
        A comment character that might appear before the header/data row, these
        are skipped. These are not checked beyond the first header/data row.
    read_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``read_method``.
    write_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the ``write_method``.
    csv_kwargs : `dict`, optional, default: `NoneType`
        Any keyword arguments that will be passed to the
        ``csv.DictReader``/``csv.DictWriter``.
    """

    TEMP_PREFIX = "csv_dict_heapq_merge_"
    """A prefix that will be added to any intermediate merge files should the
    merge not be completed in a single pass (`str`).
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, csv_kwargs=None, **kwargs):
        super().__init__(*args, **kwargs)
        csv_kwargs = csv_kwargs or dict()
        self._csv_read_kwargs = csv_kwargs
        self._csv_write_kwargs = dict(csv_kwargs)

        # Remove any DictWriter only keyword arguments
        for i in ['restval', 'extrasaction']:
            try:
                self._csv_read_kwargs.pop(i)
            except KeyError:
                pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_header(self, filename):
        """Grab a header row from a input file. This will open read the first
        row and close the file.

        Parameters
        ----------
        filename : `str`
            The path to the file.

        Returns
        -------
        header : `str`
            The first row from the file.
        """
        try:
            infile = self.open_file(filename)
            header = infile.csv.fieldnames
        finally:
            infile.close()
        return header

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def read_file_batch_header(self, files):
        """Read a single row from a batch of files.

        Parameters
        ----------
        files : `list` of `File`
            A list of files objects to read from.
        """
        pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_file(self, filename):
        """Open a file for reading.

        Parameters
        ----------
        filename : `str`
            The path to the file.

        Returns
        -------
        csv_wrapper : `merge_sort.common.CsvDictWrapper`
            A file object for reading.
        """
        file_obj = super().open_file(filename)
        csv_reader = common.CsvDictWrapper(
            file_obj, mode='r', **self._csv_read_kwargs
        )
        return csv_reader.open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_merge_file(self):
        """Open a merge file file for writing.

        Returns
        -------
        csv_wrapper : `merge_sort.common.CsvDictWrapper`
            A file object for writing.
        """
        file_obj = super().open_merge_file()
        csv_reader = common.CsvDictWrapper(
            file_obj, mode='w', fieldnames=self.header,
            **self._csv_write_kwargs
        )
        return csv_reader.open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_header(self, file):
        """Writer a header row to the file.
        """
        file.csv.writeheader()
