import yaml
import os
from smolagents import GradioUI, CodeAgent
from openai import OpenAI  # 我们这里用 OpenAI 作为接口调用的示范

# 获取当前目录路径
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

# 引入你的工具类
from tools.web_search import DuckDuckGoSearchTool as WebSearch
from tools.visit_webpage import VisitWebpageTool as VisitWebpage
from tools.suggest_menu import SimpleTool as SuggestMenu
from tools.catering_service_tool import SimpleTool as CateringServiceTool
from tools.superhero_party_theme_generator import SuperheroPartyThemeTool as SuperheroPartyThemeGenerator
from tools.final_answer import FinalAnswerTool as FinalAnswer

# 获取火山平台配置
VOLC_API_KEY = os.getenv("VOLC_API_KEY")
VOLC_BASE_URL = os.getenv("VOLC_BASE_URL")
VOLC_BOT = os.getenv("VOLC_BOT")

# 创建火山平台客户端（OpenAI兼容方式）
client = OpenAI(api_key=VOLC_API_KEY, base_url=VOLC_BASE_URL)

def _to_openai_messages(messages):
    out = []
    for m in messages:
        if isinstance(m, dict):
            out.append({"role": m.get("role"), "content": m.get("content")})
            continue

        role = getattr(m, "role", None) or "user"
        content = getattr(m, "content", None)

        if content is None:
            # 常见兜底字段
            content = getattr(m, "text", None) or getattr(m, "message", None)

        if content is None:
            content = str(m)

        out.append({"role": role, "content": content})
    return out

class VolcanoModel:
    def _normalize_messages(self, messages):
        allowed = {"system", "user", "assistant", "tool"}
        role_map = {
            # smolagents / 其他框架常见内部角色
            "tool-call": "assistant",
            "tool_call": "assistant",
            "function_call": "assistant",
            "observation": "tool",     # 有些框架用 observation 表示工具输出
        }

        out = []
        for m in messages:
            # m 可能是 ChatMessage 对象，也可能是 dict
            if isinstance(m, dict):
                role = m.get("role", "user")
                content = m.get("content", "")
            else:
                role = getattr(m, "role", "user")
                content = getattr(m, "content", "") or ""

            role = role_map.get(role, role)

            # 兜底：任何不在白名单的 role 都当 assistant（更安全）
            if role not in allowed:
                role = "assistant"

            out.append({"role": role, "content": content})

        return out

    def generate(self, messages, **kwargs):
        norm_messages = self._normalize_messages(messages)

        response = client.chat.completions.create(
            model=VOLC_BOT,
            messages=norm_messages,
            temperature=kwargs.get("temperature", None),
            max_tokens=kwargs.get("max_tokens", None),
        )

        class Result:
            content = response.choices[0].message.content
            token_usage = None  # 或按你之前那样映射 usage

        return Result()

# 读取prompt配置
with open(os.path.join(CURRENT_DIR, "prompts.yaml"), 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

# 初始化 CodeAgent
agent = CodeAgent(
    model=VolcanoModel(),  # 使用火山模型
    tools=[WebSearch(), VisitWebpage(), SuggestMenu(), CateringServiceTool(), SuperheroPartyThemeGenerator()],
    managed_agents=[],
    max_steps=10,
    verbosity_level=2,
    planning_interval=None,
    name=None,
    description=None,
    executor_type='local',
    executor_kwargs={},
    max_print_outputs_length=None,
    prompt_templates=prompt_templates
)

# 启动 Gradio UI11
if __name__ == "__main__":
    GradioUI(agent).launch(share=False, ssr_mode=False)