import asyncio
import os
from typing import List, Optional

# 使用新的仓储层
from app.repositories.neo4j.neo4j_connector import Neo4jConnector
from app.repositories.neo4j.add_nodes import add_dialogue_nodes, add_statement_nodes, add_chunk_nodes
from app.repositories.neo4j.cypher_queries import (
    STATEMENT_ENTITY_EDGE_SAVE,
    ENTITY_RELATIONSHIP_SAVE,
    EXTRACTED_ENTITY_NODE_SAVE,
    CHUNK_STATEMENT_EDGE_SAVE,
    STATEMENT_ENTITY_EDGE_SAVE,
    ENTITY_RELATIONSHIP_SAVE,
    EXTRACTED_ENTITY_NODE_SAVE,
)
from app.core.memory.models.graph_models import (
    DialogueNode,
    ChunkNode,
    StatementChunkEdge,
    StatementEntityEdge,
    StatementNode,
    ExtractedEntityNode,
    EntityEntityEdge,
    PerceptualNode,
    PerceptualEdge,
)
import logging

logger = logging.getLogger(__name__)


async def save_entities_and_relationships(
        entity_nodes: List[ExtractedEntityNode],
        entity_entity_edges: List[EntityEntityEdge],
        connector: Neo4jConnector
):
    """Save entities and their relationships using graph models"""
    all_entities = [entity.model_dump() for entity in entity_nodes]
    all_relationships = []

    for edge in entity_entity_edges:
        relationship = {
            'source_id': edge.source,
            'target_id': edge.target,
            'predicate': edge.relation_type,
            'statement_id': edge.source_statement_id,
            'value': edge.relation_value,
            'statement': edge.statement,
            'valid_at': edge.valid_at.isoformat() if edge.valid_at else None,
            'invalid_at': edge.invalid_at.isoformat() if edge.invalid_at else None,
            'created_at': edge.created_at.isoformat() if edge.created_at else None,
            'expired_at': edge.expired_at.isoformat() if edge.expired_at else None,
            'run_id': edge.run_id,
            'end_user_id': edge.end_user_id,
        }
        all_relationships.append(relationship)

    # Save entities
    if all_entities:
        entity_uuids = await connector.execute_query(EXTRACTED_ENTITY_NODE_SAVE, entities=all_entities)
        if entity_uuids:
            print(f"Successfully saved {len(entity_uuids)} entity nodes to Neo4j")
        else:
            print("Failed to save entity nodes to Neo4j")
    else:
        print("No entity nodes to save")

    # Create relationships
    if all_relationships:
        relationship_uuids = await connector.execute_query(ENTITY_RELATIONSHIP_SAVE, relationships=all_relationships)
        if relationship_uuids:
            print(f"Successfully saved {len(relationship_uuids)} entity relationships (edges) to Neo4j")
        else:
            print("Failed to save entity relationships to Neo4j")
    else:
        print("No entity relationships to save")


async def save_chunk_nodes(
        chunk_nodes: List[ChunkNode],
        connector: Neo4jConnector
):
    """Save chunk nodes using graph models"""
    if not chunk_nodes:
        print("No chunk nodes to save")
        return

    chunk_uuids = await add_chunk_nodes(chunk_nodes, connector)
    if chunk_uuids:
        print(f"Successfully saved {len(chunk_uuids)} chunk nodes to Neo4j")
    else:
        print("Failed to save chunk nodes to Neo4j")


async def save_statement_chunk_edges(
        statement_chunk_edges: List[StatementChunkEdge],
        connector: Neo4jConnector
):
    """Save statement-chunk edges using graph models"""
    if not statement_chunk_edges:
        return

    all_sc_edges = []
    for edge in statement_chunk_edges:
        all_sc_edges.append({
            "id": edge.id,
            "source": edge.source,
            "target": edge.target,
            "end_user_id": edge.end_user_id,
            "run_id": edge.run_id,
            "created_at": edge.created_at.isoformat() if edge.created_at else None,
            "expired_at": edge.expired_at.isoformat() if edge.expired_at else None,
        })

    try:
        await connector.execute_query(
            CHUNK_STATEMENT_EDGE_SAVE,
            chunk_statement_edges=all_sc_edges
        )
    except Exception:
        pass


