# Copyright (C) 2020-2022 HERE Global B.V. and its affiliate(s).
# All rights reserved.
#
# This software and other materials contain proprietary information
# controlled by HERE and are protected by applicable copyright legislation.
# Any use and utilization of this software and other materials and
# disclosure to any third parties is conditional upon having a separate
# agreement with HERE for the access, use, utilization or disclosure of this
# software. In the absence of such agreement, the use of the software is not
# allowed.

"""HERE platform catalog layer partition abstraction."""

import asyncio
import logging
import time
import uuid
from contextlib import contextmanager
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, BinaryIO, Dict, Optional, Union

from aiohttp import ClientSession
from here.platform.utils.deprecation import _deprecate_remove_parameters

if TYPE_CHECKING:
    from here.platform.layer import IndexLayer, StreamLayer, VersionedLayer, VolatileLayer

logger = logging.getLogger(__name__)


class Partition:
    """HERE platform partition abstraction."""

    def __init__(
        self,
        data_handle: Optional[str] = None,
        layer=None,  # type: ignore
        id: Optional[str] = None,
        checksum: Optional[str] = None,
        data_size: Optional[int] = None,
        crc: Optional[str] = None,
    ):
        """Instantiate Partition given a data handle, layer and other parameters."""
        assert layer
        self.data_handle = data_handle
        self.layer = layer
        self.id = id
        self.checksum = checksum
        self.data_size = data_size
        self.crc = crc
        self.billing_tag = self.layer.billing_tag  # type: ignore

    def get_blob(self) -> bytes:
        """
        Get blob (raw bytes) inside the partition for this layer with given
        data_handle.

        :return: Content of the blob referenced by the data handle.
        :raises ValueError: Unsupported content encoding or invalid data handle
        """
        if not self.data_handle:
            raise ValueError(f"data_handle not set for partition {self.id}")
        volume_type = self.layer.configuration.json["volume"]["volumeType"]
        if volume_type == "volatile":
            blob = self.layer._data_volatile_blob_api.get_volatile_blob(
                layer_id=self.layer.id,
                data_handle=self.data_handle,
                billing_tag=self.billing_tag,
            )
        else:
            blob = self.layer._data_blob_api.get_blob_by_handle(
                layer_id=self.layer.id,
                data_handle=self.data_handle,
                billing_tag=self.billing_tag,
            )
        if self.layer.configuration.content_encoding not in ("gzip", "identity", None):
            raise ValueError(
                f"Unsupported content encoding : {self.layer.configuration.content_encoding}"
            )
        blob_bytes: bytes = blob
        return blob_bytes

    @staticmethod
    def generate_data_handle() -> str:
        """
        Generate a unique data handle.

        :return: A string representing a unique blob identifier
        """
        return str(uuid.uuid4())

    def __current_time(self):
        """This is a time.time() wrapper to make it testable in a deterministic way"""
        return time.time()

    def __generate_new_handle(self):
        start_time = self.__current_time()
        while self.__current_time() - start_time <= self.layer.catalog.platform._retry_max_time:
            data_handle = self.generate_data_handle()
            if self.layer._data_blob_api.check_handle_exists(
                self.layer.id, data_handle=data_handle
            ):
                time.sleep(self.layer.catalog.platform._polling_wait)
            else:
                self.data_handle = data_handle
                break
        if not self.data_handle:
            raise ValueError("Data handle not generated, max retry timeout reached")

    def put_blob(self, path_or_data: Union[str, bytes, Path]):
        """
        Upload data in single part.

        :param path_or_data: Path of the file or data in bytes to be uploaded as a partition.
        """
        if not self.data_handle:
            self.__generate_new_handle()

        if isinstance(path_or_data, bytes):
            data = path_or_data
        else:
            with open(path_or_data, "rb") as file_data:
                data = file_data.read()
        data_length = len(data)
        content_encoding = self.layer.configuration.content_encoding
        content_type = self.layer.configuration.content_type
        self.layer._data_blob_api.publish_blob_by_handle(
            layer_id=self.layer.id,
            data_handle=self.data_handle,
            data=data,
            content_length=data_length,
            content_encoding=content_encoding,
            content_type=content_type,
            billing_tag=self.billing_tag,
        )

    def multipart_upload(
        self,
        path_or_data: Union[str, bytes, Path],
        part_size: int,
    ):
        """
        Multipart upload data to a blob store.

        :param path_or_data: Path of the file or data in bytes to be upload as a partition.
        :param part_size: An int representing size in MB, to upload in multiple parts
            minimum value is 5MB and and Maximum is 50MB.

        :raises ValueError: Data handle not generated, max retry timeout reached
        """
        if not self.data_handle:
            self.__generate_new_handle()

        content_type = self.layer.configuration.content_type
        assert content_type
        content_encoding = self.layer.configuration.content_encoding or "identity"
        multi_parts = self.layer._data_blob_api.start_multipart_upload_by_handle(
            layer_id=self.layer.id,
            data_handle=self.data_handle,
            content_type=content_type,
            billing_tag=self.billing_tag,
            content_encoding=content_encoding,
        )

        upload_url = multi_parts["links"]["uploadPart"]["href"]
        logger.debug(f"URL: {upload_url}")
        try:
            loop = asyncio.get_event_loop()  # event loop
        except RuntimeError as e:
            if e.args[0].startswith("There is no current event loop"):
                logger.info("Trying with new event loop")
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
        headers = self.layer._base_api.headers
        future = asyncio.ensure_future(
            self.read_and_upload_parts(path_or_data, upload_url, content_type, headers, part_size)
        )
        results = loop.run_until_complete(future)
        parts = []
        for result in results:
            for k, v in result.items():
                if v.get("Etag"):
                    parts.append(dict(etag=v["Etag"], number=k))
                else:
                    # Throw error instead of returning result.
                    raise ValueError("Expected header Etag in response: ", result)

        # complete multipart upload
        self.layer._data_blob_api.complete_multipart_upload_by_handle(
            complete_href_url=multi_parts["links"]["complete"]["href"],
            parts=parts,
            billing_tag=self.billing_tag,
        )

    async def read_and_upload_parts(
        self,
        path_or_data: Union[str, bytes, Path],
        upload_url: str,
        content_type: str,
        headers: Dict[str, str],
        part_size: int,
    ):
        """
        Read input file in chunks of 50Mb each and creates urls.

        :param path_or_data: Path of the file or data to be read and upload as a partition.
        :param upload_url: An URL to upload blob data.
        :param content_type: A standard MIME type describing the format of the blob data.
        :param headers: A dict containing http headers.
        :param part_size: An int representing size in MB, to upload in multiple parts
            minimum value is 5MB and and Maximum is 50MB.
        :return: A list of dict with key as part number and value as url response body.
        """
        tasks = []
        completed_tasks = []
        with self.get_data_handler(path_or_data) as data_handler:
            async with ClientSession() as session:
                data = data_handler.read(part_size * 1024 * 1024)  # reading in 50Mb chunks
                part_num = 1
                headers["Content-Type"] = content_type
                while data:
                    logger.debug(f"Starting part number: {part_num}")
                    upd_url = f"{upload_url}?partNumber={part_num}"
                    if self.billing_tag:
                        upd_url = "".join([upd_url, f"&billingTag={self.billing_tag}"])
                    task = asyncio.ensure_future(
                        self.upload_part(session, upd_url, part_num, headers, data)
                    )
                    part_num += 1
                    tasks.append(task)  # create list of tasks
                    if len(tasks) % 5 == 0:
                        completed_tasks.extend(await asyncio.gather(*tasks))
                        logger.info(f"Completed parts: {len(completed_tasks)}")
                        tasks = []
                    data = data_handler.read(part_size * 1024 * 1024)
                completed_tasks.extend(await asyncio.gather(*tasks))
                return completed_tasks  # gather task responses

    @staticmethod
    @contextmanager
    def get_data_handler(path_or_data: Union[str, bytes, Path]):
        """Get the data handler to read from file path or data.

        :param path_or_data: Path of the file or data to be read and uploaded as a partition.
        :yield Union[BytesIO, BinaryIO]: Data Handle
        """
        try:
            if isinstance(path_or_data, bytes):
                handle = BytesIO(path_or_data)  # type: Union[BytesIO, BinaryIO]
                yield handle
            else:
                handle = open(path_or_data, "rb")
                yield handle
        finally:
            handle.close()

    async def upload_part(
        self,
        session: ClientSession,
        upload_url: str,
        part_num: int,
        headers: Dict[str, str],
        data: bytes,
        max_retry_count: int = 3,
    ):
        """
        Upload a part data to a url, using specified ClientSession.

        :param session: a ClientSession object for making requests.
        :param upload_url: an URL to upload blob data.
        :param part_num: an unique number for the part to be uploaded.
        :param headers: HTTP headers to be send with the request.
        :param data: blob data to be send with the request.
        :param max_retry_count: Max retry count in case of failure with exponential backoff -
            60 seconds.

        :return: Returns a dict with key as part number and value as response body.
        """
        upload_part_resp: Dict = {}
        logger.debug(f"Upload URL: {upload_url}")
        if self.layer._has_gz_encoding():
            headers["Content-Encoding"] = "gzip"
        data_length = len(data)
        logger.debug(f"Data Length: {data_length}")
        headers["Content-Length"] = str(data_length)
        for i in range(max_retry_count):
            async with session.post(upload_url, headers=headers, data=data, timeout=0) as response:
                await response.read()
                # To be implemented
                if response.status == 200:
                    logger.debug(f"Response Status: {response.status}")
                    logger.debug(f"Headers: {response.headers}")
                    upload_part_resp[part_num] = response.headers
                    break
                elif 500 <= response.status <= 599 or response.status == 429:
                    upload_part_resp[part_num] = {"error_response": response}
                    time.sleep(i * 60)
                else:
                    upload_part_resp[part_num] = {"error_response": response}
        logger.debug(f"Response: {upload_part_resp}")
        return upload_part_resp

    def upload_volatile_blob(self, path_or_data: Union[str, Path, bytes]):
        """
        Upload data to volatile blob.

        :param path_or_data: Path of the file or data to be read and upload as a partition
        :raises ValueError: in case data handle is not set
        """
        if not self.data_handle:
            raise ValueError(f"data_handle not set for partition {self.id}")
        with Partition.get_data_handler(path_or_data) as data_handler:
            data = data_handler.read()
        self.layer._data_volatile_blob_api.put_volatile_blob(
            layer_id=self.layer.id,
            data_handle=self.data_handle,
            body=data,
            billing_tag=self.billing_tag,
        )


