# -*- coding: utf-8 -*-
import os
import ast
import re
import asyncio
import httpx
from typing import List, Dict
from openai import AsyncOpenAI

from neo4j import GraphDatabase
from neo4j_graphrag.retrievers import HybridCypherRetriever
from neo4j_graphrag.embeddings.base import Embedder
#from modelsAPI.model_api import LocalBgeM3Embeddings, OpenaiAPI
from config import MAP_INS, NEO4J_CONFIG, INDEX_CONFIG, EMBEDDING_CONFIG
from config import LLM_CONFIG, EMBEDDING_CONFIG

# ================== 配置 ==================
URI                 = NEO4J_CONFIG["uri"]
AUTH                = (NEO4J_CONFIG["user"], NEO4J_CONFIG["password"])
NAME_PROPERTY       = INDEX_CONFIG["NAME_PROPERTY"]
FULLTEXT_PROPERTY   = INDEX_CONFIG["FULLTEXT_PROPERTY"]
EMBEDDING_PROPERTY  = INDEX_CONFIG["EMBEDDING_PROPERTY"]
SEARCH_LABEL        = INDEX_CONFIG["SEARCH_LABEL"]
VECTOR_INDEX_NAME   = INDEX_CONFIG["VECTOR_INDEX_NAME"]
FULLTEXT_INDEX_NAME = INDEX_CONFIG["FULLTEXT_INDEX_NAME"]
map_ins             = MAP_INS

# ================== 全局复用：driver / embedder 只初始化一次 ==================
_driver   = None
_embedder = None

def get_driver():
    global _driver
    if _driver is None:
        _driver = GraphDatabase.driver(URI, auth=AUTH)
    return _driver

def get_embedder():
    global _embedder
    if _embedder is None:
        _embedder = LocalBgeM3Embeddings(
            base_url=EMBEDDING_CONFIG["base_url_v2"],
            api_key=EMBEDDING_CONFIG["api_key"],
        )
    return _embedder

class LocalBgeM3Embeddings(Embedder):
    def __init__(self, base_url: str, api_key: str, model_name: str = "bge-m3"):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.model_name = model_name

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        url = self.base_url + "/embeddings"
        with httpx.Client(timeout=60) as client:
            response = client.post(
                url,
                json={"model": self.model_name, "input": texts},
                headers={"Authorization": f"Bearer {self.api_key}"},
            )
            return [item["embedding"] for item in response.json()["data"]]
