"""Prompt builder for StableToolBench ReAct evaluation.

Constructs system + user prompts for each P2P condition, with pluggable
tool descriptions and in-context examples.
"""
import json
from typing import Dict, List, Optional, Any


def build_task_description(tool_descriptions):
    """Build the task description that lists available tools (from StableToolBench rapidapi.py)."""
    desc = ('You should use functions to help handle the real time user querys. Remember:\n'
        '1.ALWAYS call "Finish" function at the end of the task. And the final answer '
        'should contain enough information to show to the user,If you can\'t handle the '
        'task, or you find that function calls always fail(the function is not valid now), '
        'use function Finish->give_up_and_restart.\n'
        '2.Do not use origin tool names, use only subfunctions\' names.\n'
        'You have access of the following tools:\n')
    seen = {}
    for std_name, tool_des in tool_descriptions:
        if std_name not in seen: seen[std_name] = tool_des
    for k, (std_name, tool_des) in enumerate(seen.items()):
        striped = (tool_des[:512].replace('\n', '').strip()) if tool_des else "None"
        if not striped: striped = "None"
        desc += f"{k+1}.{std_name}: {striped}\n"
    return desc


def build_system_prompt(task_description):
    """Build the system message for ReAct (FORMAT_INSTRUCTIONS_SYSTEM_FUNCTION)."""
    return ("You are AutoGPT, you can use many tools(functions) to do the following task.\n"
        "First I will give you the task description, and your task start.\n"
        "At each step, you need to give your thought to analyze the status now and "
        "what to do next, with a function call to actually excute your step.\n"
        "After the call, you will get the call result, and you are now in a new state.\n"
        "Then you will analyze your status now, then decide what to do next...\n"
        "After many (Thought-call) pairs, you finally perform the task, then you can "
        "give your finial answer.\nRemember: \n"
        "1.the state change is irreversible, you can't go back to one of the former state, "
        'if you want to restart the task, say "I give up and restart".\n'
        "2.All the thought is short, at most in 5 sentence.\n"
        "3.You can do more then one trys, so if your plan is to continusly try some "
        "conditions, you can do one of the conditions per try.\nLet's Begin!\n"
        f"Task description: {task_description}")


def build_user_prompt(query, examples=None):
    """Build the initial user message, optionally with in-context examples."""
    prompt = ""
    if examples:
        prompt += "Here are some examples of how to use the available tools:\n\n"
        for i, ex in enumerate(examples, 1):
            prompt += f"Example {i}:\nUser: {ex['instruction']}\n"
            prompt += f"Tool call: {json.dumps(ex['fn_call'])}\n"
            prompt += f"Tool result: {str(ex['tool_results'])[:512]}\n"
            prompt += f"Answer: {ex['answer']}\n\n"
        prompt += "Now handle the following real task:\n\n"
    prompt += f"{query}\nBegin!\n"
    return prompt


def build_initial_messages(query, tool_descriptions, examples=None):
    """Build the initial message list for a ReAct conversation."""
    task_desc = build_task_description(tool_descriptions)
    system = build_system_prompt(task_desc)
    user = build_user_prompt(query, examples)
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def get_condition_config(condition, p2p_descriptions=None, p2p_examples=None):
    """Get configuration for a P2P condition.
    
    condition: One of 'baseline', 'p2p_desc', 'p2p_demo', 'p2p_full'
    Returns dict with: use_custom_descriptions, custom_descriptions, use_examples, examples
    """
    configs = {
        "baseline": {"use_custom_descriptions": False, "custom_descriptions": None, "use_examples": False, "examples": None},
        "p2p_desc": {"use_custom_descriptions": True, "custom_descriptions": p2p_descriptions or {}, "use_examples": False, "examples": None},
        "p2p_demo": {"use_custom_descriptions": False, "custom_descriptions": None, "use_examples": True, "examples": p2p_examples or {}},
        "p2p_full": {"use_custom_descriptions": True, "custom_descriptions": p2p_descriptions or {}, "use_examples": True, "examples": p2p_examples or {}},
    }
    if condition not in configs:
        raise ValueError(f"Unknown condition: {condition}. Must be one of {list(configs.keys())}")
    return configs[condition]


def gather_examples_for_query(tool_names, api_name_reflect, functions, all_examples, max_per_tool=1):
    """Gather in-context examples relevant to a specific query's tools."""
    gathered = []
    for func in functions:
        func_name = func["function"]["name"]
        if func_name == "Finish": continue
        if func_name in all_examples:
            gathered.extend(all_examples[func_name][:max_per_tool])
    return gathered