# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from threading import Event
from typing import List, Union

try:
    from mycroft_bus_client.client.waiter import MessageWaiter as _MessageWaiterBase
except ImportError:
    # TODO - code in the wild does isinstance checks
    # this conditional subclassing should be removed ASAP, it is only here for the migration period
    # mycroft_bus_client is abandonware until further notice from MycroftAI

    class _MessageWaiterBase:
        pass


class MessageWaiter:
    """Wait for a single message.

    Encapsulate the wait for a message logic separating the setup from
    the actual waiting act so the waiting can be setup, actions can be
    performed and _then_ the message can be waited for.

    Arguments:
        bus: Bus to check for messages on
        message_type: message type(s) to wait for
    """
    def __init__(self, bus, message_type: Union[str, List[str]]):
        self.bus = bus
        if not isinstance(message_type, list):
            message_type = [message_type]
        self.msg_type = message_type
        self.received_msg = None
        # Setup response handler
        self.response_event = Event()
        for msg in self.msg_type:
            self.bus.once(msg, self._handler)

    def _handler(self, message):
        """Receive response data."""
        self.received_msg = message
        self.response_event.set()

    def wait(self, timeout=3.0):
        """Wait for message.

        Arguments:
            timeout (int or float): seconds to wait for message

        Returns:
            Message or None
        """
        self.response_event.wait(timeout)
        if not self.response_event.is_set():
            # Clean up the event handler
            try:
                for msg in self.msg_type:
                    self.bus.remove(msg, self._handler)
            except (ValueError, KeyError):
                # ValueError occurs on pyee 5.0.1 removing handlers
                # registered with once.
                # KeyError may theoretically occur if the event occurs as
                # the handler is removed
                pass
        return self.received_msg
