# neo4j-admin database load --from-path="neo4j.dump文件所在目录" --overwrite-destination=true neo4j --verbose
import os
import re
import jieba
import logging
import numpy as np
from concurrent.futures import ThreadPoolExecutor
from neo4j import GraphDatabase
from neo4j_graphrag.indexes import (
    create_vector_index,
    upsert_vectors,
    create_fulltext_index,
)
from modelsAPI.model_api import OpenaiAPI
from dotenv import load_dotenv
from typing import List, Set, Tuple, Dict, Optional
from config import LLM_CONFIG, EMBEDDING_CONFIG, NEO4J_CONFIG, SEARCH_CONFIG

# 加载 .env 文件中的环境变量
load_dotenv()

# 配置控制台日志
logger = logging.getLogger("indexing")
logger.setLevel(logging.INFO)
if not logger.handlers:
    formatter = logging.Formatter("[%(levelname)s]%(asctime)s: %(message)s")
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    
VECTOR_DIMENSION = SEARCH_CONFIG["vector_dimension"]
vector_dim = SEARCH_CONFIG["vector_dimension"]  # OpenAI bge-m3 模型的嵌入向量维度
embed_batch_size = SEARCH_CONFIG["embed_batch_size"]  # 嵌入向量计算批次大小
EMBED_MAX_WORKERS = SEARCH_CONFIG["embed_thread_num"]  # embedding 批次并行线程数（受下游 OpenaiAPI 限流约束，勿过大）
NODE_FETCH_PAGE = SEARCH_CONFIG["node_fetch_page"]  # 单次拉取待生成 embedding 节点的分页大小

# 全局索引名称（覆盖所有节点标签）
DEFAULT_URI = NEO4J_CONFIG["uri"]
DEFAULT_AUTH = (
    NEO4J_CONFIG["username"],
    NEO4J_CONFIG["password"],
)

NAME_PROPERTY = SEARCH_CONFIG["name_property"]
FULLTEXT_PROPERTY = SEARCH_CONFIG["fulltext_property"]
EMBEDDING_PROPERTY = SEARCH_CONFIG["embedding_property"]
SEARCH_LABEL = SEARCH_CONFIG["search_label"]

VECTOR_INDEX_NAME = SEARCH_CONFIG["vector_index_name"]
FULLTEXT_INDEX_NAME = SEARCH_CONFIG["fulltext_index_name"]
FULLTEXT_FIELD_INDEX_NAME = SEARCH_CONFIG["fulltext_field_index_name"]

EXCLUDED_BUSINESS_LABELS: Set[str] = set(SEARCH_CONFIG["excluded_business_labels"])
OLD_INDEX_NAMES: List[str] = [
    VECTOR_INDEX_NAME,
    FULLTEXT_INDEX_NAME,
    FULLTEXT_FIELD_INDEX_NAME,
]


def get_all_labels(driver):
    """
    获取知识图谱中所有的节点标签

    Args:
        driver: Neo4j driver

    Returns:
        list: 所有节点标签列表
    """
    query = "CALL db.labels()"
    result = driver.execute_query(query)
    labels = [record["label"] for record in result.records]
    return labels


def get_label_properties(driver, label):
    """
    获取指定标签的所有属性键（通过采样多个节点获取所有可能的属性）

    Args:
        driver: Neo4j driver
        label: 节点标签

    Returns:
        list: 属性键列表（排除内部属性如 embedding, fulltext）
    """
    # 采样多个节点来获取所有可能的属性（因为不同节点可能有不同的属性集）
    query = f"""
    MATCH (n:{label})
    WITH n LIMIT 100
    UNWIND keys(n) AS prop
    RETURN DISTINCT prop AS properties
    ORDER BY properties
    """
    result = driver.execute_query(query)

    if result.records:
        properties = [record["properties"] for record in result.records]
        # 排除内部索引属性
        excluded_props = {"embedding", "fulltext", "elementId", "id"}
        properties = [prop for prop in properties if prop not in excluded_props]
        return properties
    else:
        logger.warning(f"标签 {label} 没有节点，无法获取属性列表")
        return []


def build_text_from_all_properties(record_dict):
    """
    从节点的所有属性构建文本（用于全文索引）

    Args:
        record_dict: 节点属性字典（已排除 id）

    Returns:
        str: 构建的文本字符串
    """
    text_parts = []
    for key, value in record_dict.items():
        if value is not None:
            # 将值转换为字符串
            if isinstance(value, (list, dict)):
                value_str = str(value)
            else:
                value_str = str(value).strip()

            if value_str:
                text_parts.append(f"{key}：{value_str}")

    return " ".join(text_parts) if text_parts else None


