"""Standalone Agent Engine — 执法助手单机版核心引擎"""

import json
import os
import re
import logging
import uuid
from dataclasses import dataclass, field
from typing import Any, Generator, Optional

from openai import AsyncOpenAI
from app.agent.messages import HumanMessage, AIMessage, ToolMessage, SystemMessage

from .config import SCENE_MODEL_MAP, AGENT_SOULS, SCENE_AGENT_MAP, LIGHTWEIGHT_MODEL_CONFIG
from .prompts import build_full_prompt

logger = logging.getLogger(__name__)

# 全局引擎缓存，供 ai_config API 清除
_engine_cache = {}


def clear_engine_cache():
    """清除引擎缓存（配置更新时调用）"""
    global _engine_cache
    _engine_cache.clear()
    logger.info("Engine cache cleared due to AI config update")


@dataclass
class StreamEvent:
    """流式事件（与原 DeerFlow SSE 协议兼容）"""
    type: str
    data: dict[str, Any] = field(default_factory=dict)


def _msg_to_dict(msg) -> dict:
    """Convert custom message object to dict for OpenAI API"""
    if isinstance(msg, HumanMessage):
        return {"role": "user", "content": msg.content}
    elif isinstance(msg, AIMessage):
        d = {"role": "assistant", "content": msg.content or ""}
        if msg.tool_calls:
            d["tool_calls"] = [
                {
                    "id": tc.get("id", ""),
                    "type": "function",
                    "function": {"name": tc["name"], "arguments": json.dumps(tc.get("args", {}), ensure_ascii=False)}
                }
                for tc in msg.tool_calls
            ]
        return d
    elif isinstance(msg, ToolMessage):
        return {"role": "tool", "tool_call_id": msg.tool_call_id, "content": msg.content}
    elif isinstance(msg, SystemMessage):
        return {"role": "system", "content": msg.content}
    else:
        return {"role": "user", "content": str(msg)}


