# 链接数据库

import importlib
import yaml
import qdrant_client
from qdrant_client import QdrantClient, models
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core import VectorStoreIndex
from llama_index.core import Document
from loverooma.embedding_model import VolcanoEmbedding

from llama_index.core import PromptTemplate
import json
import os
from loverooma import logger


class EmbeddingPool():
    def __init__(self):
        self.host = os.getenv("host","localhost")
        self.port = int(os.getenv("port",6333))
        self.similarity_cutoff = float(os.getenv("similarity_cutoff",0.5))
        self.QDRANT_COLLECTION_NAME = os.getenv("collection_name","loveroom")
        self.model_name = os.getenv("model_name","doubao-embedding-text-240715")
        self.api_key = os.getenv("api_key",'')
        self.similarity_top_k = int(os.getenv("similarity_top_k",2))
        self.reload()
        

    def create_collection(self,collection_name:str = "loveroom",vector_dimension:str = 1536):
        distance_metric = models.Distance.COSINE # 使用余弦相似度

        # 2. 定义 Collection 参数

        client = qdrant_client.QdrantClient(
            host=self.host,
            port=self.port,
        )

        # 3. 创建 Collection (推荐使用 recreate_collection)
        try:
            client.recreate_collection(
                collection_name=collection_name,
                vectors_config=models.VectorParams(size=vector_dimension, distance=distance_metric),
            )
        except Exception as e:
            raise Exception("创建collection失败") from e
        finally:
            client.close()

    def reload(self):
        logger.info('reload')
        # 默认仓库要被创建的
        self.postprocess = SimilarityPostprocessor(similarity_cutoff=self.similarity_cutoff)
        client = qdrant_client.QdrantClient(
            host=self.host,
            port=self.port,
        )

        # ADD 2025年11月26日13:44:24

        collections = client.get_collections().collections
        existing_collection_names = {c.name for c in collections}
        if self.QDRANT_COLLECTION_NAME not in existing_collection_names:
            client.create_collection(
                collection_name=self.QDRANT_COLLECTION_NAME,
                vectors_config=models.VectorParams(size=2560, distance=models.Distance.COSINE),
            )

        # ADD END

        vector_store = QdrantVectorStore(client=client, collection_name=self.QDRANT_COLLECTION_NAME)
        self.embed_model = VolcanoEmbedding(model_name = self.model_name,
                                            api_key =self.api_key)
        self.index = VectorStoreIndex.from_vector_store(vector_store,embed_model=self.embed_model)
        self.retriver = self.index.as_retriever(similarity_top_k=self.similarity_top_k)

        logger.debug("== reload start ==")
        logger.debug(self.postprocess)


    def update(self,text:str,id:str)->str:
        logger.info('update')
        # splitter
        doc=Document(text = text,id_=id)
        self.index.update(document=doc)

    def delete(self,id:str):
        logger.info('delete')
        self.index.delete(doc_id = id)

    def search(self,query:str)->str:
        logger.info('search')
        result = self.retriver.retrieve(query)
        logger.debug(result)
        result_p = self.postprocess.postprocess_nodes(result)
        return '\n\n'.join([i.text for i in result_p])


import os
from openai import OpenAI


class Desensitization():
    def __init__(self):
        self.client = OpenAI(
            base_url='https://ark.cn-beijing.volces.com/api/v3',
            api_key=os.getenv("api_key",''),
        )
        self.model_name = os.getenv("chat_model_name",'')
        self.desensitization_prompt_ = os.getenv("Desensitization_prompt","")
        self.evaluation_prompt_ = os.getenv("Evaluation_prompt","")
        self.roll_time = int(os.getenv("roll_time",3))

    def _prompt_chat(self,prompt:str):
        response = self.client.chat.completions.create(
            model= self.model_name,
            messages=[
                        {
                            "role": "system",
                            "content": [
                                {"type": "text", "text": prompt },
                            ],
                        }
                    ],
        )
        result = response.choices[0].message.content
        return result
    
    def _postprocess(self,eva_result:str):
        if 'true' in eva_result:
            return True
        else:
            return False

    def desensitization(self,text):
        desensitization_prompt_ = self.desensitization_prompt_
        evaluation_prompt_ = self.evaluation_prompt_
        desensitization_prompt = PromptTemplate(desensitization_prompt_)
        evaluation_prompt = PromptTemplate(evaluation_prompt_)

        advice = ""
        for i in range(self.roll_time):
            logger.info(i)
            des_result = self._prompt_chat(desensitization_prompt.format(text = text,advice = advice))
            eva_result = self._prompt_chat(evaluation_prompt.format(des_result = des_result))
            
            eva_result_json =  json.loads(eva_result)
            status = eva_result_json.get('status')
            review = eva_result_json.get("review")
            if status == "0":
                return "success", des_result
            elif status == "1":
                return "failed", "脱敏失败 " + review
            elif status == "2":
                advice = review
                text = des_result

        return "failed", "脱敏失败 执行异常"

