# 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
import warnings
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 deprecated.sphinx import deprecated

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

logger = logging.getLogger(__name__)


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

    def __init__(
        self,
        data_handle: Optional[str],
        layer: "Layer",
        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."""
        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 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.
        :raises ValueError: Collision occurred for data handle, or data handle not set
        """
        if not self.data_handle:
            raise ValueError(f"data_handle not set for partition {self.id}")
        for _ in range(3):
            data_handle_exists = self.layer._data_blob_api.check_handle_exists(
                self.layer.id, data_handle=self.data_handle
            )
            if not data_handle_exists:
                break
            self.data_handle = self.generate_data_handle()
        else:
            raise ValueError(f"Collision occurred for data handle {self.data_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: Collision occurred for data handle, or data handle not set
        """
        if not self.data_handle:
            raise ValueError(f"data_handle not set for partition {self.id}")
        for _ in range(3):
            data_handle_exists = self.layer._data_blob_api.check_handle_exists(
                self.layer.id, data_handle=self.data_handle
            )
            if not data_handle_exists:
                break
            self.data_handle = self.generate_data_handle()
        else:
            raise ValueError(f"Collision occurred for data handle {self.data_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
    """

    def __init__(
        self,
        id: Optional[str],
        layer: "IndexLayer",
        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
        super().__init__(
            data_handle=id,
            layer=layer,
            id=id,
            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
    - ``offset``: offset is deprecated. kafka_offset is used instead of offset
    - ``offset_partition``: offset_partition. kafka_partition is used instead of offset_partition
    """

    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,
        offset: Optional[int] = None,
        offset_partition: 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
        if offset:
            warnings.warn(
                "The use of offset is deprecated. Please use kafka_offset.", DeprecationWarning
            )
            self.kafka_offset = offset
        if offset_partition:
            warnings.warn(
                "The use of offset_partition is deprecated. Please use kafka_partition.",
                DeprecationWarning,
            )
            self.kafka_partition = offset_partition

    @property  # type: ignore
    @deprecated(version="2.4", reason="Please use kafka_offset")
    def offset(self) -> Optional[int]:
        """:return: kafka_offset, new property name that replaces offset."""
        return self.kafka_offset

    @property  # type: ignore
    @deprecated(version="2.4", reason="Please use kafka_partition")
    def offset_partition(self) -> Optional[int]:
        """:return: kafka_partition, new property name that replaces offset_partition."""
        return self.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.")
