# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
from __future__ import annotations
import random
import uuid
import logging
import time
import threading
from typing import (
    Dict,
    Callable,
    List,
    Any,
    Mapping,
    Union,
    TYPE_CHECKING,
    Optional,
    Iterable,
    cast,
)
from functools import partial

from .._tracing import (
    process_context_manager,
    receive_context_manager,
    get_span_links_from_received_events,
    is_tracing_enabled,
)
from .partition_context import PartitionContext
from .in_memory_checkpoint_store import InMemoryCheckpointStore
from .ownership_manager import OwnershipManager
from .common import CloseReason, LoadBalancingStrategy
from ._eventprocessor_mixin import EventProcessorMixin

if TYPE_CHECKING:
    from datetime import datetime
    from .checkpoint_store import CheckpointStore
    from .._common import EventData
    from .._consumer import EventHubConsumer
    from .._consumer_client import EventHubConsumerClient


_LOGGER = logging.getLogger(__name__)


class EventProcessor(EventProcessorMixin):  # pylint:disable=too-many-instance-attributes
    """
    An EventProcessor constantly receives events from one or multiple partitions of the Event Hub
    in the context of a given consumer group.

    """

    def __init__(
        self,
        eventhub_client: EventHubConsumerClient,
        consumer_group: str,
        on_event: Union[
            Callable[["PartitionContext", Optional["EventData"]], None],
            Callable[["PartitionContext", List["EventData"]], None],
        ],
        *,
        batch: bool = False,
        max_batch_size: int = 300,
        max_wait_time: Optional[float] = None,
        partition_id: Optional[str] = None,
        on_error: Optional[Callable[[PartitionContext, Exception], None]] = None,
        on_partition_initialize: Optional[Callable[[PartitionContext], None]] = None,
        on_partition_close: Optional[Callable[[PartitionContext, CloseReason], None]] = None,
        checkpoint_store: Optional[CheckpointStore] = None,
        initial_event_position: Union[int, str, datetime, Mapping[str, Union[int, str, datetime]]] = "@latest",
        initial_event_position_inclusive: Union[bool, Mapping[str, bool]] = False,
        load_balancing_interval: float = 30.0,
        load_balancing_strategy: LoadBalancingStrategy = LoadBalancingStrategy.GREEDY,
        partition_ownership_expiration_interval: Optional[float] = None,
        owner_level: Optional[int] = None,
        prefetch: Optional[int] = None,
        track_last_enqueued_event_properties: bool = False,
    ) -> None:
        self._consumer_group = consumer_group
        self._eventhub_client = eventhub_client
        self._namespace = eventhub_client._address.hostname
        self._eventhub_name = eventhub_client.eventhub_name
        self._event_handler = on_event
        self._batch = batch
        self._max_batch_size = max_batch_size
        self._max_wait_time = max_wait_time
        self._partition_id = partition_id
        self._error_handler: Optional[Callable[[PartitionContext, Exception], None]] = on_error
        self._partition_initialize_handler: Optional[Callable[[PartitionContext], None]] = on_partition_initialize
        self._partition_close_handler: Optional[Callable[[PartitionContext, CloseReason], None]] = on_partition_close
        self._checkpoint_store = checkpoint_store or InMemoryCheckpointStore()
        self._initial_event_position: Union[int, str, datetime, Mapping[str, Union[int, str, datetime]]] = (
            initial_event_position
        )
        self._initial_event_position_inclusive: Union[bool, Mapping[str, bool]] = initial_event_position_inclusive

        self._load_balancing_interval: float = load_balancing_interval
        self._load_balancing_strategy = load_balancing_strategy
        self._ownership_timeout = (
            partition_ownership_expiration_interval
            if partition_ownership_expiration_interval is not None
            else self._load_balancing_interval * 6
        )

        self._partition_contexts: Dict[str, PartitionContext] = {}

        # Receive parameters
        self._owner_level: Optional[int] = owner_level
        if checkpoint_store and self._owner_level is None:
            self._owner_level = 0
        self._prefetch: Optional[int] = prefetch
        self._track_last_enqueued_event_properties = track_last_enqueued_event_properties
        self._id = str(uuid.uuid4())
        self._running = False
        self._lock = threading.RLock()

        self._consumers: Dict[str, EventHubConsumer] = {}
        self._ownership_manager = OwnershipManager(
            self._eventhub_client,
            self._consumer_group,
            self._id,
            self._checkpoint_store,
            self._ownership_timeout,
            self._load_balancing_strategy,
            self._partition_id,
        )

        self._last_received_time = time.time_ns()

    def __repr__(self) -> str:
        return "EventProcessor: id {}".format(self._id)

    def _process_error(self, partition_context, err):
        if self._error_handler:
            try:
                self._error_handler(partition_context, err)
            except Exception as err_again:  # pylint:disable=broad-except
                _LOGGER.warning(
                    "EventProcessor instance %r of eventhub %r partition %r consumer group %r. "
                    "An error occurred while running on_error. The exception is %r.",
                    self._id,
                    partition_context.eventhub_name,
                    partition_context.partition_id,
                    partition_context.consumer_group,
                    err_again,
                )

    def _cancel_tasks_for_partitions(self, to_cancel_partitions: Iterable[str]) -> None:
        with self._lock:
            _LOGGER.debug(
                "EventProcessor %r tries to cancel partitions %r",
                self._id,
                to_cancel_partitions,
            )
            for partition_id in to_cancel_partitions:
                if partition_id in self._consumers:
                    self._consumers[partition_id].stop = True
                    _LOGGER.info(
                        "EventProcessor %r has cancelled partition %r",
                        self._id,
                        partition_id,
                    )

    def _initialize_partition_consumer(self, partition_id):
        if self._partition_initialize_handler:
            try:
                self._partition_initialize_handler(self._partition_contexts[partition_id])
            except Exception as err:  # pylint:disable=broad-except
                _LOGGER.warning(
                    "EventProcessor instance %r of eventhub %r partition %r consumer group %r. "
                    "An error occurred while running on_partition_initialize. The exception is %r.",
                    self._id,
                    self._partition_contexts[partition_id].eventhub_name,
                    self._partition_contexts[partition_id].partition_id,
                    self._partition_contexts[partition_id].consumer_group,
                    err,
                )
                self._process_error(self._partition_contexts[partition_id], err)
        _LOGGER.info("EventProcessor %r has claimed partition %r", self._id, partition_id)

    def _create_tasks_for_claimed_ownership(
        self, claimed_partitions: Iterable[str], checkpoints: Optional[Dict[str, Dict[str, Any]]] = None
    ) -> None:
        with self._lock:
            _LOGGER.debug(
                "EventProcessor %r tries to claim partition %r",
                self._id,
                claimed_partitions,
            )
            for partition_id in claimed_partitions:
                if partition_id not in self._consumers:
                    if partition_id in self._partition_contexts:
                        partition_context = self._partition_contexts[partition_id]
                        partition_context._last_received_event = None  # pylint:disable=protected-access
                    else:
                        partition_context = PartitionContext(
                            self._namespace,
                            self._eventhub_name,
                            self._consumer_group,
                            partition_id,
                            self._checkpoint_store,
                        )
                        self._partition_contexts[partition_id] = partition_context

                    checkpoint = checkpoints.get(partition_id) if checkpoints else None
                    (
                        initial_event_position,
                        event_postition_inclusive,
                    ) = self.get_init_event_position(partition_id, checkpoint)
                    event_received_callback = partial(self._on_event_received, partition_context)
                    self._consumers[partition_id] = cast(
                        "EventHubConsumer",
                        self.create_consumer(
                            partition_id,
                            initial_event_position,
                            event_postition_inclusive,
                            event_received_callback,
                        ),
                    )
                    self._initialize_partition_consumer(partition_id)

    def _on_event_received(
        self, partition_context: PartitionContext, event: Union[Optional[EventData], List[EventData]]
    ) -> None:
        if event:
            try:
                partition_context._last_received_event = event[-1]  # type: ignore  #pylint:disable=protected-access
            except TypeError:
                partition_context._last_received_event = event  # type: ignore  #pylint:disable=protected-access

            links = []
            is_batch = False
            if is_tracing_enabled():
                links = get_span_links_from_received_events(event)
                if isinstance(event, list):
                    is_batch = True

            with receive_context_manager(
                self._eventhub_client, links=links, start_time=self._last_received_time
            ):
                self._last_received_time = time.time_ns()

            with process_context_manager(self._eventhub_client, links=links, is_batch=is_batch):
                self._event_handler(partition_context, event)  # type: ignore
        else:
            self._event_handler(partition_context, event)  # type: ignore

    def _load_balancing(self) -> None:
        """Start the EventProcessor.

        The EventProcessor will try to claim and balance partition ownership with other `EventProcessor`
        and start receiving EventData from EventHub and processing events.

        """
        while self._running:
            random_jitter = self._load_balancing_interval * random.random() * 0.2
            load_balancing_interval = self._load_balancing_interval + random_jitter
            try:
                claimed_partition_ids = self._ownership_manager.claim_ownership()
                if claimed_partition_ids:
                    existing_pids = set(self._consumers.keys())
                    claimed_pids = set(claimed_partition_ids)
                    to_cancel_pids = existing_pids - claimed_pids
                    newly_claimed_pids = claimed_pids - existing_pids
                    if newly_claimed_pids:
                        checkpoints = self._ownership_manager.get_checkpoints() if self._checkpoint_store else None
                        self._create_tasks_for_claimed_ownership(newly_claimed_pids, checkpoints)
                else:
                    _LOGGER.info(
                        "EventProcessor %r hasn't claimed an ownership. It keeps claiming.",
                        self._id,
                    )
                    to_cancel_pids = set(self._consumers.keys())
                if to_cancel_pids:
                    self._cancel_tasks_for_partitions(to_cancel_pids)
            except Exception as err:  # pylint:disable=broad-except
                # ownership_manager.get_checkpoints() and ownership_manager.claim_ownership() may raise exceptions
                # when there are load balancing and/or checkpointing (checkpoint_store isn't None).
                # They're swallowed here to retry every self._load_balancing_interval seconds.
                # Meanwhile this event processor won't lose the partitions it has claimed before.
                # If it keeps failing, other EventProcessors will start to claim ownership of the partitions
                # that this EventProcessor is working on. So two or multiple EventProcessors may be working
                # on the same partition for a short while.
                # Setting owner_level would create exclusive connection to the partition and
                # alleviate duplicate-receiving greatly.
                _LOGGER.warning(
                    "EventProcessor instance %r of eventhub %r consumer group %r. "
                    "An error occurred while load-balancing and claiming ownership. "
                    "The exception is %r. Retrying after %r seconds",
                    self._id,
                    self._eventhub_name,
                    self._consumer_group,
                    err,
                    load_balancing_interval,
                )
                self._process_error(None, err)  # type: ignore

            time.sleep(load_balancing_interval)

    def _close_consumer(self, partition_id: str, consumer: EventHubConsumer, reason: CloseReason) -> None:
        consumer.close()
        with self._lock:
            del self._consumers[partition_id]

        _LOGGER.info(
            "PartitionProcessor of EventProcessor instance %r of eventhub %r partition %r consumer group %r"
            " is being closed. Reason is: %r",
            self._id,
            self._partition_contexts[partition_id].eventhub_name,
            self._partition_contexts[partition_id].partition_id,
            self._partition_contexts[partition_id].consumer_group,
            reason,
        )

        if self._partition_close_handler:
            try:
                self._partition_close_handler(self._partition_contexts[partition_id], reason)
            except Exception as err:  # pylint:disable=broad-except
                _LOGGER.warning(
                    "EventProcessor instance %r of eventhub %r partition %r consumer group %r. "
                    "An error occurred while running on_partition_close. The exception is %r.",
                    self._id,
                    self._partition_contexts[partition_id].eventhub_name,
                    self._partition_contexts[partition_id].partition_id,
                    self._partition_contexts[partition_id].consumer_group,
                    err,
                )
                self._process_error(self._partition_contexts[partition_id], err)

        self._ownership_manager.release_ownership(partition_id)

    def _do_receive(self, partition_id: str, consumer: EventHubConsumer) -> None:
        """Call the consumer.receive() and handle exceptions if any after it exhausts retries.

        :param str partition_id: The partition id.
        :param ~azure.eventhub._consumer.EventHubConsumer consumer: The consumer object.
        """
        try:
            consumer.receive(self._batch, self._max_batch_size, self._max_wait_time)
        except Exception as error:  # pylint:disable=broad-except
            _LOGGER.warning(
                "EventProcessor instance %r of eventhub %r partition %r consumer group %r. "
                "An error occurred while receiving. The exception is %r.",
                self._id,
                self._partition_contexts[partition_id].eventhub_name,
                self._partition_contexts[partition_id].partition_id,
                self._partition_contexts[partition_id].consumer_group,
                error,
            )
            self._process_error(self._partition_contexts[partition_id], error)
            # TODO: close consumer if non-retryable. issue #27137
            # Does OWNERSHIP_LOST make sense for all errors?
            self._close_consumer(partition_id, consumer, CloseReason.OWNERSHIP_LOST)

    def start(self) -> None:
        if self._running:
            _LOGGER.info("EventProcessor %r has already started.", self._id)
            return

        _LOGGER.info("EventProcessor %r is being started", self._id)
        self._running = True
        thread = threading.Thread(target=self._load_balancing)
        thread.daemon = True
        thread.start()

        while self._running:
            for partition_id, consumer in list(self._consumers.items()):
                if consumer.stop:
                    self._close_consumer(partition_id, consumer, CloseReason.OWNERSHIP_LOST)
                    continue

                self._do_receive(partition_id, consumer)

        with self._lock:
            for partition_id, consumer in list(self._consumers.items()):
                self._close_consumer(partition_id, consumer, CloseReason.SHUTDOWN)

    def stop(self) -> None:
        """Stop the EventProcessor.

        The EventProcessor will stop receiving events from EventHubs and release the ownership of the partitions
        it is working on.
        Other running EventProcessor will take over these released partitions.

        A stopped EventProcessor can be restarted by calling method `start` again.

        """
        if not self._running:
            _LOGGER.info("EventProcessor %r has already been stopped.", self._id)
            return

        self._running = False
        _LOGGER.info("EventProcessor %r has been stopped.", self._id)
