from neo4j import GraphDatabase
import json
import hashlib
import json
from collections import defaultdict
import json
import numpy as np
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_similarity
from collections import defaultdict
import copy
import json
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple, Union
from neo4j import GraphDatabase
# from kgsave import MilitaryKnowledgeGraph
from neo4j import GraphDatabase
import json
import hashlib
import json
from collections import defaultdict
import json
import numpy as np
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_similarity
from collections import defaultdict
import copy
import json
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple, Union
from neo4j import GraphDatabase
import json
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple, Union
from neo4j import GraphDatabase



class MilitaryKnowledgeGraph:
    def __init__(self, uri: str, auth: Tuple[str, str], Ontology: List[str] = None):
        """
        初始化知识图谱
        
        Args:
            uri: Neo4j 数据库连接地址
            auth: 认证信息 (用户名, 密码)
            Ontology: 节点类型列表，用于创建唯一性约束和类型校验
        """
        self.driver = GraphDatabase.driver(uri, auth=auth)
        
        # 默认节点类型
        self.Ontology = Ontology or ["舰艇", "系统", "子系统", "部件", "一级零部件", "二级零部件", "故障模式", "维修方案", "操作项目", "备品备件", "维修手册", "操作使用手册"]
        
        self._setup_constraints()

    def close(self):
        """关闭数据库连接"""
        self.driver.close()

    def _setup_constraints(self):
        """为每个节点类型设置 名称 字段的唯一性约束"""
        with self.driver.session() as session:
            for node_type in self.Ontology:
                try:
                    session.run(f"CREATE CONSTRAINT IF NOT EXISTS FOR (n:{node_type}) REQUIRE n.名称 IS UNIQUE")
                except Exception as e:
                    print(f"创建约束失败 {node_type}: {e}")

    def construct_fact(self, from_name: str, to_name: str, rel_type: str, attributes: Dict) -> str:
        """构造关系的自然语言描述（fact）"""
        attr_str = ", ".join([f"{k}：{v}" for k, v in attributes.items() if v])
        if attr_str:
            return f"{from_name} {rel_type} {to_name}，{attr_str}"
        else:
            return f"{from_name} {rel_type} {to_name}"

    # def _resolve_attribute_conflicts(self, existing_attrs: Dict, new_attrs: Dict) -> Dict:
    #     """解决属性冲突（保留最新状态、最早开始时间、最晚结束时间等）"""
    #     resolved = existing_attrs.copy()
    #     current_time = datetime.now(timezone.utc).isoformat()

    #     for key, new_value in new_attrs.items():
    #         if new_value is None or new_value == "":
    #             continue

    #         if key not in existing_attrs:
    #             resolved[key] = new_value
    #         elif existing_attrs[key] != new_value:
    #             if key in ["当前状态", "执行状态", "完成状态"]:
    #                 resolved[key] = new_value
    #             elif key in ["开始时间", "下达时间", "承担时间"]:
    #                 if new_value < existing_attrs[key]:
    #                     resolved[key] = new_value
    #             elif key in ["结束时间", "执行时限"]:
    #                 if new_value > existing_attrs[key]:
    #                     resolved[key] = new_value
    #             else:
    #                 resolved[key] = new_value
    #                 if f"{key}_history" not in resolved:
    #                     resolved[f"{key}_history"] = []
    #                 resolved[f"{key}_history"].append({
    #                     "previous_value": existing_attrs[key],
    #                     "new_value": new_value,
    #                     "update_time": current_time
    #                 })

    #     return resolved
    def _resolve_attribute_conflicts(self, existing_attrs: Dict, new_attrs: Dict) -> Dict:
        resolved = existing_attrs.copy()
        current_time = datetime.now(timezone.utc).isoformat()

        for key, new_value in new_attrs.items():
            if new_value is None or new_value == "":
                continue

            if key not in existing_attrs:
                resolved[key] = new_value
            elif existing_attrs[key] != new_value:
                if key in ["当前状态", "执行状态", "完成状态"]:
                    resolved[key] = new_value
                elif key in ["开始时间", "下达时间", "承担时间"]:
                    if new_value < existing_attrs[key]:
                        resolved[key] = new_value
                elif key in ["结束时间", "执行时限"]:
                    if new_value > existing_attrs[key]:
                        resolved[key] = new_value
                else:
                    resolved[key] = new_value
                    history_key = f"{key}_history"
                    
                    # ✅ 防御：确保 history 字段始终是列表
                    if history_key not in resolved:
                        resolved[history_key] = []
                    elif isinstance(resolved[history_key], str):
                        # 兼容旧数据：JSON字符串 → list
                        try:
                            parsed = json.loads(resolved[history_key])
                            resolved[history_key] = parsed if isinstance(parsed, list) else [resolved[history_key]]
                        except (json.JSONDecodeError, TypeError):
                            resolved[history_key] = [resolved[history_key]]
                    
                    # ✅ 序列化为 JSON 字符串存储（避免 Neo4j Map 类型报错）
                    history_entry = json.dumps({
                        "previous_value": existing_attrs[key],
                        "new_value": new_value,
                        "update_time": current_time
                    }, ensure_ascii=False)
                    resolved[history_key].append(history_entry)

        return resolved
    def _create_relationship_with_timeline(self, tx, from_name: str, to_name: str,
                                           rel_type: str, attributes: Dict) -> bool:
        """创建或更新带有时间线的关系"""
        current_time = datetime.now(timezone.utc).isoformat()
        clean_attributes = {k: v for k, v in attributes.items() if v is not None and v != ""}

        fact = self.construct_fact(from_name, to_name, rel_type, attributes)
        clean_attributes["fact"] = fact

        existing_rel = tx.run("""
        MATCH (a {名称: $from_name})-[r]->(b {名称: $to_name})
        WHERE type(r) = $rel_type
        RETURN r, properties(r) as props
        """, from_name=from_name, to_name=to_name, rel_type=rel_type).single()

        if existing_rel:
            existing_props = existing_rel["props"]
            updated_props = existing_props.copy()

            for key, new_value in clean_attributes.items():
                if key == "fact":
                    continue
                if key not in existing_props or existing_props[key] != new_value:
                    timeline_key = f"{key}_timeline"
                    if timeline_key not in updated_props:
                        updated_props[timeline_key] = []

                    timeline_entry_dict = {
                        "value": new_value,
                        "timestamp": current_time,
                        "source": "text_extraction"
                    }
                    if key in existing_props:
                        timeline_entry_dict["previous_value"] = existing_props[key]

                    timeline_entry_str = json.dumps(timeline_entry_dict, ensure_ascii=False)
                    updated_props[timeline_key].append(timeline_entry_str)
                    updated_props[key] = new_value

            if "fact" in clean_attributes:
                fact_timeline_key = "fact_timeline"
                if fact_timeline_key not in updated_props:
                    updated_props[fact_timeline_key] = []

                fact_entry_str = json.dumps({
                    "value": clean_attributes["fact"],
                    "timestamp": current_time,
                    "source": "text_extraction"
                }, ensure_ascii=False)
                updated_props[fact_timeline_key].append(fact_entry_str)
                updated_props["fact"] = clean_attributes["fact"]

            updated_props["last_updated"] = current_time

            query = f"""
            MATCH (a {{名称: $from_name}})-[r:{rel_type}]->(b {{名称: $to_name}})
            SET r = $properties
            RETURN r
            """
            tx.run(query, from_name=from_name, to_name=to_name, properties=updated_props)
            print(f"更新关系: {from_name} -[{rel_type}]-> {to_name}")

        else:
            relationship_props = clean_attributes.copy()

            for key, value in clean_attributes.items():
                timeline_key = f"{key}_timeline"
                timeline_entry_dict = {
                    "value": value,
                    "timestamp": current_time,
                    "source": "text_extraction"
                }
                timeline_entry_str = json.dumps(timeline_entry_dict, ensure_ascii=False)
                relationship_props[timeline_key] = [timeline_entry_str]

            relationship_props["created_at"] = current_time
            relationship_props["last_updated"] = current_time
            relationship_props["update_count"] = 1

            query = f"""
            MATCH (a {{名称: $from_name}}), (b {{名称: $to_name}})
            CREATE (a)-[r:{rel_type} $properties]->(b)
            RETURN r
            """
            tx.run(query, from_name=from_name, to_name=to_name, properties=relationship_props)
            print(f"创建关系: {from_name} -[{rel_type}]-> {to_name}")

        return True

    def _create_or_update_node(self, tx, node_type: str, 名称: str, attributes: Dict) -> bool:
        """
        创建或更新节点，支持 knowledge_source 等复杂字段
        
        Args:
            tx: Neo4j 事务对象
            node_type: 节点类型
            名称: 节点名称（唯一标识）
            attributes: 节点属性字典
        """
        if not 名称 or not node_type:
            return False

        # 安全校验：只允许预定义的节点类型
        if node_type not in self.Ontology:
            print(f"跳过未知节点类型: {node_type}")
            return False

        # 预处理 attributes：将 knowledge_source 转为 JSON 字符串
        clean_attributes = {"名称": 名称}  # 确保名称字段存在
        for k, v in attributes.items():
            if v is None or v == "":
                continue
            if k == "knowledge_source" and isinstance(v, list):
                clean_attributes[k] = json.dumps(v, ensure_ascii=False)
            else:
                clean_attributes[k] = v

        existing_node = tx.run(
            f"MATCH (n:{node_type} {{名称: $名称}}) RETURN n, properties(n) as props",
            名称=名称
        ).single()

        current_time = datetime.now(timezone.utc).isoformat()

        if existing_node:
            existing_props = existing_node["props"]
            existing_attrs = {k: v for k, v in existing_props.items()
                             if not k.endswith('_timeline') and k not in ['created_at', 'last_updated']}

            resolved_attributes = self._resolve_attribute_conflicts(existing_attrs, clean_attributes)

            timeline_data = {k: v for k, v in existing_props.items()
                            if k.endswith('_timeline') or k in ['created_at', 'last_updated']}

            final_attributes = {**timeline_data, **resolved_attributes}
            final_attributes["last_updated"] = current_time

            query = f"""
            MATCH (n:{node_type} {{名称: $名称}})
            SET n = $attributes
            RETURN n
            """
            tx.run(query, 名称=名称, attributes=final_attributes)
            print(f"更新节点: {node_type} - {名称}")
        else:
            clean_attributes["created_at"] = current_time
            clean_attributes["last_updated"] = current_time

            query = f"""
            CREATE (n:{node_type} $attributes)
            RETURN n
            """
            tx.run(query, attributes=clean_attributes)
            print(f"创建节点: {node_type} - {名称}")

        return True

    def process_data(self, data: Dict) -> Tuple[int, int]:
        """
        处理字典数据并写入Neo4j
        
        Args:
            data: 包含 entities 和 relationships 的字典
            
        Returns:
            Tuple[int, int]: (实体数量, 关系数量)
        """
        if not data or not isinstance(data, dict):
            print("无效的数据格式，需要字典类型")
            return 0, 0

        entity_count = 0
        relationship_count = 0

        with self.driver.session() as session:
            # 创建/更新实体节点
            for entity in data.get("entities", []):
                try:
                    # 提取 properties，并注入 knowledge_source（如果存在）
                    props = entity.get("properties", {}).copy()
                    if "knowledge_source" in entity:
                        props["knowledge_source"] = entity["knowledge_source"]

                    # 从 properties 中提取名称
                    名称 = props.pop("名称", None)
                    if not 名称:
                        print(f"跳过无名称节点: {entity}")
                        continue

                    success = session.execute_write(
                        self._create_or_update_node,
                        entity["type"],
                        名称,
                        props
                    )
                    if success:
                        entity_count += 1
                except Exception as e:
                    print(f"处理节点失败: {entity}, 错误: {e}")

            # 创建/更新关系
            for relationship in data.get("relationships", []):
                try:
                    success = session.execute_write(
                        self._create_relationship_with_timeline,
                        relationship["from_entity"],
                        relationship["to_entity"],
                        relationship["type"],
                        relationship.get("properties", {})
                    )
                    if success:
                        relationship_count += 1
                except Exception as e:
                    print(f"处理关系失败: {relationship}, 错误: {e}")

            print(f"成功处理 {entity_count} 个节点和 {relationship_count} 个关系")
            return entity_count, relationship_count

    def get_relationship_history(self, from_name: str, to_name: str, rel_type: str) -> Optional[Dict]:
        """获取关系的完整历史记录"""
        with self.driver.session() as session:
            result = session.run("""
            MATCH (a {名称: $from_name})-[r]->(b {名称: $to_name})
            WHERE type(r) = $rel_type
            RETURN properties(r) as props
            """, from_name=from_name, to_name=to_name, rel_type=rel_type)

            record = result.single()
            return record["props"] if record else None

    def get_property_timeline(self, from_name: str, to_name: str, rel_type: str,
                              property_name: str) -> Optional[List]:
        """获取特定属性的时间线（返回字符串列表，需 json.loads 解析）"""
        history = self.get_relationship_history(from_name, to_name, rel_type)
        if history:
            return history.get(f"{property_name}_timeline", [])
        return None

    def clear_database(self):
        """清空数据库（谨慎使用）"""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("数据库已清空")