def _compute_embeddings_parallel(names: List[str]) -> List[list]:
    """
    并行计算一批名称的 embedding。

    将 names 切成多个 batch，用线程池并发调用 OpenaiAPI.batch_embeddings，
    再按原始顺序拼回。相比串行可显著缩短大批量节点的 embedding 耗时。

    Args:
        names: 待计算 embedding 的文本列表

    Returns:
        list: 与 names 等长、顺序一致的 embedding 向量列表

    Raises:
        任一 batch 计算失败时向上抛出异常
    """
    # 切分批次，保留批次索引以便按序拼回
    batches = [
        (idx, names[i:i + embed_batch_size])
        for idx, i in enumerate(range(0, len(names), embed_batch_size))
    ]

    results: Dict[int, list] = {}

    def _run_batch(item):
        idx, batch_names = item
        batch_result = OpenaiAPI.batch_embeddings(batch_names)
        # batch_result 返回的是字典列表，每个字典包含 'embedding' 字段
        return idx, [row["embedding"] for row in batch_result]

    workers = min(EMBED_MAX_WORKERS, len(batches)) or 1
    with ThreadPoolExecutor(max_workers=workers) as executor:
        for idx, batch_embeddings in executor.map(_run_batch, batches):
            results[idx] = batch_embeddings

    # 按批次顺序拼回，保证与输入 names 对齐
    all_embeddings: List[list] = []
    for idx in range(len(batches)):
        all_embeddings.extend(results[idx])
    return all_embeddings


def drop_constraint(driver):
    """删除所有约束"""
    records = driver.execute_query("show constraints").records
    for record in records:
        driver.execute_query(f"drop constraint {record['name']} if exists")


def drop_index_without_constraint(driver):
    """删除所有没有约束的索引"""
    records = driver.execute_query("show index").records
    for record in records:
        if not record["owningConstraint"]:
            driver.execute_query(f"drop index {record['name']} if exists")


def print_neo4j_version(driver) -> None:
    """打印当前 Neo4j 实例的版本和版本类型。"""
    with driver.session() as session:
        rec = session.run(
            "CALL dbms.components() YIELD name, versions, edition "
            "WHERE name='Neo4j Kernel' RETURN versions[0] AS v, edition"
        ).single()
        if rec:
            logger.info(f"Neo4j 版本: {rec['v']} ({rec['edition']})")


def migrate_remove_entity_label(driver, dry_run: bool = False) -> None:
    """
    从所有节点上移除 :Entity 标签。

    Parameters
    ----------
    driver   : Neo4j GraphDatabase driver
    dry_run  : 若为 True，只统计数量，不执行移除
    """
    with driver.session() as session:
        cnt = session.run("MATCH (n:Entity) RETURN count(n) AS c").single()["c"]
        logger.info(f"当前带 :Entity 标签的节点数: {cnt}")
        if cnt == 0:
            return
        if dry_run:
            logger.info("[dry_run] 未执行移除")
            return
        session.run("MATCH (n:Entity) REMOVE n:Entity")
        logger.info(f"已从 {cnt} 个节点上移除 :Entity")


def tag_searchable_nodes(driver, excluded: Set[str]) -> int:
    """
    为满足以下条件的节点打上 :Searchable 标签：
      - 拥有 `名称` 属性
      - 拥有 `embedding` 属性
      - 尚未持有 :Searchable 标签
      - 至少有一个非排除业务标签

    注意：排除标签列表通过查询参数传入，避免 Cypher 字符串注入。

    Returns
    -------
    int : 本次新打标签的节点数
    """
    # 标签名无法参数化，但排除标签集合的“值”可以参数化，避免拼接注入
    query = f"""
    MATCH (n)
    WHERE n.`{NAME_PROPERTY}` IS NOT NULL
      AND n.{EMBEDDING_PROPERTY} IS NOT NULL
      AND NOT n:{SEARCH_LABEL}
      AND any(lbl IN labels(n) WHERE NOT lbl IN $excluded)
    SET n:{SEARCH_LABEL}
    RETURN count(n) AS added
    """
    excluded_list = list(excluded)
    with driver.session() as session:
        added = session.run(query, excluded=excluded_list).single()["added"]
    logger.info(f"本次新打上 :{SEARCH_LABEL} 的节点数: {added}")

    with driver.session() as session:
        total = session.run(
            f"MATCH (n:{SEARCH_LABEL}) RETURN count(n) AS c"
        ).single()["c"]
    logger.info(f"当前带 :{SEARCH_LABEL} 标签的节点总数: {total}")
    return added


def drop_old_indexes(driver, names: List[str]) -> None:
    """删除指定名称的索引（不存在时忽略）。"""
    with driver.session() as session:
        for name in names:
            try:
                session.run(f"DROP INDEX {name} IF EXISTS")
                logger.info(f"已删除旧索引: {name}")
            except Exception as e:
                logger.warning(f"删除索引 {name} 失败（可忽略）: {e}")


