#
# 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 functools
import inspect
import os
import sys
import tempfile
import threading
import time
from concurrent import futures
import logging
import grpc
from grpc import _common, _server
from grpc._cython.cygrpc import StatusCode
from grpc._server import _serialize_response, _status, _abort, _Context, _unary_request, \
    _select_thread_pool_for_behavior, _unary_response_in_pool
from typing import Dict

from notification_service.service import NotificationService

from ai_flow.protobuf.high_availability_pb2_grpc import add_HighAvailabilityManagerServicer_to_server
from ai_flow.endpoint.server.high_availability import SimpleAIFlowServerHaManager, HighAvailableService
from ai_flow.store.db.base_model import base
from ai_flow.store.sqlalchemy_store import SqlAlchemyStore
from ai_flow.store.mongo_store import MongoStore, MongoStoreConnManager
from ai_flow.store.db.db_util import extract_db_engine_from_uri, parse_mongo_uri
from ai_flow.endpoint.server.server_config import DBType
from notification_service.proto import notification_service_pb2_grpc

from ai_flow.metadata_store.service.service import MetadataService
from ai_flow.model_center.service.service import ModelCenterService
from ai_flow.metric.service.metric_service import MetricService
from ai_flow.scheduler.scheduler_service import SchedulerService, SchedulerServiceConfig

sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "../../..")))

from ai_flow.protobuf import model_center_service_pb2_grpc, \
    metadata_service_pb2_grpc, metric_service_pb2_grpc, scheduling_service_pb2_grpc

_PORT = '50051'
_ONE_DAY_IN_SECONDS = 60 * 60 * 24


class AIFlowServer(object):
    """
    Block/Async server of an AIFlow Rest Endpoint that provides Metadata/Model/Notification function service.
    """

    def __init__(self, store_uri=None, port=_PORT,
                 start_default_notification: bool = True,
                 notification_uri=None,
                 start_meta_service: bool = True,
                 start_model_center_service: bool = True,
                 start_metric_service: bool = True,
                 start_scheduler_service: bool = True,
                 scheduler_service_config: Dict = None,
                 enabled_ha: bool = False,
                 ha_manager=None,
                 ha_server_uri=None,
                 ha_storage=None,
                 ttl_ms: int = 10000):
        self.store_uri = store_uri
        self.db_type = DBType.value_of(extract_db_engine_from_uri(store_uri))
        self.executor = Executor(futures.ThreadPoolExecutor(max_workers=10))
        self.server = grpc.server(self.executor)
        self.start_default_notification = start_default_notification
        self.enabled_ha = enabled_ha
        server_uri = 'localhost:{}'.format(port)
        if start_default_notification:
            logging.info("start default notification service.")
            notification_service_pb2_grpc.add_NotificationServiceServicer_to_server(
                NotificationService.from_storage_uri(store_uri),
                self.server)
        if start_model_center_service:
            logging.info("start model center service.")
            model_center_service_pb2_grpc.add_ModelCenterServiceServicer_to_server(
                ModelCenterService(store_uri=store_uri,
                                   notification_uri=server_uri if start_default_notification
                                                                  and notification_uri is None else notification_uri),
                self.server)
        if start_meta_service:
            logging.info("start meta service.")
            metadata_service_pb2_grpc.add_MetadataServiceServicer_to_server(
                MetadataService(db_uri=store_uri, server_uri=server_uri), self.server)
        if start_metric_service:
            logging.info("start metric service.")
            metric_service_pb2_grpc.add_MetricServiceServicer_to_server(MetricService(db_uri=store_uri), self.server)

        if start_scheduler_service:
            self._add_scheduler_service(scheduler_service_config)

        if enabled_ha:
            self._add_ha_service(ha_manager, ha_server_uri, ha_storage, store_uri, ttl_ms)

        self.server.add_insecure_port('[::]:' + str(port))

    def _add_scheduler_service(self, scheduler_service_config):
        logging.info("start scheduler service.")
        # check the `scheduler` option of scheduler service config
        if scheduler_service_config is None:
            raise Exception(
                'The `scheduler` option of scheduler service config is not configured. '
                'Please add the `scheduler` option!')
        if 'scheduler_class' not in scheduler_service_config:
            raise Exception(
                'The `scheduler_class` option of scheduler service config is not configured. '
                'Please add the `scheduler_class` option under the `scheduler` option!')
        if 'scheduler_config' not in scheduler_service_config:
            scheduler_service_config['scheduler_config'] = {}
        real_config = SchedulerServiceConfig()
        real_config.set_scheduler_config(scheduler_service_config.get('scheduler_config'))
        real_config.set_repository(scheduler_service_config.get('repository'))
        real_config.set_scheduler_class(scheduler_service_config.get('scheduler_class'))
        self.scheduler_service = SchedulerService(real_config)
        scheduling_service_pb2_grpc.add_SchedulingServiceServicer_to_server(self.scheduler_service,
                                                                            self.server)

    def _add_ha_service(self, ha_manager, ha_server_uri, ha_storage, store_uri, ttl_ms):
        if ha_manager is None:
            ha_manager = SimpleAIFlowServerHaManager()
        if ha_server_uri is None:
            raise ValueError("ha_server_uri is required with ha enabled!")
        if ha_storage is None:
            db_engine = extract_db_engine_from_uri(store_uri)
            if DBType.value_of(db_engine) == DBType.MONGODB:
                username, password, host, port, db = parse_mongo_uri(store_uri)
                ha_storage = MongoStore(host=host,
                                        port=int(port),
                                        username=username,
                                        password=password,
                                        db=db)
            else:
                ha_storage = SqlAlchemyStore(store_uri)
        self.ha_service = HighAvailableService(ha_manager, ha_server_uri, ha_storage, ttl_ms)
        add_HighAvailabilityManagerServicer_to_server(self.ha_service, self.server)

    def run(self, is_block=False):
        if self.enabled_ha:
            self.ha_service.start()
        self.server.start()
        logging.info('AIFlow server started.')
        if is_block:
            try:
                while True:
                    time.sleep(_ONE_DAY_IN_SECONDS)
            except KeyboardInterrupt:
                self.stop()
        else:
            pass

    def stop(self, clear_sql_lite_db_file=False):
        self.executor.shutdown()
        self.server.stop(0)
        if self.enabled_ha:
            self.ha_service.stop()

        if self.db_type == DBType.SQLITE and clear_sql_lite_db_file:
            store = SqlAlchemyStore(self.store_uri)
            base.metadata.drop_all(store.db_engine)
            os.remove(self.store_uri[10:])
        elif self.db_type == DBType.MONGODB:
            MongoStoreConnManager().disconnect_all()

        logging.info('AIFlow server stopped.')

    def _clear_db(self):
        if self.db_type == DBType.SQLITE:
            store = SqlAlchemyStore(self.store_uri)
            base.metadata.drop_all(store.db_engine)
            base.metadata.create_all(store.db_engine)
        elif self.db_type == DBType.MONGODB:
            MongoStoreConnManager().drop_all()