# def merge_records(record_list):
#     grouped = defaultdict(list)
    
#     for record in record_list:
#         key = (record['type'], record['properties']['名称'])
#         grouped[key].append(record)
    
#     merged_list = []
    
#     for key, group in grouped.items():
#         if len(group) == 1:
#             merged_list.append(group[0])
#         else:
#             base = group[0].copy()
#             props_merged = base['properties'].copy()
#             ks_set = set()

#             for rec in group:
#                 try:
#                     ks_value = rec['properties']['knowledge_source']
#                     # 🔧 修复:处理不同类型的 knowledge_source
#                     if isinstance(ks_value, list):
#                         ks_list = ks_value
#                     elif isinstance(ks_value, str):
#                         ks_list = json.loads(ks_value)
#                     else:
#                         ks_list = []
                        
#                     for item in ks_list:
#                         ks_set.add(json.dumps(item, sort_keys=True, ensure_ascii=False))
#                 except (json.JSONDecodeError, KeyError, TypeError):
#                     pass

#                 for k, v in rec['properties'].items():
#                     if k != 'knowledge_source':
#                         if props_merged.get(k) in ("", "香", None) and v not in ("", "香", None):
#                             props_merged[k] = v
#                         elif k not in props_merged:
#                             props_merged[k] = v

