import os
import asyncio

from agentscope.tool import Toolkit
from agentscope.tool import execute_shell_command
from agentscope.tool import view_text_file
from agentscope.tool import write_text_file
from agentscope.tool import insert_text_file
from agentscope.tool import dashscope_text_to_image
from agentscope.tool import dashscope_text_to_audio
from agentscope.tool import dashscope_image_to_text
from agentscope.tool import openai_text_to_image
from agentscope.tool import openai_text_to_audio
from agentscope.tool import openai_edit_image
from agentscope.tool import openai_create_image_variation
from agentscope.tool import openai_image_to_text
from agentscope.tool import openai_audio_to_text
from agentscope.tool import execute_python_code
from agentscope.mcp import HttpStatelessClient

toolkit = Toolkit()


def _register_tools():
    {{range .AvailableTools}}
    toolkit.register_tool_function({{.}})
    {{else}}
    pass
    {{end}}


def init_toolkit_sync():
    _register_tools()
    asyncio.run(register_all_MCP(toolkit))


async def init_toolkit_async():
    _register_tools()
    await register_all_MCP(toolkit)


async def register_single_MCP(toolkit: Toolkit, mcp_config):
    """注册单个MCP服务器"""
    headers = mcp_config.get("Headers") or None

    api_client = HttpStatelessClient(
        name=mcp_config["Name"],
        transport=mcp_config["Transport"],
        url=mcp_config["URL"],
        headers=headers,
    )

    await toolkit.register_mcp_client(api_client)


async def register_all_MCP(toolkit: Toolkit):
    """注册所有配置的MCP服务器"""
    {{- range .MCPServers }}
    await register_single_MCP(toolkit, {
        "Name": "{{ .Name }}",
        "URL": "{{ .URL }}",
        "Transport": "{{ .Transport }}",
        "Headers": {
            {{- range $key, $value := .Headers }}
            "{{ $key }}": "{{ $value }}",
            {{- end }}
        }
    })
    {{- end }}