import asyncio
import httpx
from typing import List, Union
import os
import numpy as np
from openai import OpenAI
from config import LLM_CONFIG, EMBEDDING_CONFIG
from neo4j_graphrag.embeddings.base import Embedder

from openai import AsyncOpenAI

# 模块级单例客户端，避免每次调用都创建新连接
_async_client: AsyncOpenAI = None


def _get_async_client() -> AsyncOpenAI:
    """获取或创建 AsyncOpenAI 单例客户端"""
    global _async_client
    if _async_client is None:
        _async_client = AsyncOpenAI(
            api_key=LLM_CONFIG["api_key"],
            base_url=LLM_CONFIG["base_url"],
        )
    return _async_client


class OpenaiAPI:
    @staticmethod
    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 = _get_async_client()

        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
    
    @staticmethod
    async def open_api_chat_stream(
        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
    ):
        """
        流式输出大模型响应
        
        Yields:
            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 = _get_async_client()

        kwargs = {
            "model": model,
            "messages": message_list,
            "temperature": temperature,
            "stream": True,
            "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 stream. history message", message_list)
        
        stream = await client.chat.completions.create(**kwargs)
        async for chunk in stream:
            if chunk.choices and len(chunk.choices) > 0:
                delta = chunk.choices[0].delta
                if delta.content:
                    yield delta.content

    @staticmethod   
    async def open_api_vl_without_thinking(
        image_url: str
    ) -> str:
        """
        专门用于从图片中提取【设备名称】和【故障现象】。
        自动过滤思考过程，仅返回核心结果。
        
        参数:
            image_url: Base64 格式的图片数据 (data:image/...;base64,...)
            model: 模型名称，默认为配置中的模型
            custom_instruction: 额外的特定指令 (可选)
        """
        
        # 1. 确定模型名称
        model = LLM_CONFIG.get("model")
        
        # 2. 构建强约束的 System Prompt
        # 核心目标：禁止思考标签，禁止废话，只给结果
        system_prompt = (
            "你是一个工业视觉分析专家。你的任务是从图片中识别设备并诊断故障。\n"
            "【严格约束】\n"
            "1. 绝对禁止输出 <think>, <think>, <reasoning> 等任何思考过程标签。\n"
            "2. 绝对禁止输出'好的'、'根据图片'、'分析如下'等开场白或结束语。\n"
            "3. 直接输出最终结论，格式必须严格遵守下面的模板。"
        )

        # 3. 构建针对性的 User Prompt
        default_task = (
            "请分析这张图片，描述图片内容，重点关注以下信息\n"
            "1. 设备识别：图中是什么设备？型号或标签是否可见？\n"
            "2. 状态评估：仪表盘数值是否在正常区间？指示灯颜色是什么？\n"
            "3. 异常检测：是否存在漏油、生锈、电缆破损、螺丝松动或烟雾？\n"
            "4. 环境风险：周围是否有杂物堆放或安全隐患？\n"
            "请用专业、客观、简短的语言描述。"
        )
        
        final_user_prompt = default_task

        # 4. 初始化客户端
        client = _get_async_client()

        # 5. 构建请求参数
        kwargs = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": final_user_prompt},
                        {"type": "image_url", "image_url": {"url": image_url}}
                    ]
                }
            ],
            "temperature": 0.1,  # 低温度以保证事实准确性
            "stream": False,
            "max_tokens": LLM_CONFIG.get("max_tokens"),
        }

        # 尝试通过参数关闭思考 (取决于后端支持情况)
        kwargs["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}}

        try:
            print(f"🚀 正在调用 {model} 进行设备故障分析...")
            response = await client.chat.completions.create(**kwargs)
            raw_content = response.choices[0].message.content or ""

            print(f"✅ 分析完成:\n{raw_content}")
            return raw_content

        except Exception as e:
            import traceback
            error_msg = f"❌ 设备故障分析失败：{str(e)}\n{traceback.format_exc()}"
            print(error_msg)
            return f"Error: {str(e)}"

    @staticmethod
    def get_embeddings(embedding_text: str):
        embedding_base_url = EMBEDDING_CONFIG["base_url"]

        with httpx.Client(timeout=60) as client:
            response = client.post(
            embedding_base_url,
            json={"model": EMBEDDING_CONFIG["model"], "input": embedding_text},
            headers={"Authorization": EMBEDDING_CONFIG["api_key"]},
        )
            embedding_text = response.json()["data"][0]["embedding"]
        return embedding_text

    @staticmethod
    async def get_embeddings_async(embedding_text: str) -> List[float]:
        embedding_base_url = EMBEDDING_CONFIG["base_url"]

        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(
                embedding_base_url,
                json={"model": EMBEDDING_CONFIG["model"], "input": embedding_text},
                headers={"Authorization": EMBEDDING_CONFIG["api_key"]},
            )
            embedding = response.json()["data"][0]["embedding"]
        return embedding

    @staticmethod
    async def get_embeddings_batch_async(texts: List[str]) -> List[List[float]]:
        embedding_base_url = EMBEDDING_CONFIG["base_url"]

        async with httpx.AsyncClient(timeout=120) as client:
            response = await client.post(
                embedding_base_url,
                json={"model": EMBEDDING_CONFIG["model"], "input": texts},
                headers={"Authorization": EMBEDDING_CONFIG["api_key"]},
            )
            data = response.json()["data"]
            sorted_data = sorted(data, key=lambda x: x["index"])
            embeddings = [item["embedding"] for item in sorted_data]
        return embeddings

    @staticmethod
    def cosine_similarity(vec1: Union[List[float], np.ndarray], vec2: Union[List[float], np.ndarray]) -> float:
        v1 = np.array(vec1) if not isinstance(vec1, np.ndarray) else vec1
        v2 = np.array(vec2) if not isinstance(vec2, np.ndarray) else vec2
        
        norm1 = np.linalg.norm(v1)
        norm2 = np.linalg.norm(v2)
        
        if norm1 == 0 or norm2 == 0:
            return 0.0
        
        return float(np.dot(v1, v2) / (norm1 * norm2))

    @staticmethod
    def cosine_similarity_batch(query_vec: Union[List[float], np.ndarray], 
                                  candidates_vecs: List[Union[List[float], np.ndarray]]) -> List[float]:
        query = np.array(query_vec) if not isinstance(query_vec, np.ndarray) else query_vec
        candidates = [np.array(v) if not isinstance(v, np.ndarray) else v for v in candidates_vecs]
        
        query_norm = np.linalg.norm(query)
        if query_norm == 0:
            return [0.0] * len(candidates)
        
        similarities = []
        for candidate in candidates:
            cand_norm = np.linalg.norm(candidate)
            if cand_norm == 0:
                similarities.append(0.0)
            else:
                sim = float(np.dot(query, candidate) / (query_norm * cand_norm))
                similarities.append(sim)
        
        return similarities

    @staticmethod
    async def hybrid_match_with_embeddings(
        query_text: str,
        candidates: List[dict],
        name_key: str = "display_name",
        embedding_key: str = "embedding",
        text_weight: float = 0.3,
        semantic_weight: float = 0.7,
        text_threshold: float = 0.3,
        semantic_threshold: float = 0.5
    ) -> tuple:
        if not candidates:
            return None, 0.0
        
        query_embedding = await OpenaiAPI.get_embeddings_async(query_text)
        
        scored_candidates = []
        
        for idx, candidate in enumerate(candidates):
            name = candidate.get(name_key, "")
            if not name:
                props = candidate.get("props", {})
                name = props.get("name") or props.get("名称") or props.get("设备名称") or ""
            
            text_score = OpenaiAPI._compute_text_similarity(query_text, name)
            
            props = candidate.get("props", {})
            candidate_embedding = props.get(embedding_key) or candidate.get(embedding_key)
            
            if candidate_embedding:
                semantic_score = OpenaiAPI.cosine_similarity(query_embedding, candidate_embedding)
            else:
                semantic_score = 0.0
            
            combined_score = text_weight * text_score + semantic_weight * semantic_score
            
            scored_candidates.append({
                "candidate": candidate,
                "text_score": text_score,
                "semantic_score": semantic_score,
                "combined_score": combined_score,
                "index": idx
            })
        
        scored_candidates.sort(key=lambda x: x["combined_score"], reverse=True)
        
        best = scored_candidates[0]
        best_candidate = best["candidate"]
        best_score = best["combined_score"]
        
        min_threshold = max(text_threshold * text_weight + semantic_threshold * semantic_weight, 0.3)
        
        if best_score < min_threshold:
            return None, best_score
        
        return best_candidate, best_score

    @staticmethod
    def _compute_text_similarity(a: str, b: str) -> float:
        if not a or not b:
            return 0.0
        a = str(a).strip().lower()
        b = str(b).strip().lower()
        if not a or not b:
            return 0.0
        if a == b:
            return 1.0
        if a in b or b in a:
            return 0.9
        
        set_a, set_b = set(a), set(b)
        inter = len(set_a & set_b)
        union = len(set_a | set_b) or 1
        jaccard = inter / union
        
        len_diff = abs(len(a) - len(b))
        len_penalty = max(0.0, 1.0 - len_diff / max(len(a), len(b), 1))
        
        return 0.5 * jaccard + 0.5 * len_penalty

class LocalBgeM3Embeddings(Embedder):
    def __init__(self, base_url: str, api_key: str, model_name: str = "bge-m3"):
        self.base_url = base_url
        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 
        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"]]
    

if __name__ == "__main__":
    answer = asyncio.run(OpenaiAPI.open_api_chat_without_thinking(
        query="你好",
        json_output=True,
        system_prompt="你是一个助手，请用中文回答问题。请直接给出最终答案，不要包含任何推理步骤、思考过程、解释性文字或前缀。",
        enable_thinking=True
    ))
    print(answer)
        