#             merged_ks = [json.loads(s) for s in ks_set]
#             props_merged['knowledge_source'] = json.dumps(merged_ks, ensure_ascii=False)

#             base['properties'] = props_merged
#             merged_list.append(base)
    
#     return merged_list
def merge_records(record_list):
    grouped = defaultdict(list)
    
    for record in record_list:
        # 1. 安全获取 type 和 名称，确保它们是字符串（可哈希）
        node_type = record.get('type', 'Unknown')
        if isinstance(node_type, list):
            node_type = node_type[0] if node_type else 'Unknown'
        node_type = str(node_type) # 强制转为字符串
        
        # 获取名称字段
        raw_name = record.get('properties', {}).get('名称')
        
        # 2. 处理名称可能是列表的情况 (修复 unhashable type: 'list')
        if isinstance(raw_name, list):
            if len(raw_name) > 0:
                node_name = raw_name[0]
            else:
                node_name = "Unnamed_Node"
        elif raw_name is None or raw_name == "":
            node_name = "Unnamed_Node"
        else:
            node_name = raw_name
            
        # 确保最终是字符串
        node_name = str(node_name)

        # 3. 现在 key 肯定是可哈希的 (str, str)
        key = (node_type, node_name)
        
        grouped[key].append(record)
    
    merged_list = []
    
    for key, group in grouped.items():
        if len(group) == 1:
            merged_list.append(group[0])
        else:
            base = group[0].copy()
            # 深度复制 properties 以防修改影响原数据
            props_merged = copy.deepcopy(base['properties'])
            
            # 修正：确保合并后的名称是字符串，而不是列表
            if isinstance(props_merged.get('名称'), list):
                props_merged['名称'] = props_merged['名称'][0] if props_merged['名称'] else key[1]

            ks_set = set()

            for rec in group:
                try:
                    ks_value = rec['properties'].get('knowledge_source')
                    # 处理不同类型的 knowledge_source
                    if isinstance(ks_value, list):
                        ks_list = ks_value
                    elif isinstance(ks_value, str):
                        try:
                            parsed = json.loads(ks_value)
                            ks_list = parsed if isinstance(parsed, list) else [parsed]
                        except (json.JSONDecodeError, TypeError):
                            ks_list = [ks_value] if ks_value else []
                    else:
                        ks_list = []
                        
                    for item in ks_list:
                        if isinstance(item, dict):
                            # 使用 sort_keys 确保字典顺序不影响去重
                            ks_set.add(json.dumps(item, sort_keys=True, ensure_ascii=False))
                        else:
                            ks_set.add(str(item))
                except Exception as e:
                    print(f"解析 knowledge_source 失败: {e}")
                    pass

                # 合并其他属性
                for k, v in rec['properties'].items():
                    if k == 'knowledge_source':
                        continue
                    
                    # 跳过名称字段的重复合并，保持上面处理过的字符串
                    if k == '名称':
                        continue

                    current_val = props_merged.get(k)
                    
                    # 如果当前值是空或无效值，直接替换
                    if current_val in ("", "香", None, []):
                        props_merged[k] = v
                    # 如果新值有效且当前值不同，且当前没值，则填入
                    elif k not in props_merged:
                        props_merged[k] = v
                    # 如果双方都有值且不同，可以考虑保留非空值或转为列表（视需求而定，这里保持简单覆盖或忽略）
                    elif current_val != v and v not in ("", "香", None):
                        # 策略：如果现有值是“香”或空，替换；否则保留现有值（避免冲突）
                        if current_val in ("", "香"):
                            props_merged[k] = v
            
            # 处理 knowledge_source 的最终合并
            if ks_set:
                merged_ks = []
                for s in ks_set:
                    try:
                        # 尝试还原为对象
                        obj = json.loads(s)
                        merged_ks.append(obj)
                    except:
                        merged_ks.append(s)
                
                # 存回为 JSON 字符串
                props_merged['knowledge_source'] = json.dumps(merged_ks, ensure_ascii=False)
            else:
                # 如果没有提取到有效的 KS，保留原有的或设为空列表字符串
                if 'knowledge_source' not in props_merged:
                    props_merged['knowledge_source'] = "[]"

            base['properties'] = props_merged
            merged_list.append(base)
    
    return merged_list


