# This file is part of the History Store (histore).
#
# Copyright (C) 2018-2021 New York University.
#
# The History Store (histore) is released under the Revised BSD License. See
# file LICENSE for full license details.

"""Class for sorting input documents in main-memory or using secondary storage.
The functions in this module are inspired by csvsort
(https://github.com/richardpenman/csvsort).

For HISTORE we made a few modifications that allow us to handle arbitraty input
documents. First, if the document rows can be kept in main-memory they will be
sorted directly instead of writing them to disk after sorting. Second, we use
JSON as the serialization format to retain raw data types of row values.
"""

from tempfile import NamedTemporaryFile as TempFile
from typing import Callable, List, Optional, Tuple

import heapq
import json
import os
import sys

from histore.document.base import Document, DocumentIterator
from histore.document.json.reader import JsonIterator
from histore.document.json.writer import JsonWriter
from histore.document.mem import InMemoryDocument
from histore.document.schema import DocumentSchema

import histore.config as config
import histore.util as util


class SortEngine:
    """External merge sort for input documents. Seralizes a given document into
    JSON format. The advantage of using JSON as the seriaization format is the
    ablity to maintain the data type of the document rows.
    """
    def __init__(
        self, buffersize: Optional[float] = None,
        encoder: Optional[json.JSONEncoder] = None,
        decoder: Optional[Callable] = None
    ):
        """Initialize the JSON encoder and decoder.

        Uses the default encoder and decoder if the respective argument values
        are None.

        Parameters
        ----------
        buffersize: float, default=None
            Maximum size (in bytes) of file blocks that are kept in main-memory.
            If no buffer size is given the value from the environment variable
            HISTORE_SORTBUFFER is used.
        encoder: json.JSONEncoder, default=None
            Encoder used when writing rows as JSON objects to file during
            external merge sort.
        decoder: func, default=None
            Custom decoder function when reading archive rows from file. If not
            given, the default decoder will be used.
        """
        self.buffersize = buffersize if buffersize else config.SORTBUFFER()
        self.encoder = encoder
        self.decoder = decoder

    def _mergesort(
        self, buffer: List, filenames: List[str], columns: DocumentSchema,
        keys: List[int]
    ) -> str:
        """2-way merge sort for sorted blocks of an input document.

        Adopted from
        https://github.com/richardpenman/csvsort/blob/master/__init__.py

        Returns the name of the file containg the sorted output.

        Parameters
        ----------
        buffer: list
            Buffer of rows from the input file that were not written to file.
        filename: list
            List of names for sorted temporary files that were generated by the
            split method.
        columns: list of string
            List of column names.
        keys: list
            List of index positions for sort columns.

        Returns
        -------
        str
        """
        # Merge the buffer if it contains rows.
        if buffer:
            mergefile = filenames[0]
            with TempFile(delete=False, mode='w', newline='') as f_out:
                with JsonWriter(filename=f_out, encoder=self.encoder) as writer:
                    files = [
                        decorated_buffer(buffer, keys),
                        decorated_file(mergefile, keys, self.decoder)
                    ]
                    writer.write(columns)
                    for _, row in heapq.merge(*files):
                        writer.write(row)
                filenames = filenames[1:] + [f_out.name]
            os.remove(mergefile)
        while len(filenames) > 1:
            mergefiles, filenames = filenames[:2], filenames[2:]
            with TempFile(delete=False, mode='w', newline='') as f_out:
                with JsonWriter(filename=f_out, encoder=self.encoder) as writer:
                    files = [decorated_file(f, keys, self.decoder) for f in mergefiles]
                    writer.write(columns)
                    for _, row in heapq.merge(*files):
                        writer.write(row)
                filenames.append(f_out.name)
            for filename in mergefiles:
                os.remove(filename)
        return filenames[0]

    def sorted(self, doc: Document, keys: List[int]) -> Document:
        """Sort a input document based on the given columns.

        Returns a document that wraps a temporary file with the sorted
        document rows.

        Parameters
        ----------
        doc: histore.document.base.Document
            Input document.
        keys: list
            List of index positions for sort columns.

        Returns
        -------
        histore.document.base.Document
        """
        columns = doc.columns
        try:
            reader = doc.open()
            buffer, filenames = self._split(
                reader=reader,
                columns=columns,
                keys=keys
            )
        finally:
            doc.close()
        if not filenames:
            buffer.sort(key=lambda row: util.keyvalue(row[2], keys))
            return InMemoryDocument(columns=columns, rows=buffer)
        else:
            filename = self._mergesort(
                buffer=buffer,
                filenames=filenames,
                columns=columns,
                keys=keys
            )
            # Import JsonDocument here to avoid cyclic imports.
            from histore.document.json.base import JsonDocument
            return JsonDocument(
                filename=filename,
                encoder=self.encoder,
                decoder=self.decoder,
                delete_on_close=True
            )

    def _split(
        self, reader: DocumentIterator, columns: DocumentSchema,
        keys: List[int]
    ) -> Tuple[List, List[str]]:
        """Split a input document into blocks of maximum size.

        Individual blocks are written to temporary files on disk in JSON format.
        Only the final buffer is maintained in memory. Returns the memory buffer
        and the list of names for temporary files that are created.

        Adopted from
        https://github.com/richardpenman/csvsort/blob/master/__init__.py

        Parameters
        ----------
        reader: histore.document.base.DocumentIterator
            Iterator for rowid, cell values in an input document.
        columns: list of string
            List of column names.
        keys: list
            List of index positions for sort columns.

        Returns
        -------
        list(rows), list(filenames)
        """
        # Split document rows into blocks
        buffer = list()
        current_size = 0
        tmp_filenames = list()
        for row in reader:
            buffer.append(row)
            current_size += sys.getsizeof(row)
            if current_size > self.buffersize:
                # Sort buffer.
                buffer.sort(key=lambda row: util.keyvalue(row[2], keys))
                # Write buffer to disk.
                with TempFile(delete=False, mode='w', newline='') as f_out:
                    with JsonWriter(filename=f_out, encoder=self.encoder) as writer:
                        writer.write(columns)
                        for row in buffer:
                            writer.write(row)
                    # Append file name to result file name list.
                    tmp_filenames.append(f_out.name)
                # Clear buffer.
                buffer = list()
                current_size = 0
        return buffer, tmp_filenames


# -- Helper Functions ---------------------------------------------------------


def decorated_buffer(buffer: List, columns: List[int]):
    """Iterator for in-memory buffer of data rows.

    Parameters
    ----------
    buffer: list
        Buffer of rows from the input file that were not written to file.
    columns: list
        List of index positions for sort columns.
    """
    buffer.sort(key=lambda row: util.keyvalue(row[2], columns))
    for row in buffer:
        yield util.keyvalue(row[2], columns), row


def decorated_file(filename: str, columns: List[int], decoder: Callable):
    """Iterator for sorted file block. From
    https://github.com/richardpenman/csvsort/blob/master/__init__.py

    Parameters
    ----------
    filename: string
        Temporary file (block) on disk.
    columns: list
        List of index positions for sort columns.
    decoder: func
        Custom decoder function when reading archive rows from file. If not
        given, the default decoder will be used.
    """
    with JsonIterator(filename, decoder=decoder) as reader:
        for row in reader:
            yield util.keyvalue(row[2], columns), row
