#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2023/9/13 12:29
@Author  : femto Zheng
@File    : brain.py
"""
import uuid
from datetime import datetime
from typing import Any, Dict, Union

from jinja2 import Template
from pydantic import BaseModel

# Lazy import for mem0 (heavy dependency, only load when needed)
Memory = None
def _get_memory_class():
    global Memory
    if Memory is None:
        from mem0 import Memory as _Memory
        Memory = _Memory
    return Memory
from tenacity import retry, stop_after_attempt, retry_if_exception_type

from minion import config
from minion.actions.lmp_action_node import LmpActionNode
from minion.main.input import Input
from minion.main.async_python_executor import AsyncPythonExecutor
from minion.utils.utils import process_image
from minion.main.worker import ModeratorMinion
from minion.providers import create_llm_provider
from minion.types.agent_response import AgentResponse
from minion.types.agent_state import AgentState

class Mind(BaseModel):
    id: str = "UnnamedMind"
    description: str = ""
    brain: Any = None  # Brain

    def step(self, input, selected_llm=None):
        moderator = ModeratorMinion(input=input, brain=self.brain, selected_llm=selected_llm)
        
        # 检查是否需要流式输出
        if hasattr(input, 'stream') and input.stream:
            # 流式输出：返回异步生成器
            return self._stream_step_generator(moderator)
        else:
            # 普通执行：返回协程
            return self._normal_step(moderator)
    
    async def _normal_step(self, moderator):
        """普通执行步骤"""
        agent_response = await moderator.execute()
        return agent_response
    
    async def _stream_step_generator(self, moderator):
        """真正的流式步骤生成器"""
        # 检查 moderator 是否支持流式输出
        if hasattr(moderator, 'execute_stream'):
            # 使用真正的流式输出
            async for chunk in moderator.execute_stream():
                yield chunk
        else:
            # 回退到普通执行
            agent_response = await moderator.execute()
            yield agent_response


class Brain:
    def __init__(
        self,
        id=None,
        memory=None,
        memory_config=None,
        llm=create_llm_provider(config.models.get("default")),
        llms={},
        python_env=None,
        stats_storer=None,
        tools=None,
        state = None
    ):
        self.id = id or uuid.uuid4()
        self.minds = {}
        self.add_mind(
            Mind(
                id="left_mind",
                description="""
I'm the left mind, adept at logical reasoning and analytical thinking. I excel in tasks involving mathematics, language, and detailed analysis. My capabilities include:

Solving complex mathematical problems
Understanding and processing language with precision
Performing logical reasoning and critical thinking
Analyzing and synthesizing information systematically
Engaging in tasks that require attention to detail and structure""",
            )
        )
        self.add_mind(
            Mind(
                id="right_mind",
                description="""
I'm the right mind, flourishing in creative and artistic tasks. I thrive in activities that involve imagination, intuition, and holistic thinking. My capabilities include:

Creating and appreciating art and music
Engaging in creative problem-solving and innovation
Understanding and interpreting emotions and expressions
Recognizing patterns and spatial relationships
Thinking in a non-linear and abstract manner""",
            )
        )

        self.add_mind(
            Mind(
                id="hippocampus_mind",
                description="""I'm the hippocampus mind, specializing in memory formation, organization, and retrieval. I play a crucial role in both the storage of new memories and the recall of past experiences. My capabilities include:

