import json
import textwrap
import types

from pqai_agent.data_models.agent_configuration import AgentConfiguration
from pqai_agent.toolkit import get_tools
from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool
from pqai_agent.agents.simple_chat_agent import SimpleOpenAICompatibleChatAgent

class SubAgentToolkit(BaseToolkit):
    @staticmethod
    def create_tool_from_agent(agent_config: AgentConfiguration):
        def agent_execution_func(user_input: str) -> str:
            extra_params = json.loads(agent_config.extra_params)

            agent = SimpleOpenAICompatibleChatAgent(
                model=agent_config.execution_model,
                system_prompt=agent_config.system_prompt,
                tools=get_tools(json.loads(agent_config.tools)),
                generate_cfg=extra_params.get('generate_cfg', {}),
                max_run_step=extra_params.get('max_run_step', 20)
            )
            return agent.run(user_input)

        func_doc = f"""
            {agent_config.description}
            Args:
                user_input (str): 用户输入
            Returns:
                str: Agent执行结果
"""
        func_doc = textwrap.dedent(func_doc).strip()

        dynamic_function = types.FunctionType(
            agent_execution_func.__code__,
            globals(),
            name=agent_config.name,
            argdefs=agent_execution_func.__defaults__,
            closure=agent_execution_func.__closure__
        )
        dynamic_function.__doc__ = func_doc

        return FunctionTool(dynamic_function)