# -*- 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
    # 所有实体
    entities: list[TEntity]

    def __init__(self):
        Log.info("初始化EntityHelper")
        current_dir = os.path.dirname(__file__)
        self.entities = doc_dbh.get_all_entities()
        self.doc_prompt_map = {}
        entity_names = [entity.name for entity in self.entities]
        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']:
                        if entity in entity_names:
                            continue
                        _entity = TEntity(name=entity, type=ty, doc_type=doc_ty,
                                          prompts=obj2[doc_ty]['entities'][entity])
                        doc_dbh.add_entity(_entity)
                        Log.info(f"新增Entity：{entity}，id：{_entity.id}")


entity_helper = EntityHelper()