Forming and consolidating new memories
Organizing and structuring information for easy retrieval
Facilitating the recall of past experiences and learned information
Connecting new information with existing knowledge
Supporting navigation and spatial memory""",
            )
        )

        # self.add_mind(Mind(id="hypothalamus", description="..."))
        self.mem = memory
        if not memory:
            if memory_config:
                MemoryClass = _get_memory_class()
                self.mem = memory = MemoryClass.from_config(memory_config)

        # Process default llm
        if isinstance(llm, str):
            model_config = config.models.get(llm)
            if model_config is None:
                # Use pseudo provider when no model is configured
                from minion.providers.pseudo_provider import get_pseudo_provider
                logger.warning(
                    f"Model '{llm}' not found in config. Using pseudo provider. "
                    f"See: https://github.com/femto/minion#configuration"
                )
                self.llm = get_pseudo_provider()
            else:
                self.llm = create_llm_provider(model_config)
        else:
            self.llm = llm

        # Process llms dictionary
        from minion.providers.pseudo_provider import get_pseudo_provider
        self.llms = {}
        for key, value in llms.items():
            if isinstance(value, str):
                # Single model name
                model_cfg = config.models.get(value)
                if model_cfg is None:
                    logger.warning(f"Model '{value}' not found for llms['{key}'], using pseudo provider")
                    self.llms[key] = get_pseudo_provider()
                else:
                    self.llms[key] = create_llm_provider(model_cfg)
            elif isinstance(value, list):
                # List of model names or providers
                providers = []
                for item in value:
                    if isinstance(item, str):
                        item_cfg = config.models.get(item)
                        if item_cfg is None:
                            logger.warning(f"Model '{item}' not found for llms['{key}'], using pseudo provider")
                            providers.append(get_pseudo_provider())
                        else:
                            providers.append(create_llm_provider(item_cfg))
                    else:
                        providers.append(item)
                self.llms[key] = providers
            else:
                # Assume it's already a provider instance
                self.llms[key] = value

        # 设置工具集
        self.tools = tools or []
        
        # Agent state reference (will be set by agent during execution)
        self.state = None
        
        # 默认使用 AsyncPythonExecutor，提供更好的异步支持
        self.python_env = python_env or AsyncPythonExecutor(additional_authorized_imports=["numpy", "pandas", "json", "csv", "multi_tool_use", "inspect"])

        self.stats_storer = stats_storer
        if state is None:
            from minion.types.history import History
            state = AgentState(history=History()) #when brain is used standalone, create a AgentState to bind to it, it's like a simple agent
        self.state = state

    def add_tool(self, tool):
        """
        添加工具到Brain
        """
        self.tools.append(tool)

    def add_mind(self, mind):
        self.minds[mind.id] = mind
        mind.brain = self

    def process_image_input(self, input):
        if input.images:
            if isinstance(input.images, str):
                input.images = process_image(input.images)
            elif isinstance(input.images, list):
                input.images = [process_image(img) for img in input.images]
            else:
                raise ValueError("input.images should be either a string or a list of strings/images")
        return input.images

    async def step(self, state: Union[AgentState, Input, Dict[str, Any]] = None, **config_kwargs):
        # 处理不同类型的state参数
        if state is None:
            # 如果state为None，使用Brain自身的state或创建新的AgentState
            if self.state is not None:
                state = self.state
            else:
                from minion.types.history import History
                state = AgentState(history=History())

        # 根据state类型提取参数
        if isinstance(state, Input):
            # 直接是Input对象
            input = state
            current_tools = config_kwargs.get("tools", [])
        elif isinstance(state, AgentState):
            # 强类型AgentState
            input = state.input
            current_tools = config_kwargs.get("tools", [])
        elif isinstance(state, dict):
            # 字典格式（向后兼容）
            state = AgentState.model_validate(state) #convert dict to AgentState
            input = state.input
            current_tools = config_kwargs.get("tools", [])
        else:
            raise ValueError(f"Unsupported state type: {type(state)}")

        # 注意：不在这里设置 self.state，因为：
        # 1. Brain 构造时已经通过 Brain(state=agent.state) 建立了引用关系
        # 2. agent.state 的任何修改都会自动反映到 brain.state
        # 3. 重新赋值会破坏引用关系，导致状态不同步
        # 4. 如果 brain.step() 被直接调用（不通过 agent），self.state 可能为 None，这是正常的
        # 但是，如果 self.state 为 None，我们需要设置它，以便后续访问
        if self.state is None:
            self.state = state
        # 更新state的input字段
        state.input = input #now set input back to state, make sure state.input is this.

        # 从config_kwargs提取其他参数
        query = config_kwargs.get("query", "")
        query_type = config_kwargs.get("query_type", "")
        system_prompt = config_kwargs.get("system_prompt")
        messages = config_kwargs.get("messages")
        stream = config_kwargs.get("stream", False)
        selected_llm = config_kwargs.get("llm")  # 可选的LLM BaseProvider实例
        
        # 如果selected_llm是字符串，转换为BaseProvider实例
        if isinstance(selected_llm, str):
            selected_llm = create_llm_provider(config.models.get(selected_llm))
        
        # 验证必须有input或者有query/messages
        if input is None:
            if not query and messages is None:
                raise ValueError("State must contain 'input' or query/messages must be provided in config_kwargs")
            
            # 准备Input创建参数
            input_kwargs = config_kwargs.copy()
            input_kwargs.pop('query', None)
            input_kwargs.pop('query_type', None) 
            input_kwargs.pop('system_prompt', None)
            input_kwargs.pop('tools', None)
            input_kwargs.pop('messages', None)
            input_kwargs.pop('stream', None)
            input_kwargs.pop('llm', None)
            
            # 如果有messages，优先使用messages创建Input
            if messages is not None:
                # 支持标准 OpenAI messages 格式
                input = Input(query=messages, query_type=query_type, query_time=datetime.utcnow(), **input_kwargs)
            else:
                # 否则使用query创建Input
                input = Input(query=query, query_type=query_type, query_time=datetime.utcnow(), **input_kwargs)
        
        # 如果传入了messages，从中提取system_prompt（如果没有显式提供的话）
        if messages is not None and system_prompt is None:
            if isinstance(messages, list):
                for msg in messages:
                    if isinstance(msg, dict) and msg.get("role") == "system":
                        system_prompt = msg.get("content", "")
                        break  # 只取第一个system消息
            
        input.query_id = input.query_id or uuid.uuid4()
        input.images = self.process_image_input(input)  # normalize image format to base64
        
        # 设置工具和系统提示
        input.tools = current_tools
        if system_prompt is not None:
            input.system_prompt = system_prompt
        # 只有当 config_kwargs 中明确传递了 stream 参数时才覆盖
        if "stream" in config_kwargs:
            input.stream = stream

        # 选择心智
        mind_id = input.mind_id or await self.choose_mind(input)
        if mind_id == "left_mind":
            self.llm.config.temperature = 0.1
        elif mind_id == "right_mind":
            self.llm.config.temperature = 0.7
        mind = self.minds[mind_id]
        
        # 检查是否需要流式输出
        if hasattr(input, 'stream') and input.stream:
            # 设置 stream_outputs 属性，供 UI 使用
            self.stream_outputs = True
            # 流式输出：直接返回异步生成器
            return mind.step(input, selected_llm=selected_llm)
        else:
            # 清除 stream_outputs 属性
            self.stream_outputs = False
            # 普通执行：await 结果并确保返回AgentResponse
            result = await mind.step(input, selected_llm=selected_llm)
            
            # 确保结果是AgentResponse格式
            if not isinstance(result, AgentResponse):
                result = AgentResponse.from_tuple(result)
                
            return result

    def cleanup_python_env(self, input):
        if hasattr(self.python_env, 'step'):

            self.python_env.step("RESET_CONTAINER_SPECIAL_KEYWORD")
        else:
            # LocalPythonExecutor or AsyncPythonExecutor - reset by re-initializing state only
            if hasattr(self.python_env, 'send_variables'):
                self.python_env.send_variables(variables={})
            # Don't clear tools - they should persist across executions
            # Tools are only set when starting a new session with new tools

    async def choose_mind(self, input):
        return "left_mind" #don't choose mind for now
        mind_template = Template(
            """
