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

import asyncio
import json
import logging
import re
import os
from typing import AsyncGenerator, Optional

from openai import AsyncOpenAI

from app.agent.tools import TOOL_DEFINITIONS, execute_tool
from app.config import (
    get_llm_model_config,
    get_llm_max_iterations,
    get_image_model_config,
    get_max_recent_turns,
)
from app.memory import get_memory
from app.services.image_gen import to_data_uri

logger = logging.getLogger(__name__)


def _get_client(provider: str) -> AsyncOpenAI:
    """按 provider 创建对应的 OpenAI 兼容客户端。"""
    if provider == "vectorengine":
        return AsyncOpenAI(
            api_key=os.getenv("VECTORENGINE_API_KEY"),
            base_url=os.getenv("VECTORENGINE_BASE_URL", "https://api.vectorengine.ai/v1"),
        )
    if provider == "deepseek":
        return AsyncOpenAI(
            api_key=os.getenv("DEEPSEEK_API_KEY"),
            base_url=os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
        )
    return AsyncOpenAI()

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

## 你的能力
- 根据用户的文字描述生成图片（UI图标、按钮、插画、立绘、概念图等）
- 理解用户的审美意图，将中文描述转化为高质量的英文生成 prompt
- 根据用户反馈迭代修改（调整颜色、风格、构图等）
- 如果用户提供了参考图（支持多张），将参考图的风格元素融入生成 prompt
- 理解多张参考图各自的角色（如"图1的主体 + 图2的风格/视角"），并在 prompt 中准确传达
- 理解用户在图片上的标注（框选区域 + 文字批注），精准定位需要修改的部分
- 将一张图片转换为多个不同视角（使用 transform_view 工具），适合建筑、物体等需要从不同角度查看的场景

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

## 标注理解
当用户发送带有「图片标注」的消息时，表示用户在之前生成的图片上做了标注。
标注格式为：`[区域 (x%, y%) 大小 w%×h%]: 修改意见`
- 区域坐标表示标注框在图片上的相对位置
- 你需要理解标注区域所指的图片内容，并将修改意见融入新的 prompt
- 如果附带了标注截图（参考图），仔细观察红色标注框和文字来理解用户意图
- 在调整 prompt 时，保持原图整体风格不变，只针对标注区域做修改

## 生成 prompt 要求
- 必须使用英文
- 尽量详细描述：主体内容、颜色方案、光照、构图、材质等
- 如果用户要求游戏 UI 元素，添加相关关键词如 "game UI", "icon", "button" 等
- 如果你能直接看到参考图（图片内容），可以在 prompt 中描述参考图的风格特征。多张参考图时，理解用户对各图的定位（如"图1做主体参考、图2做风格参考"），将相应特征分别融入 prompt
- 如果你无法看到参考图（只收到了文字提示说有参考图），参考图会由生图工具自动处理风格融合。此时你不要自行猜测画风/艺术风格关键词（如 pixel art、watercolor、oil painting 等），把风格交给参考图来决定。但如果用户在消息中明确指定了风格（如"赛博朋克风"、"水彩风"等），应保留并翻译到 prompt 中——尊重用户的主动意图