def fetch_all_nodes_and_rels(tx):
    query = """
    MATCH (n)
    OPTIONAL MATCH (n)-[r]->()
    WITH 
        collect(DISTINCT {
            id: toString(id(n)),
            labels: labels(n),
            properties: properties(n)
        }) AS nodes,
        collect(DISTINCT r) AS rels
    UNWIND rels AS r
    WITH nodes, r
    WHERE r IS NOT NULL
    RETURN 
        nodes,
        collect({
            id: toString(id(r)),
            source: toString(id(startNode(r))),
            target: toString(id(endNode(r))),
            type: type(r),
            properties: properties(r)
        }) AS relationships
    """
    result = tx.run(query)
    record = result.single()
    if record:
        return record["nodes"], record["relationships"]
    else:
        return [], []


#从已有数据库导出已存在的关系
def reconstruct_relationships(entities, relationships):
    # 构建 id -> entity 的映射字典
    entity_map = {entity["id"]: entity for entity in entities}

    reconstructed_rela = []
    for rel in relationships:
        source_id = rel.get("source")
        target_id = rel.get("target")
        rel_type = rel.get("type")

        # 检查 source 和 target 是否存在
        source_entity = entity_map.get(source_id)
        target_entity = entity_map.get(target_id)

        if not source_entity or not target_entity:
            continue  # 跳过任一实体缺失的关系

        # 安全获取“名称”字段（注意：根据你的数据，“名称”在顶层，不在 properties 中！）
        from_name = source_entity["properties"].get("名称")
        to_name = target_entity["properties"].get("名称")

        # 如果任一名称缺失，也跳过
        if from_name is None or to_name is None:
            continue

        # 构建新关系
        reconstructed_rela.append({
            "type": rel_type,
            "from_entity": from_name,
            "to_entity": to_name
        })
    
    return reconstructed_rela

