
from typing import List, Any, Tuple, Dict, Optional
import json
import uuid
import re
from turbo_agent_core.schema.states import Conversation, Message, MessageRole
from loguru import logger

def parameters_to_prompt(parameters: List[Any], depth=1) -> str:
    if not parameters:
        return "无"
    indent = "  " * depth
    param_descs = []
    for param in parameters:
        # 处理 Pydantic 模型或字典
        p_data = param.model_dump() if hasattr(param, "model_dump") else param
        
        name = p_data.get("name", "")
        p_type = p_data.get("type", "string")
        required = "必填" if p_data.get("required", False) else "非必填"
        desc = p_data.get("description", "")
        
        param_desc = f"{indent}{name}: {p_type}, （{required}。{desc}。"
        
        if p_type == "object":
            sub_params = p_data.get("parameters", [])
            param_desc += "该参数为对象类型，其属性及含义如下：\n " + parameters_to_prompt(sub_params, depth + 1)
        elif p_type == "array":
            type_ref = p_data.get("type_ref", "")
            if type_ref == "object":
                sub_params = p_data.get("parameters", [])
                param_desc += "该参数为数组类型，其元素包含如下字段：\n " + parameters_to_prompt(sub_params, depth + 1)
            elif type_ref == "enum":
                enum_vals = p_data.get("enum_values", [])
                param_desc += f"该参数为数组类型，其元素为枚举类型，具体可选枚举值如下：\n [{', '.join(enum_vals)}]"
            else:
                param_desc += f"该参数为数组类型，其元素均为\"{type_ref}\"类型。"
        elif p_type == "enum":
            enum_vals = p_data.get("enum_values", [])
            param_desc += f"该参数为枚举类型，其枚举值如下：\n [{', '.join(enum_vals)}]"
        else:
            default = p_data.get("default")
            if default:
                param_desc += f"默认值为：{default}"

        param_desc += "）"
        param_descs.append(param_desc)
    return "\n".join(param_descs)

def construct_json_correction_conversation(tool_id: str, output_schema: Optional[Dict], bad_json: str, error_msg: str) -> Tuple[Conversation, str]:
    schema_str = json.dumps(output_schema, ensure_ascii=False) if output_schema else "No schema provided"
    
    prompt = f"""
以下是预期的输出schema：
{schema_str}

以下错误的json格式数据：
{bad_json}

格式错误原因是：{error_msg}

请基于预期的输出schema和格式错误的json数据，生成符合预期的JSON格式数据。请注意，输出的JSON必须符合预期的输出schema，并且不能包含任何额外的注释或说明。
"""
    
    messages = []
    messages.append(Message(id=str(uuid.uuid4()), role=MessageRole.system, content="你是一个JSON格式校验和修正助手，请帮助用户修正JSON格式错误"))
    msg_input = Message(id=str(uuid.uuid4()), role=MessageRole.user, content=prompt)
    messages.append(msg_input)
    
    conversation = Conversation(
        id=str(uuid.uuid4()),
        messages=messages
    )
    return conversation, msg_input.id

def construct_tool_messages(
    tool_id: str,
    tool_name: str,
    input_params: List[Any],
    output_params: List[Any],
    setting: Any,
    input_data: Dict[str, Any]
) -> Tuple[Conversation, str]:
    # 1. 准备模板变量
    input_name_list = [p.get("name") if isinstance(p, dict) else p.name for p in input_params]
    input_name_list_str = ",".join(input_name_list)
    
    input_desc = ""
    if input_params:
        input_desc = "我提供给你的输入参数含义如下：\n" + parameters_to_prompt(input_params, depth=4)
        
    output_desc = parameters_to_prompt(output_params, depth=4)
    
    business_target = getattr(setting, "target", "") if setting else ""
    business_principles = getattr(setting, "principles", []) if setting else []
    
    rules_prompt = ""
    if business_principles:
        # 如果可用，按索引排序规则
        sorted_principles = sorted(business_principles, key=lambda x: x.index) if hasattr(business_principles[0], "index") else business_principles  
        rules_str = "\n".join([f"规则{r.index}: {r.content}" for r in sorted_principles])
        rules_prompt = f"""
            在完成以本任务的过程中同时也一定要注意遵守以下行为规则：
            {rules_str}
            """

    # 2. 构建模板
    template = f"""
你好，希望你一切都好。我正在寻求你的帮助，想要解决一个特定的功能。我知道你有处理信息和执行各种任务的能力，这是基于提供的指示。为了帮助你更容易地理解我的请求，我将使用一个模板来描述函数、输入数据schema和对输入的处理方法,其中输入输出的格式与含义使用json schema表示。请在下面找到详细信息：
        function_name：[{tool_name}]
        input：[{input_name_list_str}]
        rule：[
            函数输入格式是:
            {tool_name}({input_name_list_str})。
            {input_desc}

                我需要你作为一个功能函数的任务核心目标描述如下：
            {business_target}
            {rules_prompt}

            请以JSON形式输出最终结果，不需要要提供任何额外注释和说明，具体输出结果的应该包含的字段名称、类型要求及含义请遵循以下结构:
            {output_desc}
        ]
-----------------------------------------------------------------
我恳请你根据我提供的细节为这个函数提供输出。非常感谢你的帮助！谢谢！
我将使用方括号内的相关信息替换函数所需执行的内容。这个详细的介绍应该能够帮助你更高效地理解我的请求并提供所需的输出。
以上是对需求的全部描述，如果你理解了，请仅用一个词"OK"回答。
"""
    # 3. 构造输入字符串
    # input_data 是字典。我们需要将其格式化为 func(val1, val2...)
    param_strs = []
    logger.debug(f"Constructing tool input for tool '{tool_name}' with input data: {input_data}, using parameters: {input_name_list}")
    for name in input_name_list:
        val = input_data.get(name)
        param_strs.append(json.dumps(val, ensure_ascii=False))
    input_str = f"{tool_name}({','.join(param_strs)})"
    conversation = Conversation(
        id=str(uuid.uuid4()),
    )
    
    pre = conversation.add_child(
        Message(id=str(uuid.uuid4()), role=MessageRole.system, content="你是一个乐于助人的助手，请尽全力帮助用户完成他需要你帮助他完成的功能")
    )
    pre = conversation.add_child(Message(id=str(uuid.uuid4()), role=MessageRole.user, content=template),parent_id=pre.id)
    pre = conversation.add_child(Message(id=str(uuid.uuid4()), role=MessageRole.assistant, content="OK"), parent_id=pre.id)
    
    msg_input = Message(id=str(uuid.uuid4()), role=MessageRole.user, content=input_str, parent_id=pre.id)
    pre = conversation.add_child(msg_input, parent_id=pre.id)
    return conversation, pre.id

def parse_json_from_llm_output(result_text: str) -> Dict:
    # 移除 <think> 标签
    clean_text = re.sub(r'<think>.*?</think>', '', result_text, flags=re.DOTALL)
    
    # 提取 JSON 块
    begin = clean_text.find("```json")
    end = clean_text.rfind("```")
    
    json_str = clean_text
    if begin >= 0:
        json_str = clean_text[begin+7:end]
    
    try:
        # 尝试解析
        return json.loads(json_str, strict=False)
    except Exception as e:
        raise ValueError(f"Failed to parse JSON: {e}")
