#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
#
import asyncio
import time
import traceback

from notification_service.base_notification import BaseEvent
from notification_service.event_storage import BaseEventStorage, DbEventStorage
from notification_service.high_availability import NotificationServerHaManager
from notification_service.mongo_event_storage import MongoEventStorage
from notification_service.proto import notification_service_pb2_grpc, notification_service_pb2
from notification_service.util.utils import event_to_proto, event_list_to_proto, member_to_proto, event_proto_to_event, \
    proto_to_member, count_list_to_proto
from notification_service.util.db import extract_db_engine_from_uri, DBType, parse_mongo_uri


class NotificationService(notification_service_pb2_grpc.NotificationServiceServicer):

    def __init__(self, storage: BaseEventStorage):
        self.storage = storage
        self.notification_conditions = {}
        self.lock = asyncio.Lock()
        self.write_condition = asyncio.Condition()

    @classmethod
    def from_storage_uri(cls, storage_uri: str) -> 'NotificationService':
        """
        Construct the notification service with the given storage uri
        :param storage_uri: uri of the backend storage to use
        :type storage_uri: str
        :rtype: NotificationService
        """
        db_engine = extract_db_engine_from_uri(storage_uri)
        if DBType.value_of(db_engine) == DBType.MONGODB:
            username, password, host, port, db = parse_mongo_uri(storage_uri)
            storage = MongoEventStorage(host=host,
                                        port=int(port),
                                        username=username,
                                        password=password,
                                        db=db)
            return cls(storage=storage)
        else:
            return cls(storage=DbEventStorage(storage_uri))

    def start(self):
        pass

    def stop(self):
        pass

    @asyncio.coroutine
    def sendEvent(self, request, context):
        try:
            return self._send_event(request)
        except Exception as e:
            print(e)
            traceback.print_stack()
            return notification_service_pb2.SendEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _send_event(self, request):
        event_proto = request.event
        event = BaseEvent(
            key=event_proto.key,
            value=event_proto.value,
            event_type=None if event_proto.event_type == "" else event_proto.event_type,
            context=None if event_proto.context == "" else event_proto.context,
            namespace=None if event_proto.namespace == "" else event_proto.namespace,
            sender=None if event_proto.sender == "" else event_proto.sender
        )
        uuid = request.uuid
        enable_idempotence = request.enable_idempotence
        key = event.key
        return_msg = ''
        # Lock conditions dict for get/check/update of key
        await self.lock.acquire()
        if self.notification_conditions.get(key) is None:
            self.notification_conditions.update({(key, asyncio.Condition())})
        # Release lock after check/update key of notification conditions dict
        self.lock.release()
        async with self.notification_conditions.get(key), self.write_condition:
            if enable_idempotence and self.storage.get_event_by_uuid(uuid) is not None:
                return_msg = 'Ignored because event already exists.'
            else:
                event: BaseEvent = self.storage.add_event(event, uuid)
            self.notification_conditions.get(key).notify_all()
            self.write_condition.notify_all()

        result_event_proto = event_to_proto(event)
        return notification_service_pb2.SendEventsResponse(
            event=result_event_proto,
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg=return_msg)

    @asyncio.coroutine
    def listEvents(self, request, context):
        try:
            return self._list_events(request)
        except Exception as e:
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _list_events(self, request):
        keys = request.keys
        event_type = request.event_type
        start_time = request.start_time
        start_version = request.start_version
        namespace = None if request.namespace == '' else request.namespace
        sender = None if request.sender == '' else request.sender
        timeout_seconds = request.timeout_seconds

        if timeout_seconds == 0:
            event_models = self._query_events(keys, event_type, start_time, start_version, namespace, sender)
            event_proto_list = event_list_to_proto(event_models)
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                events=event_proto_list)
        else:
            start = time.time()
            # Lock conditions dict for get/check/update of key
            await self.lock.acquire()
            for key in keys:
                if self.notification_conditions.get(key) is None:
                    self.notification_conditions.update({(key, asyncio.Condition())})
            # Release lock after check/update key of notification conditions dict
            self.lock.release()
            event_models = []
            if len(keys) == 1:
                key = keys[0]
                condition = self.notification_conditions.get(key)
            else:
                condition = self.write_condition
            async with condition:
                while time.time() - start < timeout_seconds and len(event_models) == 0:
                    try:
                        await asyncio.wait_for(condition.wait(),
                                               timeout_seconds - time.time() + start)
                        event_models = self._query_events(
                            keys, event_type, start_time, start_version, namespace, sender)
                    except asyncio.TimeoutError:
                        pass
                if len(event_models) == 0:
                    event_models = self._query_events(
                        keys, event_type, start_time, start_version, namespace, sender)
            event_proto_list = event_list_to_proto(event_models)
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                events=event_proto_list)

    def _query_events(self, keys, event_type, start_time, start_version, namespace, sender):
        return self.storage.list_events(keys, start_version, event_type, start_time, namespace, sender)

    def countEvents(self, request, context):
        keys = request.keys
        event_type = request.event_type
        start_time = request.start_time
        start_version = request.start_version
        namespace = None if request.namespace == '' else request.namespace
        sender = None if request.sender == '' else request.sender
        try:
            event_counts = self.storage.count_events(keys, start_version, event_type, start_time, namespace, sender)
            event_count, count_proto_list = count_list_to_proto(event_counts)
            return notification_service_pb2.CountEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                event_count=event_count,
                sender_event_counts=count_proto_list)
        except Exception as e:
            return notification_service_pb2.CountEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    @asyncio.coroutine
    def listAllEvents(self, request, context):
        try:
            return self._list_all_events(request)
        except Exception as e:
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _list_all_events(self, request):
        start_time = request.start_time
        start_version = request.start_version
        end_version = request.end_version
        timeout_seconds = request.timeout_seconds
        if 0 == timeout_seconds:
            event_models = self._query_all_events(start_time, start_version, end_version)
            event_proto_list = event_list_to_proto(event_models)
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                events=event_proto_list)
        else:
            start = time.time()
            event_models = self._query_all_events(start_time, start_version, end_version)
            async with self.write_condition:
                while time.time() - start < timeout_seconds and len(event_models) == 0:
                    try:
                        await asyncio.wait_for(self.write_condition.wait(),
                                               timeout_seconds - time.time() + start)
                        event_models = self._query_all_events(
                            start_time, start_version, end_version)
                    except asyncio.TimeoutError:
                        pass
            event_proto_list = event_list_to_proto(event_models)
            return notification_service_pb2.ListEventsResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                events=event_proto_list)

    def _query_all_events(self, start_time, start_version, end_version):
        if start_version > 0:
            return self.storage.list_all_events_from_version(start_version, end_version)
        else:
            return self.storage.list_all_events(start_time)

    @asyncio.coroutine
    def getLatestVersionByKey(self, request, context):
        try:
            return self._get_latest_version_by_key(request)
        except Exception as e:
            return notification_service_pb2.GetLatestVersionResponse(
                return_code=str(notification_service_pb2.ReturnStatus.ERROR), return_msg=str(e))

    async def _get_latest_version_by_key(self, request):
        key = request.key
        namespace = request.namespace
        timeout_seconds = request.timeout_seconds
        if 0 == timeout_seconds:
            latest_version = self._query_latest_version_by_key(key, namespace)
            return notification_service_pb2.GetLatestVersionResponse(
                return_code=str(notification_service_pb2.ReturnStatus.SUCCESS),
                return_msg='',
                version=latest_version)
        else:
            start = time.time()
            latest_version = self._query_latest_version_by_key(key, namespace)
            async with self.write_condition:
                while time.time() - start < timeout_seconds and latest_version == 0:
                    try:
                        await asyncio.wait_for(self.write_condition.wait(), timeout_seconds - time.time() + start)
                        latest_version = self._query_latest_version_by_key(key, namespace)
                    except asyncio.TimeoutError:
                        pass
            return notification_service_pb2.ListEventsResponse(
                return_code=str(notification_service_pb2.ReturnStatus.SUCCESS),
                return_msg='',
                version=latest_version)

    def _query_latest_version_by_key(self, key, namespace):
        if len(namespace) == 0:
            namespace = None
        return self.storage.get_latest_version(key=key, namespace=namespace)

    @asyncio.coroutine
    def notify(self, request, context):
        try:
            return self._notify(request)
        except Exception as e:
            return notification_service_pb2.CommonResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _notify(self, request):
        for notify in request.notifies:
            if notify.key in self.notification_conditions:
                async with self.notification_conditions.get(notify.key):
                    self.notification_conditions.get(notify.key).notify_all()
        async with self.write_condition:
            self.write_condition.notify_all()
        return notification_service_pb2.CommonResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='')

    @asyncio.coroutine
    def listMembers(self, request, context):
        try:
            return self._list_members(request)
        except Exception as e:
            return notification_service_pb2.ListMembersResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _list_members(self, request):
        # this method is used in HA mode, so we just return an empty list here.
        return notification_service_pb2.ListMembersResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='',
            members=[])

    @asyncio.coroutine
    def notifyNewMember(self, request, context):
        try:
            return self._notify_new_member(request)
        except Exception as e:
            return notification_service_pb2.CommonResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _notify_new_member(self, request):
        # this method is used in HA mode, so we just return here.
        return notification_service_pb2.CommonResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='')

    @asyncio.coroutine
    def registerClient(self, request, context):
        try:
            return self._register_client(request)
        except Exception as e:
            return notification_service_pb2.RegisterClientResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _register_client(self, request):
        client_id = self.storage.register_client(request.client_meta.namespace, request.client_meta.sender)
        return notification_service_pb2.RegisterClientResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='',
            client_id=client_id)

    @asyncio.coroutine
    def deleteClient(self, request, context):
        try:
            return self._delete_client(request)
        except Exception as e:
            return notification_service_pb2.CommonResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e))

    async def _delete_client(self, request):
        self.storage.delete_client(request.client_id)
        return notification_service_pb2.CommonResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='')

    @asyncio.coroutine
    def isClientExists(self, request, context):
        try:
            return self._is_client_exists(request)
        except Exception as e:
            return notification_service_pb2.isClientExistsResponse(
                return_code=notification_service_pb2.ReturnStatus.ERROR, return_msg=str(e), is_exists=False)

    async def _is_client_exists(self, request):
        is_exists = self.storage.is_client_exists(request.client_id)
        return notification_service_pb2.isClientExistsResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='',
            is_exists=is_exists)


