import signal
import logging
from concurrent import futures

import grpc

from grpc_py_gen import ner_pb2_grpc, segment_pb2_grpc, dependency_pb2_grpc
from src.service.ner import NERService
from src.service.dependency_parser import DependencyService
from src.service.bert_word_segmenter import WordSegmentationService


class GrpcServer:
    def __init__(self, logger: logging.Logger, host: str = "[::]", port: int = 50051, max_workers: int = 200):
        self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=max_workers))
        self.host = host
        self.port = port
        self.logger = logger
        segment_pb2_grpc.add_WordSegmentationServicer_to_server(WordSegmentationService(logger), self.server)
        dependency_pb2_grpc.add_DependencyServiceServicer_to_server(DependencyService(logger), self.server)
        ner_pb2_grpc.add_NerServiceServicer_to_server(NERService(logger), self.server)

    def start(self):
        addr = f"{self.host}:{self.port}"
        self.server.add_insecure_port(addr)
        self.server.start()
        self.logger.info(f"Server started on {addr}")

        def handle_signal(_signum, _frame):
            self.logger.info("Received signal to terminate")
            self.stop()

        signal.signal(signal.SIGTERM, handle_signal)
        signal.signal(signal.SIGINT, handle_signal)

        try:
            self.server.wait_for_termination()
        except KeyboardInterrupt:
            self.stop()

    def stop(self):
        self.logger.info("Stopping server...")
        self.server.stop(grace=5)
        self.logger.info("Server stopped")
