import textwrap


def get_ner_name_prompt(text: str) -> str:
    """
    Returns a prompt that generates a mapping of original names/companies to their masked versions.
    
    Args:
        text (str): The input text to be analyzed for masking
        
    Returns:
        str: The formatted prompt that will generate a mapping dictionary
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体识别助手。请从以下文本中抽取出所有需要脱敏的敏感信息，并按照指定的类别进行分类。请严格按照JSON格式输出结果。

实体类别包括:
- 人名 (不包括律师、法官、书记员、检察官等公职人员)
- 英文人名


待处理文本:
{text}

输出格式:
{{
"entities": [
    {{"text": "原始文本内容", "type": "人名"}},
    {{"text": "原始文本内容", "type": "英文人名"}},
    ...
  ]
}}



请严格按照JSON格式输出结果。
    
    """)
    return prompt.format(text=text)


def get_ner_company_prompt(text: str) -> str:
    """
    Returns a prompt that generates a mapping of original companies to their masked versions.
    
    Args:
        text (str): The input text to be analyzed for masking
        
    Returns:
        str: The formatted prompt that will generate a mapping dictionary
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体识别助手。请从以下文本中抽取出所有需要脱敏的敏感信息，并按照指定的类别进行分类。请严格按照JSON格式输出结果。

实体类别包括:
- 公司名称
- 英文公司名称
- Company with English name
- 公司名称简称
- 公司英文名称简称


待处理文本:
{text}  

输出格式:
{{
"entities": [
    {{"text": "原始文本内容", "type": "公司名称"}},
    {{"text": "原始文本内容", "type": "英文公司名称"}},
    {{"text": "原始文本内容", "type": "公司名称简称"}},
    {{"text": "原始文本内容", "type": "公司英文名称简称"}},
    ...
  ]
}}

请严格按照JSON格式输出结果。
    """)
    return prompt.format(text=text)


def get_ner_address_prompt(text: str) -> str:
    """
    Returns a prompt that generates a mapping of original addresses to their masked versions.
    
    Args:
        text (str): The input text to be analyzed for masking
        
    Returns:
        str: The formatted prompt that will generate a mapping dictionary
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体识别助手。请从以下文本中抽取出所有需要脱敏的敏感信息，并按照指定的类别进行分类。请严格按照JSON格式输出结果。

实体类别包括:
- 地址


待处理文本:
{text}  

输出格式:
{{
"entities": [
    {{"text": "原始文本内容", "type": "地址"}},
    ...
  ]
}}

请严格按照JSON格式输出结果。
    """)
    return prompt.format(text=text)


def get_ner_project_prompt(text: str) -> str:
    """
    Returns a prompt that generates a mapping of original project names to their masked versions.
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体识别助手。请从以下文本中抽取出所有需要脱敏的敏感信息，并按照指定的类别进行分类。请严格按照JSON格式输出结果。

实体类别包括:
- 项目名(此处项目特指商业、工程、合同等项目)

待处理文本:
{text}  

输出格式:
{{
"entities": [
    {{"text": "原始文本内容", "type": "项目名"}},
    ...
  ]
}}

请严格按照JSON格式输出结果。
    """)
    return prompt.format(text=text)


def get_ner_case_number_prompt(text: str) -> str:
    """
    Returns a prompt that generates a mapping of original case numbers to their masked versions.
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体识别助手。请从以下文本中抽取出所有需要脱敏的敏感信息，并按照指定的类别进行分类。请严格按照JSON格式输出结果。

实体类别包括:
- 案号

待处理文本:
{text}  

输出格式:
{{
"entities": [
    {{"text": "原始文本内容", "type": "案号"}},
    ...
  ]
}}

请严格按照JSON格式输出结果。
    """)
    return prompt.format(text=text)


def get_entity_linkage_prompt(entities_text: str) -> str:
    """
    Returns a prompt that identifies related entities and groups them together.
    
    Args:
        entities_text (str): The list of entities to be analyzed for linkage
        
    Returns:
        str: The formatted prompt that will generate entity linkage information
    """
    prompt = textwrap.dedent("""
你是一个专业的法律文本实体关联分析助手。请分析以下实体列表，识别出相互关联的实体（如全称与简称、中文名与英文名等），并将它们分组。

关联规则：
1. 公司名称关联：
   - 全称与简称（如："阿里巴巴集团控股有限公司" 与 "阿里巴巴"）
   - 中文名与英文名（如："腾讯科技有限公司" 与 "Tencent Technology Ltd."）
   - 母公司与子公司（如："腾讯" 与 "腾讯音乐"）


2. 每个组中应指定一个主要实体（is_primary: true），通常是：
   - 对于公司：选择最正式的全称
   - 对于人名：选择最常用的称呼

待分析实体列表:
{entities_text}

输出格式:
{{
"entity_groups": [
    {{
        "group_id": "group_1",
        "group_type": "公司名称",
        "entities": [
            {{
                "text": "阿里巴巴集团控股有限公司",
                "type": "公司名称",
                "is_primary": true
            }},
            {{
                "text": "阿里巴巴",
                "type": "公司名称简称",
                "is_primary": false
            }}
        ]
    }}
]
}}

注意事项：
1. 只对确实有关联的实体进行分组
2. 每个实体只能属于一个组
3. 每个组必须有且仅有一个主要实体（is_primary: true）
4. 如果实体之间没有明显关联，不要强制分组
5. group_type 应该是 "公司名称" 

请严格按照JSON格式输出结果。
    """)
    return prompt.format(entities_text=entities_text)