class VersionedPartition(Partition):
    """
    Partition subclass used for ``VersionedLayer``s.

    In addition to the fields present in the base class, this class adds:
    - ``version``:  version of the partition
    """

    def __init__(
        self,
        data_handle: Optional[str],
        layer: "VersionedLayer",
        id: str,
        checksum: Optional[str] = None,
        data_size: Optional[int] = None,
        compressed_data_size: Optional[int] = None,
        crc: Optional[str] = None,
        version: Optional[int] = None,
    ):
        """Instantiate VersionedPartition given a data handle, layer and other parameters."""
        super().__init__(
            data_handle=data_handle,
            layer=layer,
            id=id,
            checksum=checksum,
            data_size=data_size,
            crc=crc,
        )
        self.version = version
        self.compressed_data_size = compressed_data_size


class VolatilePartition(Partition):
    """
    Partition subclass used for ``VolatileLayer``s.

    """

    def __init__(
        self,
        data_handle: Optional[str],
        layer: "VolatileLayer",
        id: Optional[str] = None,
        checksum: Optional[str] = None,
        data_size: Optional[int] = None,
        compressed_data_size: Optional[int] = None,
        crc: Optional[str] = None,
    ):
        """Instantiate VolatilePartition given a data handle, layer and other parameters."""
        super().__init__(
            data_handle=data_handle,
            layer=layer,
            id=id,
            checksum=checksum,
            data_size=data_size,
            crc=crc,
        )
        self.compressed_data_size = compressed_data_size