def reconstruct_entities(entities):
    # 要排除的字段（无论在顶层还是 properties 内部）
    excluded_keys = {"last_updated", "fulltext", "created_at"}
    
    reconstructed_rela = []
    for entity in entities:
        labels = entity.get("labels")
        if not isinstance(labels, list) or len(labels) == 0:
            continue
        
        entity_type = labels[0]
        
        # 合并顶层字段 + properties 内部字段
        combined_props = {}
        
        # 添加顶层字段（除了 id, labels, 和可能存在的 properties）
        for k, v in entity.items():
            if k in ("id", "labels", "properties"):
                continue
            combined_props[k] = v
        
        # 添加 properties 内部字段（如果存在）
        inner_props = entity.get("properties", {})
        if isinstance(inner_props, dict):
            for k, v in inner_props.items():
                combined_props[k] = v
        
        # 剔除不需要的字段
        filtered_props = {
            k: v for k, v in combined_props.items()
            if k not in excluded_keys
        }
        
        # 只有当 properties 非空时才保留（可选）
        if not filtered_props:
            continue
        
        reconstructed_rela.append({
            "type": entity_type,
            "properties": filtered_props
        })
    
    return reconstructed_rela
def deduplicate_relations(relations):
    seen = set()
    unique_relations = []
    for rel in relations:
        # 构造唯一键（必须是可哈希的）
        key = (rel["type"], rel["from_entity"], rel["to_entity"])
        if key not in seen:
            seen.add(key)
            unique_relations.append(rel)
    return unique_relations

def parse_knowledge_source(ks) -> list:
    if ks is None or ks == "" or ks == []:
        return []
    if isinstance(ks, list):
        result = []
        for item in ks:
            if isinstance(item, dict):
                result.append(item)
            elif isinstance(item, str):
                try:
                    parsed = json.loads(item)
                    if isinstance(parsed, list):
                        result.extend(parsed)
                    elif isinstance(parsed, dict):
                        result.append(parsed)
                except (json.JSONDecodeError, TypeError):
                    pass
        return result
    if isinstance(ks, str):
        try:
            parsed = json.loads(ks)
            if isinstance(parsed, list):
                return parsed
            if isinstance(parsed, dict):
                return [parsed]
        except (json.JSONDecodeError, TypeError):
            pass
    return []