async def save_statement_entity_edges(
        statement_entity_edges: List[StatementEntityEdge],
        connector: Neo4jConnector
):
    """Save statement-entity edges using graph models"""
    if not statement_entity_edges:
        print("No statement-entity edges to save")
        return

    all_se_edges = []
    for edge in statement_entity_edges:
        edge_data = {
            "source": edge.source,
            "target": edge.target,
            "end_user_id": edge.end_user_id,
            "run_id": edge.run_id,
            "connect_strength": edge.connect_strength,
            "created_at": edge.created_at.isoformat() if edge.created_at else None,
            "expired_at": edge.expired_at.isoformat() if edge.expired_at else None,
        }
        all_se_edges.append(edge_data)

    if all_se_edges:
        try:
            await connector.execute_query(
                STATEMENT_ENTITY_EDGE_SAVE,
                relationships=all_se_edges
            )
        except Exception:
            pass


async def save_dialog_and_statements_to_neo4j(
        dialogue_nodes: List[DialogueNode],
        chunk_nodes: List[ChunkNode],
        statement_nodes: List[StatementNode],
        entity_nodes: List[ExtractedEntityNode],
        perceptual_nodes: List[PerceptualNode],
        entity_edges: List[EntityEntityEdge],
        statement_chunk_edges: List[StatementChunkEdge],
        statement_entity_edges: List[StatementEntityEdge],
        perceptual_edges: List[PerceptualEdge],
        connector: Neo4jConnector,
) -> bool:
    """Save dialogue nodes, chunk nodes, statement nodes, entities, and all relationships to Neo4j using graph models.

    只负责数据写入，不触发聚类。聚类由调用方在写入成功后通过
    _trigger_clustering_sync() 显式触发。

    Args:
        dialogue_nodes: List of DialogueNode objects to save
        chunk_nodes: List of ChunkNode objects to save
        statement_nodes: List of StatementNode objects to save
        entity_nodes: List of ExtractedEntityNode objects to save
        perceptual_nodes: List of PerceptualNode objects to save
        entity_edges: List of EntityEntityEdge objects to save
        statement_chunk_edges: List of StatementChunkEdge objects to save
        statement_entity_edges: List of StatementEntityEdge objects to save
        perceptual_edges: List of PerceptualEdge objects to save
        connector: Neo4j connector instance

    Returns:
        bool: True if successful, False otherwise
    """
    # TODO 需要在去重消歧节阶段，做以下逻辑的处理
    # 预处理：对特殊实体（"用户"、"AI助手"）复用 Neo4j 中已有节点的 ID，
    # 确保同一个 end_user_id 下只有一个"用户"节点和一个"AI助手"节点。
    if entity_nodes:
        _SPECIAL_NAMES = {"用户", "我", "user", "i", "ai助手", "助手", "ai assistant", "assistant"}
        end_user_id = entity_nodes[0].end_user_id if entity_nodes else None
        if end_user_id:
            try:
                # 查询已有的特殊实体
                cypher = """
                MATCH (e:ExtractedEntity)
                WHERE e.end_user_id = $end_user_id AND toLower(e.name) IN $names
                RETURN e.id AS id, e.name AS name
                """
                existing = await connector.execute_query(
                    cypher,
                    end_user_id=end_user_id,
                    names=list(_SPECIAL_NAMES),
                )
                # 建立 name(lower) → existing_id 映射
                existing_id_map = {}
                for record in (existing or []):
                    name_lower = (record.get("name") or "").strip().lower()
                    if name_lower and record.get("id"):
                        existing_id_map[name_lower] = record["id"]

                if existing_id_map:
                    # 替换新实体的 ID 为已有 ID，同时更新所有引用该 ID 的边
                    for ent in entity_nodes:
                        name_lower = (ent.name or "").strip().lower()
                        if name_lower in existing_id_map:
                            old_id = ent.id
                            new_id = existing_id_map[name_lower]
                            if old_id != new_id:
                                ent.id = new_id
                                # 更新 statement_entity_edges 中的引用
                                for edge in statement_entity_edges:
                                    if edge.target == old_id:
                                        edge.target = new_id
                                    if edge.source == old_id:
                                        edge.source = new_id
                                # 更新 entity_edges 中的引用
                                for edge in entity_edges:
                                    if edge.source == old_id:
                                        edge.source = new_id
                                    if edge.target == old_id:
                                        edge.target = new_id
                                logger.info(
                                    f"特殊实体 '{ent.name}' ID 复用: {old_id[:8]}... → {new_id[:8]}..."
                                )
            except Exception as e:
                logger.warning(f"特殊实体 ID 复用查询失败（不影响写入）: {e}")

    # 定义事务函数，将所有写操作放在一个事务中
    async def _save_all_in_transaction(tx):
        """在单个事务中执行所有保存操作，避免死锁"""
        results = {}

        # 1. Save all dialogue nodes in batch
        if dialogue_nodes:
            from app.repositories.neo4j.cypher_queries import DIALOGUE_NODE_SAVE
            dialogue_data = [node.model_dump() for node in dialogue_nodes]
            result = await tx.run(DIALOGUE_NODE_SAVE, dialogues=dialogue_data)
            dialogue_uuids = [record["uuid"] async for record in result]
            results['dialogues'] = dialogue_uuids
            logger.info(f"Dialogues saved to Neo4j with UUIDs: {dialogue_uuids}")

        # 2. Save all chunk nodes in batch
        if chunk_nodes:
            from app.repositories.neo4j.cypher_queries import CHUNK_NODE_SAVE
            chunk_data = [node.model_dump() for node in chunk_nodes]
            result = await tx.run(CHUNK_NODE_SAVE, chunks=chunk_data)
            chunk_uuids = [record["uuid"] async for record in result]
            results['chunks'] = chunk_uuids
            logger.info(f"Successfully saved {len(chunk_uuids)} chunk nodes to Neo4j")

        if perceptual_nodes:
            from app.repositories.neo4j.cypher_queries import PERCEPTUAL_NODE_SAVE
            perceptual_data = [node.model_dump() for node in perceptual_nodes]
            result = await tx.run(PERCEPTUAL_NODE_SAVE, perceptuals=perceptual_data)
            perceptual_uuids = [record["uuid"] async for record in result]
            results["perceptuals"] = perceptual_uuids
            logger.info(f"Successfully saved {len(perceptual_uuids)} perceptual nodes to Neo4j")

        # 3. Save all statement nodes in batch
        if statement_nodes:
            from app.repositories.neo4j.cypher_queries import STATEMENT_NODE_SAVE
            statement_data = [node.model_dump() for node in statement_nodes]
            result = await tx.run(STATEMENT_NODE_SAVE, statements=statement_data)
            statement_uuids = [record["uuid"] async for record in result]
            results['statements'] = statement_uuids
            logger.info(f"Successfully saved {len(statement_uuids)} statement nodes to Neo4j")

        # 4. Save entities
        if entity_nodes:
            from app.repositories.neo4j.cypher_queries import EXTRACTED_ENTITY_NODE_SAVE
            entity_data = [entity.model_dump() for entity in entity_nodes]
            result = await tx.run(EXTRACTED_ENTITY_NODE_SAVE, entities=entity_data)
            entity_uuids = [record["uuid"] async for record in result]
            results['entities'] = entity_uuids
            logger.info(f"Successfully saved {len(entity_uuids)} entity nodes to Neo4j")

        # 5. Create entity relationships
        if entity_edges:
            from app.repositories.neo4j.cypher_queries import ENTITY_RELATIONSHIP_SAVE
            relationship_data = []
            for edge in entity_edges:
                relationship_data.append({
                    'source_id': edge.source,
                    'target_id': edge.target,
                    'predicate': edge.relation_type,
                    'statement_id': edge.source_statement_id,
                    'value': edge.relation_value,
                    'statement': edge.statement,
                    'valid_at': edge.valid_at.isoformat() if edge.valid_at else None,
                    'invalid_at': edge.invalid_at.isoformat() if edge.invalid_at else None,
                    'created_at': edge.created_at.isoformat() if edge.created_at else None,
                    'expired_at': edge.expired_at.isoformat() if edge.expired_at else None,
                    'run_id': edge.run_id,
                    'end_user_id': edge.end_user_id,
                })
            result = await tx.run(ENTITY_RELATIONSHIP_SAVE, relationships=relationship_data)
            rel_uuids = [record["uuid"] async for record in result]
            results['entity_relationships'] = rel_uuids
            logger.info(f"Successfully saved {len(rel_uuids)} entity relationships to Neo4j")

        # 6. Save statement-chunk edges
        if statement_chunk_edges:
            from app.repositories.neo4j.cypher_queries import CHUNK_STATEMENT_EDGE_SAVE
            sc_edge_data = []
            for edge in statement_chunk_edges:
                sc_edge_data.append({
                    "id": edge.id,
                    "source": edge.source,
                    "target": edge.target,
                    "created_at": edge.created_at.isoformat() if edge.created_at else None,
                    "expired_at": edge.expired_at.isoformat() if edge.expired_at else None,
                    "run_id": edge.run_id,
                    "end_user_id": edge.end_user_id,
                })
            result = await tx.run(CHUNK_STATEMENT_EDGE_SAVE, chunk_statement_edges=sc_edge_data)
            sc_uuids = [record["uuid"] async for record in result]
            results['statement_chunk_edges'] = sc_uuids
            logger.info(f"Successfully saved {len(sc_uuids)} statement-chunk edges to Neo4j")

        # 7. Save statement-entity edges
        if statement_entity_edges:
            from app.repositories.neo4j.cypher_queries import STATEMENT_ENTITY_EDGE_SAVE
            se_edge_data = []
            for edge in statement_entity_edges:
                se_edge_data.append({
                    "source": edge.source,
                    "target": edge.target,
                    "created_at": edge.created_at.isoformat() if edge.created_at else None,
                    "expired_at": edge.expired_at.isoformat() if edge.expired_at else None,
                    "run_id": edge.run_id,
                    "end_user_id": edge.end_user_id,
                    "connect_strength": getattr(edge, "connect_strength", "strong"),
                })
            result = await tx.run(STATEMENT_ENTITY_EDGE_SAVE, relationships=se_edge_data)
            se_uuids = [record["uuid"] async for record in result]
            results['statement_entity_edges'] = se_uuids
            logger.info(f"Successfully saved {len(se_uuids)} statement-entity edges to Neo4j")

        if perceptual_edges:
            from app.repositories.neo4j.cypher_queries import PERCEPTUAL_CHUNK_EDGE_SAVE
            perceptual_edge_data = []
            for edge in perceptual_edges:
                print(edge.source, edge.target)
                perceptual_edge_data.append({
                    "perceptual_id": edge.source,
                    "chunk_id": edge.target,
                    "end_user_id": edge.end_user_id,
                    "created_at": edge.created_at.isoformat() if edge.created_at else None,
                })
            result = await tx.run(PERCEPTUAL_CHUNK_EDGE_SAVE, edges=perceptual_edge_data)
            perceptual_edges_uuids = [record["uuid"] async for record in result]
            results['perceptual_chunk_edges'] = perceptual_edges_uuids
            logger.info(f"Successfully saved {len(perceptual_edges_uuids)} perceptual-chunk edges to Neo4j")

        return results

    try:
        # 使用显式写事务执行所有操作，避免死锁
        results = await connector.execute_write_transaction(_save_all_in_transaction)
        summary = {
            key: len(value)
            for key, value in results.items()
            if isinstance(value, (list, tuple, set))
        }
        logger.info("Transaction completed. Summary: %s", summary)
        logger.debug("Full transaction results: %r", results)

        return True

    except Exception as e:
        logger.error(f"Neo4j integration error: {e}", exc_info=True)
        print(f"Neo4j integration error: {e}")
        print("Continuing without database storage...")
        return False