I have minds:
{% for mind in minds %}
1. **ID:** {{ mind.id }}  
   **Description:** 
   "{{ mind.description }}"
{% endfor %}
According to the current user's query, 
which is of query type: {{ input.query_type }},
and user's query: {{ input.query }}
help me choose the right mind to process the query.
return the id of the mind, please note you *MUST* return exactly case same as I provided here, do not uppercase or downcase yourself.
"""
        )

        # Create the filled template
        filled_template = mind_template.render(minds=self.minds.values(), input=input)

        try:
            lmp_action_node = LmpActionNode(llm=self.llm)
            result = await lmp_action_node.execute_answer(filled_template)

            # Ensure the result is a valid mind ID
            if result not in self.minds:
                result = "left_mind"
                #raise ValueError(f"Invalid mind ID returned: {result}")

            return result
        except Exception as e:
            return "left_mind" #EXISTING_ANSWER_PROMPT for llama3.2 which can't return valid json

    def run(self, query: str, **kwargs) -> AgentResponse:
        """
        Synchronous interface for running the brain.
        
        Args:
            query: The query string to process
            **kwargs: Additional parameters
            
        Returns:
            AgentResponse: The response from the brain
        """
        import asyncio
        return asyncio.run(self.run_async(query, **kwargs))
    
    async def run_async(self, query: str, **kwargs):
        """
        Asynchronous interface for running the brain.
        
        Args:
            query: The query string to process
            **kwargs: Additional parameters
            
        Returns:
            AgentResponse or AsyncGenerator: Response or stream based on stream parameter
        """
        # Create Input object
        input_obj = Input(query=query)
        
        # Set stream parameter
        stream = kwargs.pop('stream', False)
        input_obj.stream = stream
        
        # Pass through other parameters
        for key, value in kwargs.items():
            if hasattr(input_obj, key):
                setattr(input_obj, key, value)
        
        # Call step method with the input
        state = {"input": input_obj}
        return await self.step(state, **kwargs)
    
    async def run_stream(self, query: str, **kwargs):
        """
        Streaming interface for running the brain.
        
        Args:
            query: The query string to process
            **kwargs: Additional parameters
            
        Returns:
            AsyncGenerator: Stream of responses
        """
        # Force stream=True for this method
        kwargs['stream'] = True
        return await self.run_async(query, **kwargs)


Mind.model_rebuild()