class StandaloneAgentEngine:
    """执法助手单机版 Agent 引擎 — 替代 DeerFlow"""

    def __init__(self, base_url: str = None,
                 default_model: str = None,
                 api_key: str = None,
                 enable_subagent: bool = True,
                 enable_summary: bool = True):
        # 从 Settings 读取默认值
        from .config import LLM_API_BASE, LLM_MODEL, LLM_API_KEY
        self.base_url = base_url or LLM_API_BASE
        self.default_model = default_model or LLM_MODEL
        self.api_key = api_key or LLM_API_KEY
        self.enable_subagent = enable_subagent
        self.enable_summary = enable_summary

        # Memory-level checkpointer for multi-turn conversations
        self.checkpointer = {}  # {thread_id: [messages]}

        # Agent cache {cache_key: agent_config}
        self._agent_cache = _engine_cache

        # Middleware instances (will be initialized later)
        self._title_gen = None
        self._summarizer = None
        self._loop_detector = None
        self._memory_injector = None
        self._subagent_engine = None

    def _create_llm(self, scene_key: str = "general_chat") -> AsyncOpenAI:
        """Create LLM instance based on scene"""
        cfg = SCENE_MODEL_MAP.get(scene_key, SCENE_MODEL_MAP.get("general_chat", {}))
        return AsyncOpenAI(
            base_url=cfg.get("base_url", self.base_url),
            api_key=cfg.get("api_key", ""),
        )

    def _build_system_prompt(self, agent_name: str, subagent_enabled: bool = False,
                              memory_context: str = "") -> str:
        """Build complete system prompt"""
        soul = AGENT_SOULS.get(agent_name, "")
        return build_full_prompt(
            agent_name=agent_name,
            soul=soul,
            memory_context=memory_context,
            subagent_enabled=subagent_enabled,
        )

    def _get_agent_config(self, agent_name: str, scene_key: str, tools: list = None):
        """Create or get cached Agent config"""
        cache_key = f"{agent_name}:{scene_key}"
        if cache_key in self._agent_cache:
            return self._agent_cache[cache_key]

        if tools is None:
            from .tools import get_all_tools
            tools = get_all_tools(include_task=self.enable_subagent)

        prompt = self._build_system_prompt(agent_name, self.enable_subagent)
        cfg = SCENE_MODEL_MAP.get(scene_key, SCENE_MODEL_MAP.get("general_chat", {}))

        config = {
            "model": cfg.get("model", self.default_model),
            "tools": tools,
            "prompt": prompt,
            "temperature": cfg.get("temperature", 0.7),
        }
        self._agent_cache[cache_key] = config
        logger.info(f"Agent config created: agent={agent_name}, scene={scene_key}")
        return config

    async def chat(self, message: str, thread_id: str = None,
             agent_name: str = "general-assistant",
             scene_key: str = "general_chat") -> str:
        """Non-streaming chat"""
        last_text = ""
        async for event in self.stream(message, thread_id, agent_name, scene_key):
            if event.type == "messages-tuple" and event.data.get("type") == "ai":
                content = event.data.get("content", "")
                if content:
                    last_text = content
        return last_text

    async def stream(self, message: str, thread_id: str = None,
               agent_name: str = "general-assistant",
               scene_key: str = "general_chat"):
        """Streaming chat (SSE event generator)"""
        from app.agent.tools.decorator import get_tool_schema

        if thread_id is None:
            thread_id = str(uuid.uuid4())

        # Resolve agent name from scene key
        if scene_key in SCENE_AGENT_MAP:
            agent_name = SCENE_AGENT_MAP[scene_key]

        config = self._get_agent_config(agent_name, scene_key)
        llm = self._create_llm(scene_key)
        tools = config["tools"]
        prompt = config["prompt"]
        tool_schemas = [get_tool_schema(t) for t in tools] if tools else None

        # Get or create conversation history (stored as custom message objects)
        key = f"thread:{thread_id}"
        if key not in self.checkpointer:
            self.checkpointer[key] = [SystemMessage(content=prompt)]
        messages = self.checkpointer[key]
        messages.append(HumanMessage(content=message))

        seen_ids = set()

        try:
            for turn in range(50):  # max turns
                # Convert messages to dicts for API call
                api_messages = [_msg_to_dict(m) for m in messages]

                # === Middleware: Summary compression ===
                if self.enable_summary and self._summarizer:
                    state = {"messages": messages}
                    state = self._summarizer.summarize(state)
                    messages = state["messages"]
                    api_messages = [_msg_to_dict(m) for m in messages]

                response = await llm.chat.completions.create(
                    model=config["model"],
                    messages=api_messages,
                    tools=tool_schemas,
                    temperature=config["temperature"],
                    stream=True,
                )

                # Collect streaming response
                full_content = ""
                tool_calls_data = []
                async for chunk in response:
                    delta = chunk.choices[0].delta if chunk.choices else None
                    if delta is None:
                        continue
                    if delta.content:
                        full_content += delta.content
                    if delta.tool_calls:
                        for tc in delta.tool_calls:
                            # Build tool calls from streaming chunks
                            while len(tool_calls_data) <= tc.index:
                                tool_calls_data.append({"id": "", "function": {"name": "", "arguments": ""}})
                            if tc.id:
                                tool_calls_data[tc.index]["id"] = tc.id
                            if tc.function:
                                if tc.function.name:
                                    tool_calls_data[tc.index]["function"]["name"] = tc.function.name
                                if tc.function.arguments:
                                    tool_calls_data[tc.index]["function"]["arguments"] += tc.function.arguments

                if tool_calls_data:
                    # AI wants to call tools
                    parsed_tool_calls = []
                    for tc in tool_calls_data:
                        try:
                            args = json.loads(tc["function"]["arguments"])
                        except (json.JSONDecodeError, KeyError):
                            args = {}
                        parsed_tool_calls.append({
                            "name": tc["function"]["name"],
                            "args": args,
                            "id": tc["id"],
                        })

                    ai_msg = AIMessage(
                        content=full_content,
                        id=str(uuid.uuid4()),
                        tool_calls=parsed_tool_calls,
                    )

                    # === Middleware: Loop detection ===
                    if self._loop_detector:
                        state = {"messages": messages + [ai_msg]}
                        loop_result = self._loop_detector.check(state, thread_id)
                        if loop_result:
                            prev_len = len(messages)
                            messages = loop_result.get("messages", state["messages"])
                            # Yield any new messages added by loop detector (e.g., warning)
                            for lm in messages[prev_len:]:
                                if isinstance(lm, AIMessage):
                                    yield StreamEvent(type="messages-tuple", data={
                                        "type": "ai", "content": self._extract_text(lm.content),
                                        "id": lm.id,
                                    })
                            # Skip tool execution for loop-detected messages
                            self.checkpointer[key] = messages
                            continue

                    yield StreamEvent(type="messages-tuple", data={
                        "type": "ai", "content": full_content,
                        "id": ai_msg.id,
                        "tool_calls": parsed_tool_calls,
                    })
                    messages.append(ai_msg)

                    # Execute tools
                    tool_map = {t.name: t for t in tools}
                    for tc in tool_calls_data:
                        tool_name = tc["function"]["name"]
                        try:
                            tool_args = json.loads(tc["function"]["arguments"])
                        except (json.JSONDecodeError, KeyError):
                            tool_args = {}

                        tool_func = tool_map.get(tool_name)
                        if tool_func:
                            try:
                                # Check if the tool is async
                                import inspect
                                if inspect.iscoroutinefunction(tool_func):
                                    tool_result = await tool_func(**tool_args)
                                else:
                                    tool_result = tool_func(**tool_args)
                            except Exception as e:
                                tool_result = f"[工具执行错误] {e}"
                        else:
                            tool_result = f"[工具未找到] {tool_name}"

                        tool_msg = ToolMessage(
                            content=str(tool_result),
                            name=tool_name,
                            tool_call_id=tc["id"],
                        )
                        yield StreamEvent(type="messages-tuple", data={
                            "type": "tool",
                            "content": str(tool_result),
                            "name": tool_name,
                            "id": tc["id"],
                        })
                        messages.append(tool_msg)
                else:
                    # Final answer
                    text = self._extract_text(full_content)
                    text = self._filter_think_tags(text)
                    ai_msg = AIMessage(content=text, id=str(uuid.uuid4()))
                    if text:
                        yield StreamEvent(type="messages-tuple", data={
                            "type": "ai", "content": text, "id": ai_msg.id,
                        })
                    messages.append(ai_msg)
                    break

                # Yield values snapshot
                yield StreamEvent(type="values", data={
                    "messages": [self._serialize(m) for m in messages],
                })

            self.checkpointer[key] = messages

        except Exception as e:
            logger.exception(f"Agent stream error: {e}")
            yield StreamEvent(type="error", data={"error": str(e)})

        yield StreamEvent(type="end", data={})

    @staticmethod
    def _filter_think_tags(text: str) -> str:
        """Filter R1 thinking tags"""
        if not text:
            return text
        # blocks
        text = re.sub(r' thinking.*? response', '', text, flags=re.DOTALL)
        # Remove partial think tags
        text = re.sub(r'</?think>', '', text)
        return text.strip()

    @staticmethod
    def _extract_text(content) -> str:
        """Extract text from various content formats"""
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            parts = []
            for item in content:
                if isinstance(item, str):
                    parts.append(item)
                elif isinstance(item, dict):
                    parts.append(item.get("text", ""))
            return "".join(parts)
        return str(content)

    @staticmethod
    def _serialize(msg) -> dict:
        """Serialize a message to dict for SSE output"""
        t = getattr(msg, 'type', 'unknown')
        d = {
            "type": t,
            "content": StandaloneAgentEngine._extract_text(getattr(msg, 'content', '')),
            "id": getattr(msg, 'id', None),
        }
        if hasattr(msg, 'tool_calls') and msg.tool_calls:
            d["tool_calls"] = [
                {"name": tc["name"], "args": tc["args"], "id": tc.get("id")}
                for tc in msg.tool_calls
            ]
        if hasattr(msg, 'usage_metadata') and msg.usage_metadata:
            d["usage_metadata"] = msg.usage_metadata
        return d