import os
import logging

import grpc
import stanza

from grpc_py_gen import dependency_pb2, dependency_pb2_grpc

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


class DependencyParser:
    def __init__(self, logger: logging.Logger):
        self.logger = logger
        self.default_lang = {"zh-hans": None, "en": None}
        for lang in self.default_lang:
            nlp = stanza.Pipeline(lang=lang, processors="tokenize,pos,lemma,depparse", model_dir=cache_model_path)
            self.default_lang[lang] = nlp

    def process(self, lang: str, text: str):
        if lang not in self.default_lang:
            raise ValueError(f"Unsupported language: {lang}")
        self.logger.info(f"Processing text in {lang}")
        doc = self.default_lang[lang](text)
        return doc


class DependencyService(dependency_pb2_grpc.DependencyServiceServicer):
    def __init__(self, logger: logging.Logger):
        self.logger = logger
        self.dependency_parser = DependencyParser(logger)

    def ParseDependencies(
        self, request: dependency_pb2.ParseDependenciesRequest, context
    ) -> dependency_pb2.ParseDependenciesResponse:
        """
        Get dependency relations
        :param request:
        :param context:
        :return:
        """
        if not request.lang or not request.text:
            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
            context.set_details("Both input texts cannot be empty")
            return dependency_pb2.ParseDependenciesResponse()

        try:
            doc = self.dependency_parser.process(request.lang, request.text)
            ret = []
            for sent in doc.sentences:
                dependency_list = dependency_pb2.DependencyList()
                for dep_edge in sent.dependencies:
                    triples = dependency_pb2.Dependency(word=dep_edge[2].text, gov_idx=dep_edge[0].id, rel=dep_edge[1])
                    dependency_list.dependencies.append(triples)
                ret.append(dependency_list)

            return dependency_pb2.ParseDependenciesResponse(dependency_list=ret)
        except Exception as e:
            context.abort(grpc.StatusCode.INTERNAL, f"Failed to get dependencies: {str(e)}")
            return dependency_pb2.ParseDependenciesResponse()
