"""
Interface to Mongo for using Mongo as a buffer for chunked data transfer.

We use Mongo as a buffer because we cannot guarantee that all of the requests
for individual chunks will hit the same worker. While we could use streamed responses for the download,
they're not available for upload, so for simplicity we use the same approach in both directions.

The intended usage pattern is that a single message is stored as a single "chunk set".
The chunk set is first "filled" (either by the client or the server), then "emptied" to obtain the data
(again, by either the client or the server).

Note that while this can be used for pagination, it is not in itself pagination.

We plan to gather most of this material into ew_common after the chunked transfer interface has been rolled out
to its three main users (ew_geostore, ew_specklia, ew_online_processing_service) and proven effective for each.
At that point, this entire module will move into ew_common. Note that the chunked transfer interface will always
require MongoDB or a similar provision to work correctly.

IMPORTANT: THE VERSION HERE IN THE SPECKLIA PACKAGE MUST NOT BE MADE DEPENDENT UPON EW_COMMON SINCE EW_COMMON
IS PRIVATE BUT THIS PACKAGE IS PUBLIC!
"""

from enum import Enum
from http import HTTPStatus
from io import BytesIO
from logging import Logger
import struct
from typing import List, Tuple, Union

from geopandas import GeoDataFrame, read_feather as read_geofeather
from pandas import DataFrame, read_feather
import requests

CHUNK_DB_NAME = "data_transfer_chunks"
CHUNK_METADATA_COLLECTION_NAME = "chunk_metadata"
MAX_CHUNK_AGE_SECONDS = 3600
MAX_CHUNK_SIZE_BYTES = 5 * 1024 ** 2  # must be small enough to fit into an HTTP GET Request


class ChunkSetStatus(Enum):
    """
    Chunk set status.

    Prevents the accidental access of chunk sets that have not yet received all of their data.
    """

    FILLING = 0
    EMPTYING = 1


def upload_chunks(api_address: str, chunks: List[Tuple[int, bytes]], logger: Logger) -> str:
    """
    Upload data chunks.

    Upload a series of data chunks through the chunked transfer mechanism.
    This method is for use on the client, not the server.

    Parameters
    ----------
    api_address : str
        The full URL of the API, including port but not including endpoint, e.g. "http://127.0.0.1:9999"
    chunks : List[Tuple[int, bytes]]
        A list of tuples containing the ordinal number of the chunk and each chunk
    logger : Logger
        A logger with which to log the upload.

    Returns
    -------
    str
        The chunk set uuid of the uploaded chunks
    """
    # post the first chunk to start the upload
    response = requests.post(
        api_address + f"/chunk/upload/{chunks[0][0]}-of-{len(chunks)}",
        data=chunks[0][1])
    logger.info("response from very first /chunk/upload was '%s'", response.json())
    assert response.status_code == HTTPStatus.OK, response.text
    chunk_set_uuid = response.json()['chunk_set_uuid']

    # post the rest of the chunks in a random order
    for i, chunk in chunks[1:]:
        response = requests.post(
            api_address + f"/chunk/upload/{chunk_set_uuid}/{i}-of-{len(chunks)}", data=chunk)
        logger.info("response from subsequent /chunk/upload/uuid call was '%s'", response.text)
        assert response.status_code == HTTPStatus.OK, response.text

    return chunk_set_uuid


def download_chunks(api_address: str, chunk_set_uuid: str) -> List[Tuple[int, bytes]]:
    """
    Download data chunks.

    Download a series of data chunks through the chunked transfer mechanism.
    This method is for use on the client, not the server.

    Parameters
    ----------
    api_address : str
        The full URL of the API, including port but not including endpoint, e.g. "http://127.0.0.1:9999"
    chunk_set_uuid : str
        The uuid of the chunk set to download.

    Returns
    -------
    chunks : List[Tuple[int, bytes]]
        A list of tuples containing the ordinal number of the chunk and each chunk
    """
    # fetch the data
    data_chunks = []
    finished = False
    while not finished:
        this_chunk_response = requests.get(api_address + f"/chunk/download/{chunk_set_uuid}")
        if this_chunk_response.status_code == HTTPStatus.NO_CONTENT:
            finished = True
        else:
            data_chunks.append((
                struct.unpack('i', this_chunk_response.content[:4])[0],
                this_chunk_response.content[4:]))

    return data_chunks


def split_into_chunks(data: bytes, chunk_size: int = MAX_CHUNK_SIZE_BYTES) -> List[Tuple[int, bytes]]:
    """
    Split data into compressed chunks for transport.

    Parameters
    ----------
    data : bytes
        The data to be split into chunks.
    chunk_size: int
        The maximum number of bytes allowed in each chunk.

    Returns
    -------
    List[Tuple[int, bytes]]
        A list of tuples containing the ordinal number of the chunk and each chunk
    """
    return list(
        enumerate((data[i:i + chunk_size] for i in range(0, len(data), chunk_size)), start=1))


def merge_from_chunks(chunks: List[Tuple[int, bytes]]) -> bytes:
    """
    Merge data that has been split into compressed chunks back into a single message.

    Parameters
    ----------
    chunks : List[Tuple[int, bytes]]
        A list of tuples containing the ordinal number of the chunk and each chunk

    Returns
    -------
    bytes
        The merged data
    """
    return b''.join([dc[1] for dc in sorted(chunks, key=lambda x: x[0])])


def deserialise_dataframe(data: bytes) -> Union[DataFrame, GeoDataFrame]:
    """
    Convert a binary serialised feather table to pandas dataframe.

    Parameters
    ----------
    data : bytes
        Binary serialised feather table.

    Returns
    -------
    Union[DataFrame, GeoDataFrame]
        Input table converted to a pandas dataframe.

    Raises
    ------
    ValueError
        When bytes can't be interpreted as meaningful dataframe.
    """
    try:
        buffer = BytesIO(data)
        df = read_geofeather(buffer)
    except ValueError as e:
        # First attempt to deserialise as a geodataframe. If geo meta is missing, we expect a clear ValueError
        # and we then load as a plain dataframe instead.
        if "Missing geo meta" in e.args[0] or "'geo' metadata" in e.args[0]:
            try:
                df = read_feather(BytesIO(data))
            except ValueError as e:
                raise ValueError("Couldn't deserialise table format") from e
        else:
            raise ValueError("Couldn't deserialise table format") from e
    return df


def serialise_dataframe(df: Union[DataFrame, GeoDataFrame]) -> bytes:
    """
    Serialise a dataframe using the feather table format.

    Parameters
    ----------
    df : DataFrame
        Input dataframe

    Returns
    -------
    bytes
        Serialised feather table.
    """
    feather_buffer = BytesIO()
    # Browser implementations of feather do not support compressed feather formats.
    df.to_feather(feather_buffer, compression='uncompressed')
    feather_buffer.seek(0)
    return feather_buffer.getvalue()