class IndexPartition(Partition):
    """
    Partition subclass used for ``IndexLayer``s.

    In addition to the fields present in the base class, this class adds:
    - ``timestamp``: timestamp of the partition
    - ``fields``: fields according to the index layer configuration
    - ``additional_metadata``: free-form additional metadata
    """

    @_deprecate_remove_parameters(
        remove_params=["id"],
        since_version="2.12",
        use_new_params=["data_handle"],
    )
    def __init__(
        self,
        layer: "IndexLayer",
        id: Optional[str] = None,
        data_handle: Optional[str] = None,
        checksum: Optional[str] = None,
        data_size: Optional[int] = None,
        crc: Optional[str] = None,
        timestamp: Optional[int] = None,
        fields: Dict[str, Union[str, int, bool]] = {},
        additional_metadata: Dict[str, str] = {},
    ):
        """Instantiate IndexPartition given a data handle, layer and other parameters."""
        # Index layers seems not to have partition ids,
        # the convention is to replicate data handles as ids
        # data handle is called id in the REST API
        data_handle = data_handle or id
        super().__init__(
            data_handle=data_handle,
            layer=layer,
            id=data_handle,
            checksum=checksum,
            data_size=data_size,
            crc=crc,
        )
        self.timestamp = timestamp
        self.fields = fields
        self.additional_metadata = additional_metadata


