import os
from typing import List, Dict, Any, Optional, Tuple, Union
from io import BytesIO

import numpy as np

from llama_cpp import Llama

from langchain_core.embeddings import Embeddings
from pydantic import BaseModel

import base64
from PIL import Image

import logging

# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

class LlamaCppCustomEmbeddings(Embeddings):
    """
    自定义LlamaCpp嵌入类，绕开langchain_community封装，直接控制底层参数
    适配LangChain的Embeddings接口（需实现embed_query和embed_documents）
    """
    def __init__(
        self,
        model_path: str,
        n_ctx: int = 512,
        n_threads: int = 4,
        n_gpu_layers: int = 0,
        n_seq_max: int = 128,  # 底层支持的批量序列数参数
        n_threads_batch = None,
        flash_attn = False,
        verbose: bool = False
    ):
        try:
            # 直接初始化llama-cpp-python的Llama类，传递所有底层参数
            self.model = Llama(
                model_path=model_path,
                n_ctx=n_ctx,          # 上下文长度（匹配文档长度）
                n_threads=n_threads,  # CPU线程数
                n_gpu_layers=n_gpu_layers,  # GPU层数量（RTX 5090可设20+）
                n_seq_max=n_seq_max,  # 允许批量处理的最大序列数（解决seq_id超限）
                n_threads_batch = n_threads_batch,  
                flash_attn= flash_attn,
                verbose=verbose,
                embedding=True,       # 强制启用嵌入模式（底层关键参数）
                chat_format="none"    # 禁用聊天模板（嵌入无需生成式格式）
            )
        
        # 验证模型是否正常加载（获取向量维度，确保嵌入功能可用）
        # test_emb = self.embed_query("测试嵌入")
        # self.embedding_dim = len(test_emb)
        # logger.info(f"✅ 自定义嵌入模型加载成功！向量维度：{self.embedding_dim}，n_seq_max：{n_seq_max}")
        except Exception as e:
            logger.error(f"❌ llamacpp嵌入模型初始化失败: {str(e)}")
            return None

    def _validate_embedding_support(self) -> None:
        """验证模型是否支持嵌入功能，避免加载非嵌入模型"""
        model_meta = self.model.metadata or {}
        # 嵌入模型通常包含"embedding"关键字或特定架构（如bge、mistral-embed）
        if "embedding" not in model_meta.get("general.name", "").lower() and \
           model_meta.get("general.architecture", "") not in ["bge", "mistral-embed"]:
            logger.warning(
                "当前模型可能不是专用嵌入模型，可能导致嵌入结果异常！"
                "建议使用bge-small-zh-v1.5-GGUF或mistral-embed-v0.1-GGUF"
            )
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed a list of documents using the Llama model.

        Args:
            texts: The list of texts to embed.

        Returns:
            List of embeddings, one for each text.
        """
        # embeddings = self.model.create_embedding(texts)
        embeddings = [self.model.embed(text) for text in texts]
        final_embeddings = []
        # for e in embeddings["data"]:
        #     try:
        #         if isinstance(e["embedding"][0], list):
        #             for data in e["embedding"]:
        #                 final_embeddings.append(list(map(float, data)))
        #         else:
        #             final_embeddings.append(list(map(float, e["embedding"])))
        #     except (IndexError, TypeError):
        #         final_embeddings.append(list(map(float, e["embedding"])))
        for e in embeddings:
            try:
                if isinstance(e[0], list):
                    for data in e:
                        final_embeddings.append(list(map(float, data)))
                else:
                    final_embeddings.append(list(map(float, e)))
            except (IndexError, TypeError):
                final_embeddings.append(list(map(float, e["embedding"])))
        return final_embeddings
    
    def embed_query(self, text: str) -> List[float]:
        """Embed a query using the Llama model.

        Args:
            text: The text to embed.

        Returns:
            Embeddings for the text.
        """
        embedding = self.model.embed(text)
        if embedding and isinstance(embedding, list) and isinstance(embedding[0], list):
            return list(map(float, embedding[0]))
        else:
            return list(map(float, embedding))
        
    def embed_query2(self, text: str) -> list[float]:
        """实现单句查询嵌入（LangChain必选方法）"""
        # 底层llama.embed返回字典，key为"embedding"
        return self.model.embed(text)#["embedding"]

    def embed_documents2(self, texts: list[str]) -> list[list[float]]:
        """实现多文档批量嵌入（LangChain必选方法，优化批量性能）"""
        # 直接调用底层批量嵌入接口，避免循环调用embed_query（提升速度）
        # self.llm.embed(text)["embedding"] 
        return [self.model.embed(text) for text in texts]
    
    def embed_documents1(self, texts: list[str]) -> list[list[float]]:
        """实现多文档批量嵌入（LangChain必选方法，优化批量性能）"""
        # 直接调用底层批量嵌入接口，避免循环调用embed_query（提升速度）
        # self.llm.embed(text)["embedding"] 
        # return [self.model.embed(text) for text in texts]
        return self.model.embed(texts)

    def embed_image(self, image: Union[Image.Image, str],n_ctx_size) -> List[float]:
        """为图像生成嵌入向量"""
        if isinstance(image, Image.Image):
            # 将图像转换为base64编码
            buffer = BytesIO()
            image.save(buffer, format="PNG")
            image_base64 = base64.b64encode(buffer.getvalue()).decode()
        else:
            image_base64 = image  # 假设是base64字符串
            
        # 构造图像嵌入提示
        prompt = f"Embed this image: data:image/png;base64,{image_base64}"
        if len(prompt) > n_ctx_size:
            raise ValueError("图像编码过长，超过模型上下文限制")
            
        return self.model.create_embedding(prompt)["data"][0]["embedding"]
    
    def normalize_embeddings(self, embeddings: list[np.ndarray]) -> list[np.ndarray]:
        """
        嵌入向量归一化（可选但推荐：提升后续检索的相似度计算精度）
        :param embeddings: 原始嵌入向量列表
        :return: 归一化后的嵌入向量列表
        """
        normalized = []
        for emb in embeddings:
            norm = np.linalg.norm(emb)  # 计算向量的 L2 范数
            if norm == 0:
                normalized.append(emb)
            else:
                normalized.append(emb / norm)  # 归一化到单位向量
        return normalized

# class LlamaCppEmbeddings(BaseModel,Embeddings):
class LlamaCppEmbeddings:

    """
    基于llama-cpp-python实现的嵌入工具类
    支持GGUF格式嵌入模型，本地运行，兼容CPU/GPU加速
    """
    def __init__(
        self,
        model_path: str,
        n_ctx: int = 512,
        n_gpu_layers: int = 0,
        n_threads: Optional[int] = None,
        n_batch=512,
        verbose: bool = False,
        embedding: bool = True  # 强制启用嵌入模式（关键参数）
    ):
        """
        初始化参数说明：
        :param model_path: GGUF格式嵌入模型的本地路径（必填）
        :param n_ctx: 模型上下文窗口大小（需 ≤ 模型的context_length，默认512）
        :param n_gpu_layers: 分配到GPU的层数量（0=全CPU，≥1=GPU加速，根据显存调整）
        :param n_threads: CPU推理线程数（默认=CPU核心数//2，平衡性能与开销）
        :param verbose: 是否打印模型加载日志（调试用True，上线用False）
        :param embedding: 必须设为True（启用嵌入模式，否则模型以生成模式运行）
        :param n_batch  # 批量处理大小，加速embed_documents
        """
        # 1. 校验模型路径
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"模型文件不存在：{model_path}")
        if not model_path.endswith(".gguf"):
            logger.warning("模型文件后缀非.gguf，可能不是llama.cpp兼容格式！")

        # 2. 初始化默认参数
        self.n_threads = n_threads or (os.cpu_count() // 2) if os.cpu_count() else 4

        # 3. 加载llama.cpp模型（核心步骤）
        self.model = Llama(
            model_path=model_path,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu_layers,
            n_threads=self.n_threads,
            n_batch=n_batch,  # 批量处理大小，加速embed_documents
            verbose=verbose,
            embedding=embedding # 强制启用嵌入模式，否则无法生成向量
        )

        # 4. 验证模型是否支持嵌入（通过元数据判断）
        # self._validate_embedding_support()

    def _validate_embedding_support(self) -> None:
        """验证模型是否支持嵌入功能，避免加载非嵌入模型"""
        model_meta = self.model.metadata or {}
        # 嵌入模型通常包含"embedding"关键字或特定架构（如bge、mistral-embed）
        if "embedding" not in model_meta.get("general.name", "").lower() and \
           model_meta.get("general.architecture", "") not in ["bge", "mistral-embed"]:
            logger.warning(
                "当前模型可能不是专用嵌入模型，可能导致嵌入结果异常！"
                "建议使用bge-small-zh-v1.5-GGUF或mistral-embed-v0.1-GGUF"
            )

    def embed_document(self, document: str) -> np.ndarray:
        """
        生成单条文本的嵌入向量
        :param document: 输入文本（建议预处理：去除多余空格、换行）
        :return: 384 维嵌入向量（numpy 数组）
        """
        # 预处理文本（避免特殊字符影响嵌入质量）
        document = document.strip().replace("\n", " ").replace("\r", "")
        # 调用 Llama.cpp 生成嵌入（返回 list，需转为 numpy 数组）
        embedding = self.model.create_embedding(document)["data"][0]["embedding"]
        return np.array(embedding, dtype=np.float32)

    def embed_documents(self, documents: list[str]) -> list[np.ndarray]:
        """
        批量生成多条文本的嵌入向量（比单条调用更高效）
        :param documents: 文本列表
        :return: 嵌入向量列表（每个元素为 384 维 numpy 数组）
        """
        return [self.embed_document(document) for document in documents]

    def normalize_embeddings(self, embeddings: list[np.ndarray]) -> list[np.ndarray]:
        """
        嵌入向量归一化（可选但推荐：提升后续检索的相似度计算精度）
        :param embeddings: 原始嵌入向量列表
        :return: 归一化后的嵌入向量列表
        """
        normalized = []
        for emb in embeddings:
            norm = np.linalg.norm(emb)  # 计算向量的 L2 范数
            if norm == 0:
                normalized.append(emb)
            else:
                normalized.append(emb / norm)  # 归一化到单位向量
        return normalized
    
    def embed_query(self, text: str) -> list[float]:
        """处理单个查询文本的嵌入"""
        return self.embed_documents([text])
    
    def embed_image(self, image: Union[Image.Image, str],n_ctx_size) -> List[float]:
        """为图像生成嵌入向量"""
        if isinstance(image, Image.Image):
            # 将图像转换为base64编码
            buffer = BytesIO()
            image.save(buffer, format="PNG")
            image_base64 = base64.b64encode(buffer.getvalue()).decode()
        else:
            image_base64 = image  # 假设是base64字符串
            
        # 构造图像嵌入提示
        prompt = f"Embed this image: data:image/png;base64,{image_base64}"
        if len(prompt) > n_ctx_size:
            raise ValueError("图像编码过长，超过模型上下文限制")
            
        return self.llm.create_embedding(prompt)["data"][0]["embedding"]