## 注意事项
- 用中文和用户交流
- 生成图片后简要说明你使用的 prompt 思路
- 主动建议迭代方向
- **禁止模拟工具调用**：生成图片时必须实际调用 generate_image 工具，绝不能用文字描述"已生成"或假装工具已执行。如果需要生成多张图片，每张都必须单独调用工具
- **禁止在回复中嵌入图片链接**：不要在回复文字中使用 Markdown 图片语法（如 `![](...)` 或 `sandbox:` 链接）。图片展示由系统自动处理，你只需用文字描述结果即可
"""


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

    事件类型：
    - text_delta: LLM 文字增量
    - tool_start: 开始执行工具
    - image_result: 图片生成结果
    - memory_warning: 记忆存储异常（不中断对话）
    - done: 完成
    - error: 错误
    """
    # ── 滑动窗口：截断过早的消息 ──
    max_turns = get_max_recent_turns()
    if len(messages) > max_turns:
        messages = messages[-max_turns:]

    # ── 记忆检索：用最新 user 消息语义检索相关记忆 ──
    memory_block = ""
    try:
        memory = get_memory()
        last_user_content = ""
        for m in reversed(messages):
            if m["role"] == "user":
                last_user_content = m["content"]
                break

        if last_user_content:
            search_kwargs = {"query": last_user_content, "user_id": user_id, "limit": 10}
            if session_id:
                search_kwargs["run_id"] = session_id
            relevant = memory.search(**search_kwargs)

            results = relevant.get("results", []) if isinstance(relevant, dict) else relevant
            if results:
                items = "\n".join(f"- {m['memory']}" for m in results if m.get("memory"))
                if items:
                    memory_block = f"\n\n## 用户记忆（来自历史对话）\n{items}"
    except Exception as e:
        logger.error("Mem0 记忆检索失败: %s", e, exc_info=True)
        yield {"type": "error", "data": {"message": f"记忆系统检索失败: {e}"}}
        return

    # ── 解析 LLM 模型配置 ──
    llm_config = get_llm_model_config(llm_model)
    llm_provider = llm_config["provider"]
    llm_model_id = llm_config["model_id"]
    vision_capable = llm_config.get("vision", False)

    # ── 构建 system prompt ──
    img_model_config = get_image_model_config(image_model)
    current_model_name = img_model_config.get('name', '未知')
    current_model_id = img_model_config.get('id', '未知')
    model_hint = (
        f"\n\n## 当前生图模型（重要）\n"
        f"本次对话用户选择的生图模型是 **{current_model_name}**"
        f"（model_id: {current_model_id}）。\n"
        f"**注意**：对话历史中可能包含之前使用其他模型的记录，忽略那些旧模型名。"
        f"本次生成使用的是 {current_model_name}，在回复中只能使用这个名称。"
    )
    api_messages = [{"role": "system", "content": SYSTEM_PROMPT + model_hint + memory_block}]

    effective_refs = ref_image_urls or []

    for msg in messages:
        if msg["role"] == "user" and effective_refs and msg is messages[-1]:
            if vision_capable:
                content_parts: list[dict] = [{"type": "text", "text": msg["content"]}]
                for ref_url in effective_refs:
                    content_parts.append({
                        "type": "image_url",
                        "image_url": {"url": to_data_uri(ref_url)},
                    })
                api_messages.append({"role": "user", "content": content_parts})
            else:
                n_refs = len(effective_refs)
                hint = (
                    f"{msg['content']}\n\n"
                    f"【系统提示：用户上传了 {n_refs} 张参考图，已自动传递给图片生成工具。"
                    "生图工具会根据模型能力自动处理参考图的风格融合。"
                    "你无法看到这些参考图，因此在生成 prompt 时：\n"
                    "1. 描述画面内容（主体、构图、光照、材质等）\n"
                    "2. 不要自行猜测画风/艺术风格——但如果用户明确指定了风格，保留到 prompt 中\n"
                    "3. 用户未指定风格时，风格完全由参考图决定】"
                )
                api_messages.append({"role": "user", "content": hint})
            continue
        api_messages.append({"role": msg["role"], "content": msg["content"]})

    client = _get_client(llm_provider)

    for _ in range(get_llm_max_iterations()):
        try:
            response = await client.chat.completions.create(
                model=llm_model_id,
                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

        # 如果没有工具调用：检测 LLM 是否在用文字模拟生图
        if not tool_calls_data:
            if _looks_like_fake_generation(collected_text):
                # LLM 用文字假装调用了工具，丢弃这段文字，注入纠正消息强制重试
                yield {
                    "type": "text_delta",
                    "data": {"text": "\n\n[系统：检测到未调用生图工具，正在重试...]\n"},
                }
                api_messages.append({"role": "assistant", "content": collected_text})
                api_messages.append({
                    "role": "user",
                    "content": (
                        "你刚才没有调用 generate_image 工具，只是用文字描述了生成过程。"
                        "请立即调用 generate_image 工具来实际生成图片。"
                        "不要解释，直接调用工具。"
                        f"当前使用的生图模型是 {img_model_config.get('name', '未知')}。"
                    ),
                })
                continue
            # 对话正常结束，异步存储记忆
            async for evt in _store_and_done(messages, session_id, user_id):
                yield evt
            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"]

            tool_message = (
                "正在进行视角变换..." if tool_name == "transform_view"
                else "正在生成图片..."
            )
            yield {
                "type": "tool_start",
                "data": {"tool": tool_name, "message": tool_message},
            }

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

            result = await execute_tool(tool_name, arguments, effective_refs or None, image_model)

            used_model = result.get("model_name", "")

            if result.get("errors"):
                yield {
                    "type": "tool_error",
                    "data": {
                        "tool": tool_name,
                        "errors": result["errors"],
                        "model_name": used_model,
                    },
                }

            if result.get("error") and not result.get("success"):
                yield {
                    "type": "tool_error",
                    "data": {
                        "tool": tool_name,
                        "errors": [result["error"]],
                        "model_name": used_model,
                    },
                }

            if tool_name == "transform_view" and result.get("images"):
                # transform_view 返回带视角信息的图片列表
                image_urls = [v["url"] for v in result["images"]]
                yield {
                    "type": "image_result",
                    "data": {
                        "images": image_urls,
                        "prompt_used": "",
                        "model_name": used_model,
                        "view_info": result["images"],
                    },
                }
            elif result.get("images"):
                yield {
                    "type": "image_result",
                    "data": {
                        "images": result["images"],
                        "prompt_used": result.get("prompt_used", ""),
                        "model_name": used_model,
                    },
                }

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

        # ref_image_url 不清除：消息构建只在循环外执行一次，
        # 后续工具调用仍需参考图（InstantStyle 等模型必须有 style_image）

    # 迭代次数用尽，存储记忆后结束
    async for evt in _store_and_done(messages, session_id, user_id):
        yield evt


# ─── 异步记忆存储 ─────────────────────────────────────────


async def _store_and_done(
    messages: list[dict], session_id: Optional[str], user_id: str = "default_user"
) -> AsyncGenerator[dict, None]:
    """触发 Mem0 异步存储后 yield done。存储失败时 yield warning 但不中断。"""
    try:
        memory = get_memory()
        recent = messages[-4:] if len(messages) >= 4 else messages
        add_kwargs: dict = {"user_id": user_id}
        if session_id:
            add_kwargs["run_id"] = session_id

        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, lambda: memory.add(recent, **add_kwargs))
    except Exception as e:
        logger.error("Mem0 记忆存储失败: %s", e, exc_info=True)
        yield {
            "type": "memory_warning",
            "data": {"message": f"记忆存储失败: {e}"},
        }

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


# ─── 假生成检测 ──────────────────────────────────────────

_FAKE_GEN_PATTERNS = re.compile(
    r"已生成|生成完成|开始生成|正在生成|图片已|prompt.*?设计思路",
    re.IGNORECASE,
)


def _looks_like_fake_generation(text: str) -> bool:
    """判断 LLM 的纯文字回复是否在假装已经调用了生图工具。"""
    if not text or len(text) < 50:
        return False
    matches = _FAKE_GEN_PATTERNS.findall(text)
    return len(matches) >= 2