async def open_api_chat_without_thinking(
        query: str = None,
        model: str = None,
        json_output: bool = False,
        system_prompt: str = None,
        messages: list = None,
        enable_thinking: bool = True,
        temperature: float = 0.1
    ) -> str:
        if model is None:
            model = LLM_CONFIG["model"]
        if system_prompt is None:
            system_prompt = "我是一个专业的AI助手，请用中文回答问题。请直接给出最终答案，不要包含任何推理步骤、思考过程、解释性文字或前缀。"


        # 构建消息
        message_list = [{"role": "system", "content": system_prompt}]
    
        if messages is not None:
            message_list.extend(messages)
        if query is not None:
            message_list.append({"role": "user", "content": query})

        client = AsyncOpenAI(
            api_key=LLM_CONFIG["api_key"],
            base_url=LLM_CONFIG["base_url"],
        )

        kwargs = {
            "model": model,
            "messages": message_list,
            "temperature": temperature,
            "stream": False,
            "max_tokens":LLM_CONFIG["max_tokens"]
        }

        kwargs["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}}

        if json_output:
            kwargs["response_format"] = {"type": "json_object"}
        
        print("model API. history message", message_list)
        response = await client.chat.completions.create(**kwargs)
        return response.choices[0].message.content
# ================== 辅助：解析 Record 字符串 ==================
def parse_record_string(s: str):
    name_match   = re.search(r"name='([^']*)'",      s)
    labels_match = re.search(r"labels=(\[[^\]]*\])", s)
    score_match  = re.search(r"score=([\d.]+)",       s)

    name   = name_match.group(1) if name_match else "未知名称"
    labels = ["未知标签"]
    if labels_match:
        try:
            labels = ast.literal_eval(labels_match.group(1))
        except (ValueError, SyntaxError):
            pass
    score = float(score_match.group(1)) if score_match else 0.0
    return name, labels, score


# ================== 构建 Retriever ==================
def build_retriever(sync_driver, embedder: Embedder) -> HybridCypherRetriever:
    return HybridCypherRetriever(
        driver=sync_driver,
        vector_index_name=VECTOR_INDEX_NAME,
        fulltext_index_name=FULLTEXT_INDEX_NAME,
        embedder=embedder,
        retrieval_query=f"""
            RETURN
                node.`{NAME_PROPERTY}` AS name,
                [lbl IN labels(node) WHERE lbl <> '{SEARCH_LABEL}'] AS labels,
                score
        """,
    )


# ================== 同步检索（在线程池中执行，不阻塞事件循环）==================
def _hybrid_search_sync(sentence: str, top_k: int = 10) -> List[Dict]:
    retriever   = build_retriever(get_driver(), get_embedder())
    results_raw = retriever.search(query_text=sentence, top_k=top_k)

    hybrid_results: List[Dict] = []
    for item in results_raw.items:
        content = item.content
        name    = "未知名称"
        labels  = ["未知标签"]
        score   = 0.0

        if hasattr(content, "keys"):
            name   = content.get("name",   name)
            labels = content.get("labels", labels)
            score  = content.get("score",  score)
        elif isinstance(content, str):
            if content.startswith("<Record"):
                name, labels, score = parse_record_string(content)
            else:
                try:
                    parsed = ast.literal_eval(content)
                    if isinstance(parsed, dict):
                        name   = parsed.get("name",   name)
                        labels = parsed.get("labels", labels)
                        score  = parsed.get("score",  score)
                except (ValueError, SyntaxError):
                    name = content

        if score == 0.0 and hasattr(item, "metadata") and item.metadata:
            score = item.metadata.get("score", score)

        hybrid_results.append({
            "标签": labels,
            "名称": name,
            "得分": score,
            "来源": "hybrid",
            "fulltext_snippet": "",
        })

    return hybrid_results


# ================== Prompt 模板 ==================
PROMPT_TEMPLATE = """
请根据给出的问题，以及问题相关的检索结果，以及实体类型和sql表字段映射表，筛选出与问题相关的信息并输出。
要求：
1. 请仔细分析问题和检索结果，确定问题中出现的实体类型和sql表字段。
2. 输出内容可以为充分分析后的一段文本，即你对这个问题中所用到的实体类型和sql表字段的描述。
3. 给出的信息尽量精简，不能省略任何重要信息。
实体类型和sql表字段映射表：
    "故障模式" : "fault",
    "系统" : "system_name",
    "设备" : "device_name",
    "舷号" : "ship_number",
其中键表示实体类型，值表示sql表字段。

检索结果：
{results}

问题：{user_question}


案例1：
用户问题为：扫气箱内发现润滑油泄漏发生了多少次？
检索知识库为：   
   [ 1] score=1.0000 | 活塞杆填料函密封失效导致扫气箱内发现润滑油泄漏  标签=['故障模式']
   [ 2] score=0.9298 | 排除喷油器雾化不良故障  标签=['维修项目']
   [ 3] score=0.9115 | 执行发动机燃油泄漏应急封堵操作  标签=['操作项目']
   [ 4] score=0.9082 | 活塞环磨损或断裂导致机座与机架组件润滑油消耗异常增加  标签=['故障模式']
   [ 5] score=0.9058 | 检测气缸套内径磨损量  标签=['维修项目']
   [ 6] score=0.9049 | 拆检活塞头磨损情况  标签=['维修项目']
   [ 7] score=0.9048 | 检查排气阀杆密封面磨损  标签=['维修项目']
   [ 8] score=0.9015 | 润滑油压力骤降导致轴承组件紧急停机  标签=['故障模式']
   [ 9] score=0.8985 | 喷油器组件的维修工作  标签=['维修工作']
   [10] score=0.8979 | 发动机的安全警告  标签=['安全警告']
输出为：
问题中扫气箱内发现润滑油泄漏是一个故障现象，属于表中fault字段。
"""

SYSTEM_PROMPT = "你是一个专业信息筛选助手，根据给出的问题和问题检索到的相关实体关系，筛选出与问题相关的信息。"


# ================== 主函数（传参调用）==================
async def get_extra_info(user_question: str, top_k: int = 20) -> str:
    """
    参数:
        user_question: 用户问题
        top_k:         检索返回条数
    返回:
        LLM 筛选结果字符串
    """
    # 同步检索放入线程池，不阻塞事件循环
    results = await asyncio.to_thread(_hybrid_search_sync, user_question, top_k)

    # 构建 prompt 并异步调用 LLM
    res = await open_api_chat_without_thinking(
        system_prompt=SYSTEM_PROMPT,
        query=PROMPT_TEMPLATE.format(results=results, user_question=user_question),
    )

    return res.strip()


# ================== 入口 ==================
if __name__ == "__main__":
    answer = asyncio.run(get_extra_info("发动机这个月发生了多少次故障"))
    print(111111111111111111)
    print(answer)