class HighAvailableNotificationService(NotificationService):

    def __init__(self,
                 storage,
                 ha_manager,
                 server_uri,
                 ha_storage,
                 ttl_ms: int = 10000,
                 min_notify_interval_ms: int = 100):
        super(HighAvailableNotificationService, self).__init__(storage)
        self.server_uri = server_uri
        self.ha_storage = ha_storage
        self.ttl_ms = ttl_ms
        self.min_notify_interval_ms = min_notify_interval_ms
        self.ha_manager = ha_manager  # type: NotificationServerHaManager
        self.member_updated_condition = asyncio.Condition()

    def start(self):
        self.ha_manager.start(self.server_uri,
                              self.ha_storage,
                              self.ttl_ms,
                              self.min_notify_interval_ms,
                              self.member_updated_condition)

    def stop(self):
        self.ha_manager.stop()

    async def _send_event(self, request):
        response = await super(HighAvailableNotificationService, self)._send_event(request)
        try:
            if response.return_code == notification_service_pb2.ReturnStatus.SUCCESS:
                self.ha_manager.notify_others(event_proto_to_event(response.event))
        finally:
            return response

    async def _list_members(self, request):
        timeout_seconds = request.timeout_seconds
        if 0 == timeout_seconds:
            return notification_service_pb2.ListMembersResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                members=[member_to_proto(member) for member in self.ha_manager.get_living_members()])
        else:
            start = time.time()
            members = self.ha_manager.get_living_members()
            async with self.member_updated_condition:
                while time.time() - start < timeout_seconds:
                    try:
                        await asyncio.wait_for(self.member_updated_condition.wait(),
                                               timeout_seconds - time.time() + start)
                        members = self.ha_manager.get_living_members()
                        break
                    except asyncio.TimeoutError:
                        pass
            return notification_service_pb2.ListMembersResponse(
                return_code=notification_service_pb2.ReturnStatus.SUCCESS,
                return_msg='',
                members=[member_to_proto(member) for member in members])

    async def _notify_new_member(self, request):
        self.ha_manager.add_living_member(proto_to_member(request.member))
        async with self.member_updated_condition:
            self.member_updated_condition.notify_all()
        return notification_service_pb2.CommonResponse(
            return_code=notification_service_pb2.ReturnStatus.SUCCESS,
            return_msg='')