def merge_knowledge_source(master_ks, slave_ks) -> list:
    master_list = parse_knowledge_source(master_ks)
    slave_list  = parse_knowledge_source(slave_ks)
    seen = {(item.get("filename", ""), item.get("url", "")) for item in master_list}
    for item in slave_list:
        key = (item.get("filename", ""), item.get("url", ""))
        if key not in seen:
            master_list.append(item)
            seen.add(key)
    return master_list


# ============================================================
# Step 1: 加载数据
# ============================================================
def load_data(path: str):
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    entities      = data["entities"]
    relationships = data.get("relationships", [])
    print(f"[Step 1] 共加载 {len(entities)} 个实体，"
          f"Embedding 维度: {np.array(entities[0]['properties']['embedding']).shape[0]}，"
          f"{len(relationships)} 条关系")
    return data, entities, relationships


# ============================================================
# Step 2: DBSCAN 聚类（仅用于目标类型）
# ============================================================
def cluster_entities(entities: list, eps: float, min_samples: int) -> np.ndarray:
    embeddings = np.array([e["properties"]["embedding"] for e in entities], dtype=np.float32)
    embeddings_norm = normalize(embeddings, norm="l2")
    labels = DBSCAN(
        eps=eps,
        min_samples=min_samples,
        metric="cosine",
        algorithm="brute",
        n_jobs=-1
    ).fit_predict(embeddings_norm)
    n_clusters = len(set(labels)) - (1 if -1 in labels else 0)
    print(f"[Step 2] DBSCAN 聚类完成，共发现 {n_clusters} 个簇（eps={eps}）")
    return labels


# ============================================================
# Step 3: 属性合并（单对主从）
# ============================================================
PROTECTED_KEYS = {"名称", "型号", "组成编码"}

def merge_properties(master: dict, slave: dict) -> list:
    changes = []

    for key, val in slave["properties"].items():
        if key == "embedding":
            continue

        if key == "knowledge_source":
            master_ks = master["properties"].get("knowledge_source")
            merged_ks = merge_knowledge_source(master_ks, val)
            master["properties"]["knowledge_source"] = json.dumps(merged_ks, ensure_ascii=False)
            changes.append(f"knowledge_source extend {len(merged_ks)} items")
            continue

        if key in PROTECTED_KEYS:
            continue

        existing = master["properties"].get(key)
        if existing in (None, "", []):
            master["properties"][key] = val
            changes.append(f"{key}={val}")
        elif val not in (None, "", []) and existing != val:
            merged = existing if isinstance(existing, list) else [existing]
            for item in (val if isinstance(val, list) else [val]):
                if item not in merged:
                    merged.append(item)
            master["properties"][key] = merged
            changes.append(f"{key} merged={merged}")

    return changes


# ============================================================
# 辅助：从一组实体中选主实体
#   优先有"组成编码"，其次属性最丰富
# ============================================================
def pick_master(indices: list, entities: list) -> int:
    master_idx = next(
        (idx for idx in indices if entities[idx]["properties"].get("组成编码")),
        None
    )
    if master_idx is None:
        master_idx = max(
            indices,
            key=lambda idx: sum(
                1 for k, v in entities[idx]["properties"].items()
                if k != "embedding" and v not in (None, "", [])
            )
        )
    return master_idx


