import logging
import time
from zhipuai import ZhipuAI
from app.utils.prompt_repository import PromptRepository

logger = logging.getLogger(__name__)

class ZhipuKbService:
    def __init__(self):
        self.model_name = "glm-4-plus"
        self.app_secret_key = "d54f764a1d67c17d857bd3983b772016.GRjowY0fyiMNurLc"
        logger.info("ZhipuKbService initialized with model: %s", self.model_name)


    def retrive(self, message, knowledge_id, prompt_template):
        logger.info("Starting retrive call with knowledge_id: %s", knowledge_id)
        start_time = time.time()
        client = ZhipuAI(api_key=self.app_secret_key)
        default_prompt = "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找问题\n\"\"\"\n{{question}}\n\"\"\"\n的答案，找到答案就仅使用文档语句回答问题，找不到答案就用自身知识回答并且告诉用户该信息不是来自文档。\n不要复述问题，直接开始回答。你是一个智能助理，请用该角色的语言风格和对话方式回答问题。"

        if prompt_template is None or prompt_template == "":
            prompt_template = default_prompt
        try:
            response = client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "user", "content": message},
                ],
                tools=[
                    {
                        "type": "retrieval",
                        "retrieval": {
                            "knowledge_id": knowledge_id,
                            "prompt_template": prompt_template
                        }
                    }
                ],
                stream=False,
                max_tokens=4095,
                temperature=0.01,  # default=0.01
                top_p=0.1, #default=0.1
            )
            result = response.choices[0].message.content
            end_time = time.time()
            logger.info("retrive call completed in %.2f seconds", end_time - start_time)
            return result
        except Exception as e:
            logger.error("Error in retrive: %s", str(e))
            raise

    def retrive_sse(self, message, knowledge_id, prompt_template=None,system_prompt=None):
        logger.info("Starting retrive_sse call with knowledge_id: %s, message:%s", knowledge_id, message)
        start_time = time.time()
        client = ZhipuAI(api_key=self.app_secret_key)
        default_prompt = "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找问题\n\"\"\"\n{{question}}\n\"\"\"\n的答案，找到答案就仅使用文档语句回答问题，找不到答案就告诉用户知识库中没有该信息。\n不要复述问题，直接开始回答。"
        messages = [{"role": "user", "content": message}]
        # if system_prompt != None:
        #     messages.append({"role": "system", "content": system_prompt})
        if prompt_template is None or prompt_template == "":
            prompt_template = default_prompt
        try:
            response = client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                tools=[
                    {
                        "type": "retrieval",
                        "retrieval": {
                            "knowledge_id": knowledge_id,
                            "prompt_template": prompt_template
                        }
                    }
                ],
                stream=True,
                max_tokens=4095,
                temperature=0.01,  # default=0.01
                top_p=0.1, #default=0.1
            )
            for chunk in response:
                yield chunk.choices[0].delta.content
            end_time = time.time()
            logger.info("retrive_sse call completed in %.2f seconds", end_time - start_time)
        except Exception as e:
            logger.error("Error in retrive_sse: %s", str(e))
            raise