def create_indexes(driver) -> None:
    """
    创建三个索引：
      1. 向量索引            —— 用于语义相似度检索
      2. 名称全文索引        —— 用于节点名称关键词检索
      3. fulltext 字段全文索引 —— 用于节点 fulltext 属性关键词检索
    """
    with driver.session() as session:

        # 1. 向量索引
        session.run(f"""
        CREATE VECTOR INDEX {VECTOR_INDEX_NAME} IF NOT EXISTS
        FOR (n:{SEARCH_LABEL}) ON (n.{EMBEDDING_PROPERTY})
        OPTIONS {{
            indexConfig: {{
                `vector.dimensions`: {VECTOR_DIMENSION},
                `vector.similarity_function`: 'cosine'
            }}
        }}
        """)
        logger.info(f"向量索引 '{VECTOR_INDEX_NAME}' 已就绪")

        # 2. 名称全文索引
        session.run(f"""
        CREATE FULLTEXT INDEX {FULLTEXT_INDEX_NAME} IF NOT EXISTS
        FOR (n:{SEARCH_LABEL}) ON EACH [n.`{NAME_PROPERTY}`]
        """)
        logger.info(f"全文索引（名称） '{FULLTEXT_INDEX_NAME}' 已就绪")

        # 3. fulltext 字段全文索引
        session.run(f"""
        CREATE FULLTEXT INDEX {FULLTEXT_FIELD_INDEX_NAME} IF NOT EXISTS
        FOR (n:{SEARCH_LABEL}) ON EACH [n.`{FULLTEXT_PROPERTY}`]
        """)
        logger.info(f"全文索引（fulltext 字段） '{FULLTEXT_FIELD_INDEX_NAME}' 已就绪")


# ================== 连接解析工具 ==================

def _resolve_connection(
    uri: Optional[str],
    user: Optional[str],
    password: Optional[str],
) -> Tuple[str, Tuple[str, str]]:
    """
    按优先级解析连接参数：传参 > 环境变量 > 默认常量。

    Returns
    -------
    (uri, (user, password))
    """
    resolved_uri = uri or os.getenv("NEO4J_URI", DEFAULT_URI)
    resolved_user = user or os.getenv("NEO4J_USER", DEFAULT_AUTH[0])
    resolved_pw = password or os.getenv("NEO4J_PASSWORD", DEFAULT_AUTH[1])
    return resolved_uri, (resolved_user, resolved_pw)


# ================== 主入口函数 ==================

def build_hybrid_indexes(
    driver=None,
    *,
    uri: Optional[str] = None,
    user: Optional[str] = None,
    password: Optional[str] = None,
    remove_entity_label: bool = False,
    dry_run_entity: bool = False,
    drop_old: bool = True,
) -> dict:
    """
    完整的混合索引构建流程。

    连接参数优先级：传参 > 环境变量(NEO4J_URI/NEO4J_USER/NEO4J_PASSWORD) > 默认常量

    Parameters
    ----------
    driver              : 已有的 Neo4j driver（传入则直接使用，不再创建新连接）
    uri                 : Neo4j 连接地址，如 "bolt://host:7687"
    user                : 数据库用户名
    password            : 数据库密码
    remove_entity_label : 是否执行移除 :Entity 标签的迁移步骤
    dry_run_entity      : remove_entity_label=True 时，是否仅统计而不实际移除
    drop_old            : 是否在创建前先删除旧索引（建议保持 True）

    Returns
    -------
    dict : {"success": bool, "message": str, "searchable_added": int}
           失败时 success=False 且 message 含错误信息
    """
    # 若未传入 driver，则根据参数/环境变量/默认值自动创建
    _owns_driver = driver is None
    if _owns_driver:
        resolved_uri, resolved_auth = _resolve_connection(uri, user, password)
        logger.info(f"连接 Neo4j: {resolved_uri}  用户: {resolved_auth[0]}")
        driver = GraphDatabase.driver(resolved_uri, auth=resolved_auth)

    try:
        logger.info("=" * 60)
        logger.info("开始构建混合索引")
        logger.info("=" * 60)

        print_neo4j_version(driver)

        if remove_entity_label:
            migrate_remove_entity_label(driver, dry_run=dry_run_entity)

        searchable_added = tag_searchable_nodes(driver, EXCLUDED_BUSINESS_LABELS)

        if drop_old:
            drop_old_indexes(driver, OLD_INDEX_NAMES)

        create_indexes(driver)

        return {
            "success": True,
            "message": "混合索引构建完成",
            "searchable_added": searchable_added,
        }

    except Exception as e:
        error_msg = f"构建混合索引过程中发生错误: {e}"
        logger.error(error_msg, exc_info=True)
        return {
            "success": False,
            "message": error_msg,
            "searchable_added": 0,
        }

    finally:
        # 仅关闭由本函数自己创建的 driver
        if _owns_driver:
            driver.close()


# 向后兼容别名：保留旧拼写，避免其他模块仍引用 build_hyrid_indexes 时报错
build_hyrid_indexes = build_hybrid_indexes


if __name__ == "__main__":

    result = build_hybrid_indexes(
        drop_old=True
    )

    if result["success"]:
        logger.info("混合索引创建成功")
    else:
        logger.error(result["message"])