"""Agent Loop 主循环：LLM 对话 -> 工具调用 -> 结果回传 -> 继续。"""

import json
from typing import AsyncGenerator, Optional

from openai import AsyncOpenAI

from app.agent.tools import TOOL_DEFINITIONS, execute_tool


def _get_client() -> AsyncOpenAI:
    return AsyncOpenAI()

SYSTEM_PROMPT = """\
你是一个专业的游戏美术 AI 助手。你的工作是帮助美术人员通过对话生成游戏美术资源。

## 你的能力
- 根据用户的文字描述生成图片（UI图标、按钮、插画、立绘、概念图等）
- 理解用户的审美意图，将中文描述转化为高质量的英文生成 prompt
- 根据用户反馈迭代修改（调整颜色、风格、构图等）
- 如果用户提供了参考图，将参考图的风格元素融入生成 prompt

## 工作流程
1. 理解用户需求，必要时追问细节（尺寸、风格、用途等）
2. 将需求转化为详细的英文 prompt，调用 generate_image 工具生成图片
3. 向用户展示结果并询问反馈
4. 根据反馈调整 prompt 并重新生成

## 生成 prompt 要求
- 必须使用英文
- 尽量详细描述：主体内容、风格、颜色方案、光照、构图、材质等
- 如果用户要求游戏 UI 元素，添加相关关键词如 "game UI", "icon", "button" 等
- 如果用户提供了参考图，在 prompt 中描述参考图的风格特征

## 注意事项
- 用中文和用户交流
- 生成图片后简要说明你使用的 prompt 思路
- 主动建议迭代方向
"""


async def run_agent_loop(
    messages: list[dict],
    ref_image_url: Optional[str] = None,
) -> AsyncGenerator[dict, None]:
    """
    运行 Agent Loop，以 SSE 事件流形式 yield 结果。

    事件类型：
    - text_delta: LLM 文字增量
    - tool_start: 开始执行工具
    - image_result: 图片生成结果
    - done: 完成
    - error: 错误
    """
    # 构建消息列表
    api_messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for msg in messages:
        if msg["role"] == "user" and ref_image_url:
            # 最后一条用户消息附加参考图（GPT-4o vision）
            if msg == messages[-1] or (
                msg.get("role") == "user"
                and messages.index(msg) == len(messages) - 1
            ):
                api_messages.append({
                    "role": "user",
                    "content": [
                        {"type": "text", "text": msg["content"]},
                        {
                            "type": "image_url",
                            "image_url": {"url": ref_image_url},
                        },
                    ],
                })
                continue
        api_messages.append({"role": msg["role"], "content": msg["content"]})

    client = _get_client()

    max_iterations = 5
    for _ in range(max_iterations):
        try:
            response = await client.chat.completions.create(
                model="gpt-4o-mini",
                messages=api_messages,
                tools=TOOL_DEFINITIONS,
                stream=True,
            )
        except Exception as e:
            yield {"type": "error", "data": {"message": str(e)}}
            return

        collected_text = ""
        tool_calls_data: dict[int, dict] = {}

        async for chunk in response:
            delta = chunk.choices[0].delta if chunk.choices else None
            if not delta:
                continue

            # 文字内容
            if delta.content:
                collected_text += delta.content
                yield {"type": "text_delta", "data": {"text": delta.content}}

            # 工具调用（流式累积）
            if delta.tool_calls:
                for tc in delta.tool_calls:
                    idx = tc.index
                    if idx not in tool_calls_data:
                        tool_calls_data[idx] = {
                            "id": "",
                            "name": "",
                            "arguments": "",
                        }
                    if tc.id:
                        tool_calls_data[idx]["id"] = tc.id
                    if tc.function and tc.function.name:
                        tool_calls_data[idx]["name"] = tc.function.name
                    if tc.function and tc.function.arguments:
                        tool_calls_data[idx]["arguments"] += tc.function.arguments

        finish_reason = chunk.choices[0].finish_reason if chunk.choices else None

        # 如果没有工具调用，对话结束
        if not tool_calls_data:
            yield {"type": "done", "data": {}}
            return

        # 将 assistant 消息（含 tool_calls）加入历史
        assistant_msg: dict = {"role": "assistant"}
        if collected_text:
            assistant_msg["content"] = collected_text
        else:
            assistant_msg["content"] = None

        assistant_msg["tool_calls"] = []
        for idx in sorted(tool_calls_data.keys()):
            tc_data = tool_calls_data[idx]
            assistant_msg["tool_calls"].append({
                "id": tc_data["id"],
                "type": "function",
                "function": {
                    "name": tc_data["name"],
                    "arguments": tc_data["arguments"],
                },
            })
        api_messages.append(assistant_msg)

        # 依次执行每个工具调用
        for idx in sorted(tool_calls_data.keys()):
            tc_data = tool_calls_data[idx]
            tool_name = tc_data["name"]

            yield {
                "type": "tool_start",
                "data": {"tool": tool_name, "message": "正在生成图片..."},
            }

            try:
                arguments = json.loads(tc_data["arguments"])
            except json.JSONDecodeError:
                arguments = {}

            result = await execute_tool(tool_name, arguments, ref_image_url)

            # 如果有图片结果，推送给前端
            if result.get("images"):
                yield {
                    "type": "image_result",
                    "data": {"images": result["images"]},
                }

            # 工具结果回传给 LLM
            api_messages.append({
                "role": "tool",
                "tool_call_id": tc_data["id"],
                "content": json.dumps(result, ensure_ascii=False),
            })

        # 清除 ref_image_url，避免后续轮次重复附加
        ref_image_url = None

    yield {"type": "done", "data": {}}