from datetime import datetime
import json
import logging
import re
from pydantic import BaseModel, Field

from configs.model_config import LLM_MODELS
from server.agent.tools.search_tool import search_tool
from server.chat import utils
from server.chat.policy_fun_iast import get_llm_model_response







def math_count(query: str):
    try:
        matches = re.findall(r'\{.*?\}', query)
        if len(matches)>=2:
            query = matches[0]
            parsed_query = json.loads(query)["query"]
            # 继续使用解析后的查询进行后续操作
            time_based_uuid = json.loads(matches[1])["uuid"]
            # temp = utils.get_shared_variable(time_based_uuid)
            # temp["END"] ="ok"
            # utils.set_shared_variable(time_based_uuid,temp)
            
            tip = {}
            # tip["END"]="ok"
            tip["source_docs"]=[]
            tip["num"]=0
            tip["title"]=[]
                    
            utils.set_shared_variable(time_based_uuid+"q",tip)
            first_json = {
                "query": parsed_query,
                "knowledge_name": [],
                "keywords": []
                }
            second_json = {
                "uuid": time_based_uuid+"q"
                }
            math_doc = search_tool(json.dumps(first_json) + json.dumps(second_json))
            utils.remove_shared_variable(time_based_uuid+"q")
            res = get_llm_model_response(
                    strategy_name="default_math",
                    llm_model_name=LLM_MODELS[3],
                    template_prompt_name="default_math",
                    prompt_param_dict={"input": parsed_query, "math_doc": f"{math_doc}", "time": datetime.now().strftime("%Y%m%d")},
                    temperature=0.01,
                    max_tokens=512
            )
            return f"{res}"
    except Exception as e:
                logging.error(f"Error occurred while processing math query: {e}")
                return "<system>不要再调用该工具了</system>"
                
    except Exception as e:
        return "<system>不要再调用该工具了</system>"

def code_count(query: str):
    try:
        matches = re.findall(r'\{.*?\}', query)
        if len(matches)>=2:
            query = matches[0]
            parsed_query = json.loads(query)["query"]
            # 继续使用解析后的查询进行后续操作
            time_based_uuid = json.loads(matches[1])["uuid"]
            # temp = utils.get_shared_variable(time_based_uuid)
            # temp["END"] ="ok"
            # utils.set_shared_variable(time_based_uuid,temp)
            
            tip = {}
            # tip["END"]="ok"
            tip["source_docs"]=[]
            tip["num"]=0
            tip["title"]=[]
                    
            utils.set_shared_variable(time_based_uuid+"q",tip)
            first_json = {
                "query": parsed_query,
                "knowledge_name": [],
                "keywords": []
                }
            second_json = {
                "uuid": time_based_uuid+"q"
                }
            code_doc = search_tool(json.dumps(first_json) + json.dumps(second_json))
            utils.remove_shared_variable(time_based_uuid+"q")
            res = get_llm_model_response(
                    strategy_name="default_code",
                    llm_model_name=LLM_MODELS[2],
                    template_prompt_name="default_code",
                    prompt_param_dict={"input": parsed_query, "code_doc": f"{code_doc}", "time": datetime.now().strftime("%Y%m%d")},
                    temperature=0.01,
                    max_tokens=512
            )
            res = res.replace("<think>","")
            return f"{res}"
    except Exception as e:
                logging.error(f"Error occurred while processing math query: {e}")
                return "<system>不要再调用该工具了</system>"
                
    except Exception as e:
        return "<system>不要再调用该工具了</system>"




class RagSearchInput(BaseModel):
    query: str = Field(...,description="查询对象")