def _loop(loop: asyncio.AbstractEventLoop):
    asyncio.set_event_loop(loop)
    if not loop.is_running() or loop.is_closed():
        loop.run_forever()
    pending = asyncio.all_tasks(loop=loop)
    if pending:
        loop.run_until_complete(asyncio.gather(*pending))


class Executor(futures.Executor):
    def __init__(self, thread_pool, loop=None):
        super().__init__()
        self._shutdown = False
        self._thread_pool = thread_pool
        self._loop = loop or asyncio.get_event_loop()
        if not self._loop.is_running() or self._loop.is_closed():
            self._thread = threading.Thread(target=_loop, args=(self._loop,), daemon=True)
            self._thread.start()

    def submit(self, fn, *args, **kwargs):
        if self._shutdown:
            raise RuntimeError('Cannot schedule new futures after shutdown.')
        if not self._loop.is_running():
            raise RuntimeError('Loop must be started before any function could be submitted.')
        if inspect.iscoroutinefunction(fn):
            coroutine = fn(*args, **kwargs)
            return asyncio.run_coroutine_threadsafe(coroutine, self._loop)
        else:
            func = functools.partial(fn, *args, **kwargs)
            return self._loop.run_in_executor(self._thread_pool, func)

    def shutdown(self, wait=True):
        self._shutdown = True
        if wait:
            self._thread_pool.shutdown()


async def _call_behavior_async(rpc_event, state, behavior, argument, request_deserializer):
    context = _Context(rpc_event, state, request_deserializer)
    try:
        return await behavior(argument, context), True
    except Exception as e:
        with state.condition_type:
            if e not in state.rpc_errors:
                logging.exception(e)
                _abort(state, rpc_event.operation_call, StatusCode.unknown, _common.encode(e))
        return None, False


async def _unary_response_in_pool_async(rpc_event, state, behavior, argument_thunk, request_deserializer,
                                        response_serializer):
    argument = argument_thunk()
    if argument is not None:
        response, proceed = await _call_behavior_async(rpc_event, state, behavior, argument, request_deserializer)
        if proceed:
            serialized_response = _serialize_response(rpc_event, state, response, response_serializer)
            if serialized_response is not None:
                _status(rpc_event, state, serialized_response)


def _handle_unary_unary(rpc_event, state, method_handler, default_thread_pool):
    unary_request = _unary_request(rpc_event, state, method_handler.request_deserializer)
    thread_pool = _select_thread_pool_for_behavior(method_handler.unary_unary, default_thread_pool)
    if asyncio.iscoroutinefunction(method_handler.unary_unary):
        return thread_pool.submit(_unary_response_in_pool_async, rpc_event, state, method_handler.unary_unary,
                                  unary_request, method_handler.request_deserializer,
                                  method_handler.response_serializer)
    else:
        return thread_pool.submit(_unary_response_in_pool, rpc_event, state, method_handler.unary_unary, unary_request,
                                  method_handler.request_deserializer, method_handler.response_serializer)


_server._handle_unary_unary = _handle_unary_unary

if __name__ == '__main__':
    fd, temp_db_file = tempfile.mkstemp()
    os.close(fd)
    store_uri = '%s%s' % ('sqlite:///', temp_db_file)
    server = AIFlowServer(store_uri=store_uri)
    server.run(is_block=True)
