# -*- coding: utf-8 -*-
# 
# @author: 
# @date: 
# @version: 
# @description:
import json
import os.path

from knowledgebase.db.doc_db_helper import doc_dbh
from knowledgebase.db.doc_db_models import TEntity
from knowledgebase.log import Log


class EntityHelper:
    # 文档类型和识别提示词map
    doc_prompt_map: dict

    def __init__(self):
        Log.info("初始化EntityHelper")
        current_dir = os.path.dirname(__file__)
        self.doc_prompt_map = {}
        with open(f'{current_dir}/../../tpl/entities.json', 'r', encoding='utf-8') as f:
            text = f.read()
            obj = json.loads(text)
            for ty in obj:
                obj2 = obj[ty]
                for doc_ty in obj2:
                    prompts = obj2[doc_ty]['prompts']
                    self.doc_prompt_map[doc_ty] = prompts
                    for entity in obj2[doc_ty]['entities']:
                        _entity = TEntity(name=entity, type=ty, doc_type=doc_ty,
                                          prompts=obj2[doc_ty]['entities'][entity])
                        if doc_dbh.get_entity(_entity):
                            continue
                        doc_dbh.add_entity(_entity)
                        Log.info(f"新增Entity：{entity}，id：{_entity.id}")

_entity_helper:EntityHelper = None

def get_entity_helper():
    global _entity_helper
    return _entity_helper

def init_entity_helper():
    global _entity_helper
    _entity_helper = EntityHelper()