# ============================================================
# Step 4: 目标类型实体合并
#   规则：余弦相似度 >= NAME_SIM_THRESH 即合并（不要求名称一致）
# ============================================================
def merge_target_entities(
    entities: list,
    labels: np.ndarray,
    target_types: set,
    sim_thresh: float,
) -> tuple:

    clusters = defaultdict(list)
    for idx, label in enumerate(labels):
        clusters[label].append(idx)

    to_delete:  set  = set()
    name_remap: dict = {}
    merge_log:  list = []

    for label, indices in clusters.items():
        if label == -1 or len(indices) <= 1:
            continue

        target_indices = [idx for idx in indices if entities[idx]["type"] in target_types]
        if len(target_indices) <= 1:
            continue

        # Union-Find：相似度 >= sim_thresh 则合并
        parent = {idx: idx for idx in target_indices}

        def find(x):
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        def union(x, y):
            parent[find(x)] = find(y)

        for i in range(len(target_indices)):
            for j in range(i + 1, len(target_indices)):
                ei = entities[target_indices[i]]
                ej = entities[target_indices[j]]
                emb_i = np.array(ei["properties"]["embedding"], dtype=np.float32).reshape(1, -1)
                emb_j = np.array(ej["properties"]["embedding"], dtype=np.float32).reshape(1, -1)
                sim   = cosine_similarity(emb_i, emb_j)[0][0]
                if sim >= sim_thresh:
                    union(target_indices[i], target_indices[j])

        groups = defaultdict(list)
        for idx in target_indices:
            groups[find(idx)].append(idx)

        for root, group_indices in groups.items():
            if len(group_indices) <= 1:
                continue

            master_idx  = pick_master(group_indices, entities)
            master      = entities[master_idx]
            master_name = master["properties"]["名称"]
            if isinstance(master_name, list):
                master_name = master_name[0] if master_name else ""
                master["properties"]["名称"] = master_name

            slave_indices = [idx for idx in group_indices if idx != master_idx]
            log_entry = {
                "阶段": "目标类型合并",
                "簇ID": label, "名称": master_name,
                "主实体类型": master["type"],
                "被合并实体": [], "新增属性": []
            }

            for s_idx in slave_indices:
                slave      = entities[s_idx]
                slave_name = slave["properties"].get("名称", "")
                if isinstance(slave_name, list):
                    slave_name = slave_name[0] if slave_name else ""

                log_entry["被合并实体"].append({
                    "类型":  slave["type"],
                    "原属性": copy.deepcopy(slave["properties"])
                })

                if slave_name and slave_name != master_name:
                    name_remap[slave_name] = master_name

                changes = merge_properties(master, slave)
                log_entry["新增属性"].extend(changes)
                to_delete.add(s_idx)

            merge_log.append(log_entry)

    print(f"[Step 4] 目标类型合并：{len(merge_log)} 组合并，"
          f"删除 {len(to_delete)} 个从实体，名称映射 {len(name_remap)} 条")
    return to_delete, name_remap, merge_log


# ============================================================
# Step 5: 非目标类型实体合并
#   规则：type 与 name（名称）字段值完全一致才合并
#   全局遍历，不依赖 DBSCAN 簇
# ============================================================
def merge_other_entities(
    entities: list,
    target_types: set,
    existing_to_delete: set,
) -> tuple:

    # 按 (type, name) 分组，跳过目标类型和已被删除的实体
    groups: dict = defaultdict(list)
    for idx, e in enumerate(entities):
        if idx in existing_to_delete:
            continue
        if e["type"] in target_types:
            continue
        name = e["properties"].get("名称", "")
        if isinstance(name, list):
            name = name[0] if name else ""
        key = (e["type"], name)
        if name:  # 名称为空时不参与合并
            groups[key].append(idx)

    to_delete:  set  = set()
    name_remap: dict = {}
    merge_log:  list = []

    for (etype, ename), indices in groups.items():
        if len(indices) <= 1:
            continue

        master_idx    = pick_master(indices, entities)
        master        = entities[master_idx]
        slave_indices = [idx for idx in indices if idx != master_idx]

        log_entry = {
            "阶段": "非目标类型合并",
            "类型": etype, "名称": ename,
            "被合并实体": [], "新增属性": []
        }

        for s_idx in slave_indices:
            slave      = entities[s_idx]
            slave_name = slave["properties"].get("名称", "")
            if isinstance(slave_name, list):
                slave_name = slave_name[0] if slave_name else ""

            log_entry["被合并实体"].append({
                "类型":  slave["type"],
                "原属性": copy.deepcopy(slave["properties"])
            })

            # 名称相同时无需添加映射，但保留以防万一
            if slave_name and slave_name != ename:
                name_remap[slave_name] = ename

            changes = merge_properties(master, slave)
            log_entry["新增属性"].extend(changes)
            to_delete.add(s_idx)

        merge_log.append(log_entry)

    print(f"[Step 5] 非目标类型合并：{len(merge_log)} 组合并，"
          f"删除 {len(to_delete)} 个从实体，名称映射 {len(name_remap)} 条")
    return to_delete, name_remap, merge_log


# ============================================================
# Step 6: 清理所有实体的"名称"字段（防止历史脏数据）
# ============================================================
def clean_entity_names(entities: list) -> list:
    for e in entities:
        name = e["properties"].get("名称")
        if isinstance(name, list):
            e["properties"]["名称"] = name[0] if name else ""
    return entities


# ============================================================
# Step 7: 更新 relationships 并去重
# ============================================================
def update_relationships(relationships: list, name_remap: dict) -> tuple:
    updated_count = 0
    for rel in relationships:
        orig_from, orig_to = rel["from_entity"], rel["to_entity"]
        rel["from_entity"] = name_remap.get(rel["from_entity"], rel["from_entity"])
        rel["to_entity"]   = name_remap.get(rel["to_entity"],   rel["to_entity"])
        if rel["from_entity"] != orig_from or rel["to_entity"] != orig_to:
            updated_count += 1

    seen, deduped = set(), []
    for rel in relationships:
        key = (rel["type"], rel["from_entity"], rel["to_entity"])
        if key not in seen:
            seen.add(key)
            deduped.append(rel)

    removed = len(relationships) - len(deduped)
    print(f"[Step 7] 关系更新完成：更新 {updated_count} 条，"
          f"去重删除 {removed} 条，剩余 {len(deduped)} 条")
    return deduped, updated_count, removed


