import httpx
from openai import OpenAI
import os
import requests
import json
import asyncio
import aiohttp
from typing import List, Dict, Any,Optional
from tqdm import tqdm  # 进度条支持（可选）
from openai import AsyncOpenAI
from neo4j_graphrag.embeddings.base import Embedder

#now_model = "Qwen3.5-35B-A3B"
#now_model = "Qwen3-32B"

from config import LLM_CONFIG, EMBEDDING_CONFIG


class OpenaiAPI:
    
    @staticmethod
    def get_embeddings(embedding_texts: List[str]):
        # ✅ 2. 使用 EMBEDDING_CONFIG
        base_url = EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", ""))
        embedding_base_url = os.path.join(base_url, "embeddings")
        
        with httpx.Client(timeout=60) as client:
            response = client.post(
                embedding_base_url,
                json={"model": EMBEDDING_CONFIG["model"], "input": embedding_texts},
                headers={"Authorization": f"Bearer {EMBEDDING_CONFIG['api_key']}"},
            )
            # 注意：batch_embeddings 返回的是列表，get_embeddings 原代码只取了第一个元素的 embedding
            # 这里保持原逻辑，如果输入是列表，通常期望返回所有 embedding，原代码可能有误，此处仅替换配置源
            data = response.json()["data"]
            if isinstance(data, list) and len(data) > 0:
                 return data[0]["embedding"]
            return []

    @staticmethod
    def batch_embeddings(embedding_texts: List[str]):
        # ✅ 3. 使用 EMBEDDING_CONFIG
        base_url = EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", ""))
        embedding_base_url = os.path.join(base_url, "embeddings")
        
        with httpx.Client(timeout=60) as client:
            response = client.post(
                embedding_base_url,
                json={"model": EMBEDDING_CONFIG["model"], "input": embedding_texts},
                headers={"Authorization": f"Bearer {EMBEDDING_CONFIG['api_key']}"},
            )
            embedding_list = response.json()["data"]
        return embedding_list

    @staticmethod
    def rerank_query(query: str, documents: List[str], top_n: str = 3):
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {EMBEDDING_CONFIG["api_key"]}' # 假设 Rerank 也使用相同的 Key 或在此处扩展配置
        }

        data = {
            "model": os.getenv("OPENAI_RERANKER_MODEL", "bge-reranker-v2-m3"), # Rerank 模型未在 config.py 定义，暂保留环境变量或可添加到 config
            "query": query,
            "top_n": top_n,
            "documents": documents
        }

        # 假设 Rerank 地址与 Embedding 地址相近，或者需要单独配置。此处暂用环境变量或默认逻辑
        rerank_base = os.getenv("OPENAI_API_RERANKER_BASE", EMBEDDING_CONFIG["base_url_v2"])
        response = requests.post(f"{rerank_base}/rerank", headers=headers, data=json.dumps(data))
        
        results_list = []
        if response.status_code == 200:
            result_json = response.json()
            for result in result_json["results"]:
                results_list.append({
                    "index": result["index"],
                    "text": result["document"]["text"],
                    "scores": result["relevance_score"]
                })
            return results_list
        else:
            raise Exception(f"Error: {response.status_code}\n{response.text}")

    @staticmethod
    def embedding_cosine_similarity(text_1, text_2):
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {EMBEDDING_CONFIG["api_key"]}'
        }
        data = {
            "text_1": text_1,
            "text_2": text_2,
            "model": "bge-rerank"
        }
        rerank_base = os.getenv("OPENAI_API_RERANKER_BASE", EMBEDDING_CONFIG["base_url_v2"])
        response = requests.post(f"{rerank_base}/score", headers=headers, data=json.dumps(data))
        
        if response.status_code == 200:
            result_json = response.json()
            score = result_json["data"][0]["score"]
            return score
        else:
            raise Exception(f"Error: {response.status_code}\n{response.text}")

    @staticmethod
    def open_api_chat(query: str, model: str = None):
        # ✅ 4. 使用 LLM_CONFIG
        if model is None:
            model = LLM_CONFIG["model"]
            
        response = OpenAI(
            api_key=LLM_CONFIG["api_key"],
            base_url=LLM_CONFIG["base_url"],
        ).chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query},
            ],
            temperature=LLM_CONFIG["temperature"],
            stream=False
        )
        return response.choices[0].message.content
    
    @staticmethod
    async def open_api_chat_async(query: str, model: str = None, temperature: float = None) -> str:
        if model is None:
            model = LLM_CONFIG["model"]
        if temperature is None:
            temperature = LLM_CONFIG["temperature"]

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

        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query}
            ],
            max_tokens=LLM_CONFIG["max_tokens"],
            temperature=temperature,
            stream=False,
            extra_body={"chat_template_kwargs": {"enable_thinking": False}}
        )
        return response.choices[0].message.content

    @staticmethod
    async def open_api_chat_async_json(query: str, model: str = None, temperature: float = None) -> str:
        if model is None:
            model = LLM_CONFIG["model"]
        if temperature is None:
            temperature = LLM_CONFIG["temperature"]

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

        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant that outputs only valid JSON."},
                {"role": "user", "content": query}
            ],
            temperature=temperature,
            stream=False,
            max_tokens=LLM_CONFIG["max_tokens"],
            response_format={"type": "json_object"},
            extra_body={"chat_template_kwargs": {"enable_thinking": False}}
        )
        return response.choices[0].message.content        

    @staticmethod
    def open_api_caht_without_thinking(query: str, model: str = None):
        if model is None:
            model = LLM_CONFIG["model"]
            
        response = OpenAI(
            api_key=LLM_CONFIG["api_key"],
            base_url=LLM_CONFIG["base_url"],
        ).chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query},
            ],
            temperature=LLM_CONFIG["temperature"],
            stream=False,
            extra_body={"chat_template_kwargs": {"enable_thinking": False}}
        )
        return response.choices[0].message.content

    @staticmethod
    def open_api_chat_stream(query: str, model: str = None):
        if model is None:
            model = LLM_CONFIG["model"]

        response = OpenAI(
            api_key=LLM_CONFIG["api_key"],
            base_url=LLM_CONFIG["base_url"],
        ).chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query},
            ],
            temperature=LLM_CONFIG["temperature"],
            stream=True
        )
        for chunk in response:
            content = chunk.choices[0].delta.content
            if content:
                yield content.encode('utf-8')

    @staticmethod
    def openai_chat(query: str, timeout: int = None):
        if timeout is None:
            timeout = LLM_CONFIG["timeout"]
            
        try:
            client = OpenAI(
                api_key=LLM_CONFIG["api_key"],
                base_url=LLM_CONFIG["base_url"],
                timeout=timeout
            )

            response = client.chat.completions.create(
                model=LLM_CONFIG["model"],
                messages=[
                    {"role": "user", "content": query},
                ],
                temperature=LLM_CONFIG["temperature"],
                stream=False,
                response_format={"type": "json_object"},
                extra_body={"chat_template_kwargs": {"enable_thinking": False}}
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用OpenAI API时出错: {e}")
            return None
        
    @staticmethod
    async def openai_chat_aysnc(query: str, timeout: int = None) -> str | None:
        if timeout is None:
            timeout = LLM_CONFIG["timeout"]

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

            response = await client.chat.completions.create(
                model=LLM_CONFIG["model"],
                messages=[
                    {"role": "user", "content": query},
                ],
                temperature=LLM_CONFIG["temperature"],
                stream=False,
                response_format={"type": "json_object"},
                extra_body={"chat_template_kwargs": {"enable_thinking": False}}
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用 OpenAI API 时出错: {e}")
            return None
        finally:
            await client.close()

    @staticmethod
    async def openai_chat_aysnc_nothink(query: str, timeout: int = None) -> str | None:
        if timeout is None:
            timeout = LLM_CONFIG["timeout"]

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

            # 注意：原代码中此方法使用了不同的模型 "46-qwen3.6-35B"
            # 如果需要统一，请使用 LLM_CONFIG["model"]，否则需单独配置
            model_to_use = LLM_CONFIG["model"] 

            response = await client.chat.completions.create(
                model=model_to_use,
                messages=[
                    {"role": "user", "content": query},
                ],
                temperature=LLM_CONFIG["temperature"],
                stream=False,
                extra_body={"chat_template_kwargs": {"enable_thinking": False}}
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用 OpenAI API 时出错: {e}")
            return None
        finally:
            await client.close()

    @staticmethod
    def generate_embedding(text: str) -> Optional[List[float]]:
        try:
            client = OpenAI(
                api_key=EMBEDDING_CONFIG["api_key"],
                base_url=EMBEDDING_CONFIG.get("base_url_v2", EMBEDDING_CONFIG["base_url"].replace("/embeddings", "")),
            )
            response = client.embeddings.create(
                model=EMBEDDING_CONFIG["model"],
                input=text
            )
            return response.data[0].embedding
        except Exception as e:
            print(f"生成embedding失败: {e}")
            return None

async def batch_chunk_process(
    markdown_context: str,
    chunks: List[Dict[str, Any]],
    batch_size: int = 5,
    timeout: int = 30
    ) -> List[Dict[str, Any]]:
        """
        批量处理文档块（异步并发）
        
        Args:
            markdown_context: 完整文档内容
            chunks: 文档块列表，每个块需包含"text"键
            batch_size: 并发请求数 (默认5)
            timeout: 单个请求超时时间(秒)
            
        Returns:
            处理后的文档块列表
        """
        # 初始化参数
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_API_BASE")
        model = os.getenv("OPENAI_MODEL", "qwen3.5-35b")
        
        if not api_key or not base_url:
            raise ValueError("Missing OpenAI API configuration")

        async def process_single_chunk(
            session: aiohttp.ClientSession,
            chunk: Dict[str, Any],
            semaphore: asyncio.Semaphore
        ) -> Dict[str, Any]:
            """处理单个chunk的异步函数"""
            async with semaphore:
                text = chunk.get("text", "").strip()
                if not text:
                    return chunk

                prompt = f"""
                    <document>{markdown_context}</document>
                    <chunk>{text}</chunk>
                    Please provide a succinct context for this chunk within the document.
                    Respond ONLY with the context text, no additional commentary.
                    """
                chinese_prompt = f"""
                                    <文档>
                    {markdown_context}
                    </文档>
                    以下是需要定位到整体文档中的片段
                    <片段>
                    {text}
                    </片段>
                    请简要说明该片段在整体文档中的上下文关系，以提升片段搜索效果。只需给出简洁的上下文说明，无需其他内容。
                    """
                
                try:
                    async with session.post(
                        f"{base_url}/chat/completions",
                        headers={"Authorization": f"Bearer {api_key}"},
                        json={
                            "model": model,
                            "messages": [
                                {"role": "system", "content": "You are a helpful assistant."},
                                {"role": "user", "content": chinese_prompt}
                            ],
                            "temperature": 0.1,
                            "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}
                        },
                        timeout=timeout
                    ) as response:
                        if response.status == 200:
                            data = await response.json()
                            if data.get("choices"):
                                chunk["text"] = data["choices"][0]["message"]["content"].strip()
                        else:
                            print(f"API error for chunk: {await response.text()}")
                except Exception as e:
                    print(f"Error processing chunk: {e}")
                
                return chunk

        # 创建信号量控制并发数
        semaphore = asyncio.Semaphore(batch_size)
        
        # 异步处理所有chunks
        async with aiohttp.ClientSession() as session:
            tasks = [
                process_single_chunk(session, chunk, semaphore)
                for chunk in chunks
            ]
            
            # 使用tqdm显示进度（可选）
            processed_chunks = []
            for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
                processed_chunks.append(await future)
        
        return processed_chunks

def prepare_rag_prompt(query, documents):
    docs = ""
    for i, doc in enumerate(documents):
        docs += f"文本信息{i+1}：{doc}\n"
    return f"""你作为一个文档检索问答专家，请你从以下文本信息中选取有用的内容作为参考，仔细思考后用中文回答用户的问题，如果文本中没有关于问题的答案，则回答不知道:
                            -----
                            {docs}
                            -----
                            用户提问：{query}
                            你的回答："""

def query_rewrite(query):
    prompt = f"""
            你是一个专业的搜索查询优化助手。你的任务是根据用户输入的原始查询，生成更清晰、更易被搜索引擎理解的改写版本。请遵循以下规则：

            1. **保持原意**：改写后的查询必须与用户意图一致，不能曲解原意。
            2. **简洁清晰**：去掉冗余词，优化表达，但不要过度缩写。
            3. **结构化**：如果查询适合拆解，可以适当结构化（如添加隐含的限定条件）。
            4. **自然语言**：确保改写后的查询仍然符合自然语言习惯。
            5. **多语言支持**：如果查询是中文，优先用中文改写；如果是英文，用英文优化。

            请改写以下查询，直接返回优化后的内容，不要额外解释：
            【用户查询开始】
            {query}
            【用户查询结束】
            """
    return OpenaiAPI.open_api_chat(prompt)


def RAG_PROMPT(QUERY,CONTEXT):
    PROMPT = f"""
### Task:
Respond to the user query using the provided context, incorporating inline citations in the format [source_id] **only when the <source_id>   tag is explicitly provided** in the context.

### Guidelines:
- If you don't know the answer, clearly state that.
- If uncertain, ask the user for clarification.
- Respond in the same language as the user's query.
- If the context is unreadable or of poor quality, inform the user and provide the best possible answer.
- If the answer isn't present in the context but you possess the knowledge, explain this to the user and provide the answer using your own understanding.
- When you find corresponding content in the provided sources, try to quote it verbatim as much as possible.
- When quoting from sources, reproduce the original text verbatim—do not omit, paraphrase, or add content.
- **Only include inline citations using [source_id] (e.g., [1], [2]) when a `<source_id>` tag is explicitly provided in the context.**
- If the source file has a .glb extension, directly mark the corresponding content with the source_id in the relevant position.
- Do not cite if the <source_id> or tag is not provided in the context.  
- Do not use XML tags in your response.
- **If none of the provided sources contain relevant information to answer the question, respond only with: "没有找到相关来源" (No relevant sources found). Do not add any other content.
- Ensure citations are concise and directly related to the information provided.

### Example of Citation:
If the user asks about a specific topic and the information is found in "whitepaper.pdf" with a provided <source_id> , the response should include the citation like so:  
* "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
If no <source_id> is present, the response should omit the citation.

### Output:
Provide a clear and direct response to the user's query, including inline citations in the format [source_id] only when the <source_id> tag is present in the context.

<context>
{CONTEXT}
</context>

<user_query>
{QUERY}
</user_query>
"""
    return PROMPT

#     for chunk in response:
#     content = chunk.choices[0].delta.content
#     if content:
#         print(content, end='', flush=True)
# print('\n')
    # return response
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]]:
        embedding_base_url = os.path.join(self.base_url, "embeddings")
        with httpx.Client(timeout=60) as client:
            response = client.post(
                embedding_base_url,
                json={"model": self.model_name, "input": texts},
                headers={"Authorization": f"Bearer {self.api_key}"},
            )
            embedding_list = [item["embedding"] for item in response.json()["data"]]
        return embedding_list