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:
        """
        解决属性冲突：
        1. 状态类：直接覆盖
        2. 开始时间类：取最早时间
        3. 结束时间类：取最晚时间
        4. 其他类：直接覆盖 (不再记录历史)
        """
        resolved = existing_attrs.copy()
        # current_time 在此版本中不再需要，但保留导入以防后续扩展使用
        # 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:
                    # 其他普通属性：直接覆盖，不再生成 _history 字段
                    resolved[key] = new_value

        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("数据库已清空")

# if __name__ == "__main__":
#     # 连接本地 Neo4j（请确保服务已启动）
#     try:
#         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)
#         kg = MilitaryKnowledgeGraph(
#                 uri="bolt://192.168.0.46:57687",
#                 auth=("neo4j", "zdht123@"),
#             )
#         entity_count, rel_count = kg.process_data(data)


#     finally:
#         # kg.clear_database()  # 取消注释以清理测试数据
#         kg.close()