# -*- coding: utf-8 -*-
"""
实体规范化器
负责实体的去重、ID生成和规范化
"""

import json
import hashlib
import logging
from typing import Dict, List, Any, Optional
from pathlib import Path
import yaml


class EntityNormalizer:
    """实体规范化器"""

    def __init__(self, ontology_file: str):
        """
        初始化规范化器

        Args:
            ontology_file: 本体配置文件路径
        """
        # 加载本体配置
        with open(ontology_file, 'r', encoding='utf-8') as f:
            self.ontology = yaml.safe_load(f)

        self.logger = logging.getLogger(__name__)

        # 实体注册表：{entity_id: entity_data}
        self.entity_registry = {}

        # 文本到ID的映射：{normalized_text: entity_id}
        self.text_to_id_map = {}

        # 类型计数器：{entity_type: count}
        self.type_counters = {}

        self.logger.info("EntityNormalizer初始化完成")

    def normalize_text(self, text: str) -> str:
        """
        标准化文本

        Args:
            text: 原始文本

        Returns:
            标准化后的文本
        """
        if not text:
            return ""

        # 去除首尾空格
        text = text.strip()

        # 统一全角/半角字符
        text = text.replace('　', ' ').replace('，', ',').replace('、', ',')

        # 转换为小写进行比较（保持原文用于显示）
        return text

    def calculate_similarity(self, text1: str, text2: str) -> float:
        """
        计算两个文本的相似度（基于编辑距离）

        Args:
            text1: 文本1
            text2: 文本2

        Returns:
            相似度（0-1之间）
        """
        import Levenshtein

        norm1 = self.normalize_text(text1)
        norm2 = self.normalize_text(text2)

        if not norm1 or not norm2:
            return 0.0

        max_len = max(len(norm1), len(norm2))
        if max_len == 0:
            return 1.0

        distance = Levenshtein.distance(norm1, norm2)
        similarity = 1.0 - (distance / max_len)

        return similarity

    def generate_entity_id(self, entity_type: str, entity_text: str) -> str:
        """
        生成实体ID

        Args:
            entity_type: 实体类型
            entity_text: 实体文本

        Returns:
            实体ID
        """
        # 获取类型配置
        type_config = self.ontology['entity_types'].get(entity_type, {})
        prefix = type_config.get('prefix', 'UNK')

        # 生成哈希值
        text_hash = abs(hash(entity_text)) % 100000

        # 检查是否需要计数器（某些类型可能需要）
        if entity_type not in self.type_counters:
            self.type_counters[entity_type] = 0

        # 根据类型生成ID
        if entity_type in ['Ethnic_Group', 'Geographic_Location', 'Geographic_Environment', 'Time_Period']:
            # 这些类型使用唯一哈希，不需要计数器
            entity_id = f"{prefix}-{text_hash:05d}"
        else:
            # 其他类型使用计数器
            self.type_counters[entity_type] += 1
            counter = self.type_counters[entity_type]
            entity_id = f"{prefix}-{text_hash:05d}-{counter}"

        return entity_id

    def find_similar_entity(
        self,
        entity_text: str,
        entity_type: str,
        similarity_threshold: float = 0.85
    ) -> Optional[str]:
        """
        查找相似实体

        Args:
            entity_text: 实体文本
            entity_type: 实体类型
            similarity_threshold: 相似度阈值

        Returns:
            相似实体的ID，如果不存在返回None
        """
        normalized_text = self.normalize_text(entity_text)

        # 只在相同类型的实体中查找
        for entity_id, entity_data in self.entity_registry.items():
            if entity_data['type'] != entity_type:
                continue

            # 检查文本相似度
            similarity = self.calculate_similarity(normalized_text, entity_data['canonical_name'])

            if similarity >= similarity_threshold:
                self.logger.info(f"发现相似实体: '{entity_text}' ~ '{entity_data['canonical_name']}' (相似度: {similarity:.2f})")
                return entity_id

        return None

    def is_generic_concept(self, entity_text: str, entity_type: str) -> bool:
        """
        检查是否为泛指概念（应该被过滤）

        Args:
            entity_text: 实体文本
            entity_type: 实体类型

        Returns:
            True表示是泛指概念，应该过滤
        """
        generic_patterns = [
            '东北少数民族',
            '本地土著民族',
            '当地民族',
            '少数民族',
            '土著',
            '东北民族',
            '本地民族',
            '地区民族',
        ]

        normalized_text = self.normalize_text(entity_text)

        # 检查是否匹配泛指模式
        for pattern in generic_patterns:
            if pattern in normalized_text:
                self.logger.info(f"过滤泛指概念: '{entity_text}' (类型: {entity_type})")
                return True

        # 对于Ethnic_Group类型，必须是具体的民族名称
        if entity_type == 'Ethnic_Group':
            specific_ethnic_groups = [
                '满族', '赫哲族', '鄂伦春族', '鄂温克族',
                '达斡尔族', '朝鲜族', '蒙古族', '回族',
                '汉族', '锡伯族', '柯尔克孜族'
            ]
            # 如果不在具体民族列表中，可能是泛指概念
            is_specific = any(group in normalized_text for group in specific_ethnic_groups)
            if not is_specific:
                self.logger.info(f"过滤非具体民族: '{entity_text}'")
                return True

        return False

    def normalize_entity(
        self,
        raw_entity: Dict[str, Any],
        similarity_threshold: float = 0.85,
        filter_generic: bool = True
    ) -> str:
        """
        规范化单个实体

        Args:
            raw_entity: 原始实体数据（从LLM返回）
            similarity_threshold: 相似度阈值
            filter_generic: 是否过滤泛指概念

        Returns:
            实体ID
        """
        entity_text = raw_entity.get('text', '')
        entity_type = raw_entity.get('type', '')
        attributes = raw_entity.get('attributes', {})

        if not entity_text or not entity_type:
            self.logger.error(f"实体缺少text或type: {raw_entity}")
            return None

        # 验证实体类型是否在本体中定义
        if entity_type not in self.ontology['entity_types']:
            self.logger.warning(f"未知实体类型 '{entity_type}'，跳过: '{entity_text}'")
            return None

        # 过滤泛指概念
        if filter_generic and self.is_generic_concept(entity_text, entity_type):
            return None

        # 标准化文本
        normalized_text = self.normalize_text(entity_text)

        # 检查是否已存在完全相同的实体
        if normalized_text in self.text_to_id_map:
            existing_id = self.text_to_id_map[normalized_text]
            self.logger.info(f"实体已存在: '{entity_text}' -> {existing_id}")
            return existing_id

        # 查找相似实体
        similar_id = self.find_similar_entity(entity_text, entity_type, similarity_threshold)
        if similar_id:
            # 合并到相似实体
            self.logger.info(f"合并实体: '{entity_text}' -> {similar_id}")
            self.text_to_id_map[normalized_text] = similar_id
            return similar_id

        # 生成新实体ID
        entity_id = self.generate_entity_id(entity_type, entity_text)

        # 获取标准名称
        standard_name = attributes.get('name') or entity_text

        # 保存到注册表
        self.entity_registry[entity_id] = {
            'canonical_name': normalized_text,
            'display_name': standard_name,
            'type': entity_type,
            'attributes': attributes,
            'source_projects': []
        }

        # 保存文本映射
        self.text_to_id_map[normalized_text] = entity_id

        self.logger.info(f"新建实体: {entity_id} - '{standard_name}' ({entity_type})")

        return entity_id

    def normalize_batch(
        self,
        extraction_results: List[Dict[str, Any]],
        similarity_threshold: float = 0.85
    ) -> Dict[str, str]:
        """
        批量规范化实体

        Args:
            extraction_results: LLM抽取结果列表
            similarity_threshold: 相似度阈值

        Returns:
            文本到ID的映射字典（包含所有已处理的实体）
        """
        self.logger.info(f"开始批量规范化，共{len(extraction_results)}个抽取结果")

        for idx, result in enumerate(extraction_results):
            if not result or 'entities' not in result:
                continue

            for entity in result['entities']:
                self.normalize_entity(entity, similarity_threshold)

        self.logger.info(f"批量规范化完成，生成{len(self.entity_registry)}个唯一实体")

        # 返回完整的文本到ID映射（包含所有已处理的实体，不仅仅是本批次）
        return self.text_to_id_map.copy()

    def get_entity_nodes(self) -> List[Dict[str, Any]]:
        """
        获取所有实体节点（用于生成CSV）

        Returns:
            节点列表
        """
        nodes = []

        for entity_id, entity_data in self.entity_registry.items():
            node = {
                'id': entity_id,
                'label': entity_data['display_name'],
                'type': entity_data['type'],
                'properties': json.dumps(entity_data['attributes'], ensure_ascii=False)
            }
            nodes.append(node)

        return nodes

    def get_statistics(self) -> Dict[str, Any]:
        """
        获取统计信息

        Returns:
            统计信息字典
        """
        stats = {
            'total_entities': len(self.entity_registry),
            'entities_by_type': {},
            'type_counters': self.type_counters.copy()
        }

        # 按类型统计
        for entity_data in self.entity_registry.values():
            entity_type = entity_data['type']
            stats['entities_by_type'][entity_type] = stats['entities_by_type'].get(entity_type, 0) + 1

        return stats