class StreamPartition(Partition):
    """
    Partition subclass used for ``StreamLayer``s.

    In addition to the fields present in the base class, this class adds:
    - ``data``: inline data, alternative to ``data_handle``
    - ``timestamp``: timestamp in milliseconds since the Unix epoch (1970-01-01T00:00:00 UTC)
    - ``kafka_offset``: the offset of the message in the Kafka stream partition
    - ``kafka_partition``: the Kafka stream partition number the offset is related to
    """

    def __init__(
        self,
        layer: "StreamLayer",
        data_handle: Optional[str] = None,
        id: Optional[str] = None,
        checksum: Optional[str] = None,
        data_size: Optional[int] = None,
        crc: Optional[str] = None,
        data: Optional[bytes] = None,
        timestamp: Optional[int] = None,
        kafka_partition: Optional[int] = None,
        kafka_offset: Optional[int] = None,
    ):
        """Instantiate StreamPartition given a data handle, layer and other parameters."""
        super().__init__(
            data_handle=data_handle,
            layer=layer,
            id=id,
            checksum=checksum,
            data_size=data_size,
            crc=crc,
        )
        self.data = data
        self.timestamp = timestamp
        self.kafka_offset = kafka_offset
        self.kafka_partition = kafka_partition

    def get_data(self) -> bytes:
        """
        Get the data associated with the ``StreamPartition``.

        Data can be present directly in the ``data`` field
        or retrieved via ``data_handle``. This function returns
        the data, regardless of where it is stored.

        :return: Data associated with the stream partition
        :raises ValueError: in case neither in ``data`` nor ``data_handle`` is available
        """
        if self.data:
            if self.layer.configuration.content_encoding not in ("gzip", "identity", None):
                raise ValueError(
                    f"Unsupported content encoding : {self.layer.configuration.content_encoding}"
                )
            return self.data
        elif self.data_handle:
            return self.get_blob()
        else:
            raise ValueError("Unsupported stream format. Expected data or data_handle attribute.")
