import os
import logging

from modelscope import Model
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from grpc_py_gen import ner_pb2, ner_pb2_grpc

cache_model_path = os.path.expanduser("~/nlp_models/ner")


class NER:
    def __init__(
        self, model_id="iic/nlp_raner_named-entity-recognition_chinese-base-news", model_path=cache_model_path
    ):
        """
        :param model_id:
        :param model_path:
        """
        self.model = Model.from_pretrained(model_path or model_id)
        self.pipeline_inst = pipeline(task=Tasks.named_entity_recognition, model=self.model)

    def extract(self, text: str) -> dict:
        """
        return NER result
        :param text:
        :return:
        """
        return self.pipeline_inst(text)


class NERService(ner_pb2_grpc.NerServiceServicer):
    def __init__(self, logger: logging.Logger):
        self.logger = logger
        self.ner = NER()

    @staticmethod
    def convert_to_ner_response(model_output: dict):
        ner_response = ner_pb2.NerResponse(
            entity_list=ner_pb2.EntityList(
                entities=[
                    ner_pb2.Entity(
                        type=item["type"],
                        start=int(item["start"]),  # Convert np.int64 to int
                        end=int(item["end"]),  # Convert np.int64 to int
                        prob=float(item["prob"]),  # Convert np.float32 to float
                        span=item["span"],  # string, no conversion needed
                    )
                    for item in model_output["output"]
                ]
            )
        )
        return ner_response

    def ExtractEntities(self, request: ner_pb2.NerRequest, _context):
        ret = self.ner.extract(request.text)
        return NERService.convert_to_ner_response(ret)