def filter_node_relationship(URI,USERNAME,PASSWORD,merged_output_data):
    DBSCAN_EPS      = 0.1
    DBSCAN_MIN_SPL  = 1
    TARGET_TYPES    = {"系统", "子系统", "部件", "一级零部件", "二级零部件"}
    NAME_SIM_THRESH = 0.9   # 目标类型：embedding 余弦相似度阈值
    OUTPUT_FILE1 = "./huizong_node_rela.json"

    OUTPUT_FILE = "./neo4j_export_deduplicated.json"

    driver = GraphDatabase.driver(URI, auth=(USERNAME, PASSWORD))

    try:
        with driver.session() as session:
            print("正在从 Neo4j 获取所有节点和关系...")
            nodes, relationships = session.execute_read(fetch_all_nodes_and_rels)

            print(f"原始数据: {len(nodes)} 个实体, {len(relationships)} 条关系")
            final_entity = []
            final_relation = []
            all_entity_old = reconstruct_entities(nodes)
            all_relationships_old = reconstruct_relationships(nodes, relationships)
            all_entity_new = merged_output_data["entities"]
            all_relationships_new = merged_output_data["relationships"]
            final_entity.extend(all_entity_old)
            final_entity.extend(all_entity_new)
            final_relation.extend(all_relationships_old)
            final_relation.extend(all_relationships_new)
            print(f"汇总数据: {len(final_entity)} 个实体, {len(final_relation)} 条关系")
            
            filter_data1 = {
                "entities": final_entity,
                "relationships": final_relation
            }
            with open(OUTPUT_FILE1, 'w', encoding='utf-8') as f:
                json.dump(filter_data1, f, indent=2, ensure_ascii=False)
            filter_node = merge_records(final_entity)
            filter_relationships = deduplicate_relations(final_relation)
            print(f"过滤后数据: {len(filter_node)} 个实体, {len(filter_relationships)} 条关系")

            # print("进行DBSCAN聚类")
            # entities = clean_entity_names(filter_node)

            # # Step 3: DBSCAN 聚类（全量，用于目标类型相似度筛选）
            # labels = cluster_entities(entities, eps=DBSCAN_EPS, min_samples=DBSCAN_MIN_SPL)

            # # Step 4: 目标类型合并（相似度 >= NAME_SIM_THRESH）
            # to_delete_target, name_remap_target, merge_log_target = merge_target_entities(
            #     entities, labels, TARGET_TYPES, NAME_SIM_THRESH
            # )

            # # Step 5: 非目标类型合并（type + name 完全一致）
            # to_delete_other, name_remap_other, merge_log_other = merge_other_entities(
            #     entities, TARGET_TYPES, existing_to_delete=to_delete_target
            # )

            # 汇总
            # all_to_delete = to_delete_target | to_delete_other
            # all_name_remap = {**name_remap_target, **name_remap_other}   # 目标类型映射优先级更高时可调整顺序

            # new_entities = [e for idx, e in enumerate(entities) if idx not in all_to_delete]
            # print(f"         总计删除 {len(all_to_delete)} 个从实体，剩余实体数量: {len(new_entities)}")

            # # Step 6: 更新关系
            # new_relationships, _, _ = update_relationships(filter_relationships, all_name_remap)

            filter_data = {
                "entities": filter_node,
                "relationships": filter_relationships
            }
            with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
                json.dump(filter_data, f, indent=2, ensure_ascii=False)

            return filter_data
    except Exception as e:
        print(f"? 导出失败: {e}")
        
        import traceback
        traceback.print_exc()
        return merged_output_data
    finally:
        driver.close()


# if __name__ == "__main__":
#     URI = "bolt://192.168.0.46:57687"
#     USERNAME = "neo4j"
#     PASSWORD = "zdht123@"
#     with open("/storage01/home/hdf/project/wxkgrag/kgextract/output/output3/neo4j_export_merged1.json", "r", encoding="utf-8") as f:
#         data = json.load(f)
#     filter_data = filter_node_relationship(URI=URI,USERNAME=USERNAME,PASSWORD=PASSWORD,merged_output_data=data)
#     kg = MilitaryKnowledgeGraph(
#             uri="bolt://192.168.0.46:57687",
#             auth=("neo4j", "zdht123@"),
#         )
#     entity_count, rel_count = kg.process_data(filter_data)
