import os
import logging

import grpc
from modelscope.models import Model
from modelscope.pipelines import pipeline
from modelscope.preprocessors import TokenClassificationTransformersPreprocessor
from modelscope.utils.constant import Tasks

from grpc_py_gen import segment_pb2, segment_pb2_grpc

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


class BertWordSegmenter:
    def __init__(self, model_id="damo/nlp_structbert_word-segmentation_chinese-base", model_path=cache_model_path):
        """
        Init word segmenter
        :param model_id:
        :param model_path:
        """
        self.model = Model.from_pretrained(model_path or model_id)
        self.tokenizer = TokenClassificationTransformersPreprocessor(self.model.model_dir)
        self.pipeline_inst = pipeline(task=Tasks.word_segmentation, model=self.model, preprocessor=self.tokenizer)

    def segment(self, text: str) -> dict:
        """
        segment text
        :param text:
        :return:
        """
        return self.pipeline_inst(input=text)


class WordSegmentationService(segment_pb2_grpc.WordSegmentationServicer):
    def __init__(self, logger: logging.Logger, max_batch_size: int = 120):
        self.bert_ws = BertWordSegmenter()
        self.logger = logger
        self.max_batch_size = max_batch_size

    def Segment(self, request, context) -> segment_pb2.SegmentResponse:
        """
        Segment text
        :param request:
        :param context:
        :return:
        """
        if not request.input:
            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
            context.set_details("Input text cannot be empty")
            return segment_pb2.SegmentResponse()

        if len(request.input) > self.max_batch_size:
            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
            context.set_details(f"Input text exceeds maximum length of {self.max_batch_size}")
            return segment_pb2.SegmentResponse()

        try:
            result = self.bert_ws.segment(request.input)
            return segment_pb2.SegmentResponse(output=result["output"])
        except Exception as e:
            context.abort(grpc.StatusCode.INTERNAL, f"Segmentation failed: {str(e)}")
            return segment_pb2.SegmentResponse()