async def _trigger_clustering_sync(
        entity_nodes: List,
        llm_model_id: Optional[str] = None,
        embedding_model_id: Optional[str] = None,
) -> None:
    """
    同步等待聚类完成，避免与其他 LLM 任务并发冲突。
    """
    if not entity_nodes:
        return

    clustering_enabled = os.getenv("CLUSTERING_ENABLED", "true").lower() != "false"
    if not clustering_enabled:
        logger.info("[Clustering] 聚类已禁用（CLUSTERING_ENABLED=false），跳过聚类触发")
        return

    end_user_id = entity_nodes[0].end_user_id
    new_entity_ids = [e.id for e in entity_nodes]
    logger.info(f"[Clustering] 准备触发聚类（同步），实体数: {len(new_entity_ids)}, end_user_id: {end_user_id}")
    await _trigger_clustering(new_entity_ids, end_user_id, llm_model_id=llm_model_id,
                              embedding_model_id=embedding_model_id)


async def _trigger_clustering(
        new_entity_ids: List[str],
        end_user_id: str,
        llm_model_id: Optional[str] = None,
        embedding_model_id: Optional[str] = None,
) -> None:
    """
    聚类触发函数，自动判断全量初始化还是增量更新。
    """
    connector = None
    try:
        from app.core.memory.storage_services.clustering_engine import LabelPropagationEngine
        logger.info(f"[Clustering] 开始聚类，end_user_id={end_user_id}, 实体数={len(new_entity_ids)}")
        connector = Neo4jConnector()
        engine = LabelPropagationEngine(connector, llm_model_id=llm_model_id, embedding_model_id=embedding_model_id)
        await engine.run(end_user_id=end_user_id, new_entity_ids=new_entity_ids)
        logger.info(f"[Clustering] 聚类完成，end_user_id={end_user_id}")
    except Exception as e:
        logger.error(f"[Clustering] 聚类触发失败: {e}", exc_info=True)
    finally:
        if connector:
            try:
                await connector.close()
            except Exception:
                pass