"""LangGraph Agent service with tool calling capabilities."""

import asyncio
from typing import List, Dict, Any, Optional, AsyncGenerator
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import tool

# Import logger first
from ...utils.logger import get_logger
logger = get_logger("langgraph_agent_service")

# Try to import langgraph related modules
try:
    from langgraph.prebuilt import create_react_agent
    LANGGRAPH_AVAILABLE = True
except ImportError:
    logger.warning("langgraph not available. LangGraph agent functionality will be disabled.")
    create_react_agent = None
    LANGGRAPH_AVAILABLE = False

# Try to import init_chat_model from langchain_openai
try:
    from langchain_openai import ChatOpenAI
    # Create a simple init_chat_model function as a replacement
    def init_chat_model(model_name: str, **kwargs):
        return ChatOpenAI(model=model_name, **kwargs)
except ImportError:
    logger.warning("langchain_openai not available. Chat model functionality may be limited.")
    init_chat_model = None
from pydantic import BaseModel, Field

from .base import ToolRegistry
from th_agenter.services.tools import WeatherQueryTool, TavilySearchTool, DateTimeTool
from ..postgresql_tool_manager import get_postgresql_tool
from ...core.config import get_settings
from ...utils.logger import get_logger
from ..agent_config import AgentConfigService

logger = get_logger("langgraph_agent_service")



class LangGraphAgentConfig(BaseModel):
    """LangGraph Agent configuration."""
    model_name: str = Field(default="gpt-3.5-turbo")
    model_provider: str = Field(default="openai")
    base_url: Optional[str] = Field(default=None)
    api_key: Optional[str] = Field(default=None)
    enabled_tools: List[str] = Field(default_factory=lambda: [
        "calculator", "weather", "search", "file", "image"
    ])
    max_iterations: int = Field(default=10)
    temperature: float = Field(default=0.7)
    max_tokens: int = Field(default=1000)
    system_message: str = Field(
        default="""你是一个有用的AI助手，可以使用各种工具来帮助用户解决问题。
                重要规则：
                1. 工具调用失败时，必须仔细分析失败原因，特别是参数格式问题 
                3. 在重新调用工具前，先解释上次失败的原因和改进方案
                4. 确保每个工具调用的参数格式严格符合工具的要求 """
    )
    verbose: bool = Field(default=True)


class LangGraphAgentService:
    """LangGraph Agent service using create_react_agent."""
    
    def __init__(self, db_session=None):
        self.settings = get_settings()
        self.tool_registry = ToolRegistry()
        self.config = LangGraphAgentConfig()
        
        # Check if langgraph is available
        if not LANGGRAPH_AVAILABLE:
            logger.warning("LangGraph is not available. Some features may be disabled.")
        self.tools = []
        self.db_session = db_session
        self.config_service = AgentConfigService(db_session) if db_session else None
        self._initialize_tools()
        self._load_config()
        self._create_agent()
        
    def _initialize_tools(self):
        """Initialize available tools."""
        # Use the @tool decorated functions
        self.tools = [

            WeatherQueryTool(),
            TavilySearchTool(),
            DateTimeTool()
        ]
        

            
    def _load_config(self):
        """Load configuration from database if available."""
        if self.config_service:
            try:
                db_config = self.config_service.get_active_config()
                if db_config:
                    # Update config with database values
                    config_dict = db_config.config_data
                    for key, value in config_dict.items():
                        if hasattr(self.config, key):
                            setattr(self.config, key, value)
                    logger.info("Loaded configuration from database")
            except Exception as e:
                logger.warning(f"Failed to load config from database: {e}")
                

        
    def _create_agent(self):
        """Create LangGraph agent using create_react_agent."""
        try:
            # Initialize the model
            llm_config = get_settings().llm.get_current_config()
            self.model = init_chat_model(
                model=llm_config['model'],
                model_provider='openai',
                temperature=llm_config['temperature'],
                max_tokens=llm_config['max_tokens'],
                base_url= llm_config['base_url'],
                api_key=llm_config['api_key']
            )
            

            
            # Create the react agent
            self.agent = create_react_agent(
                model=self.model,
                tools=self.tools,)
            
            logger.info("LangGraph React agent created successfully")
            
        except Exception as e:
            logger.error(f"Failed to create agent: {str(e)}")
            raise
        

            

            

        
    def _format_tools_info(self) -> str:
        """Format tools information for the prompt."""
        tools_info = []
        for tool_name in self.config.enabled_tools:
            tool = self.tool_registry.get_tool(tool_name)
            if tool:
                params_info = []
                for param in tool.get_parameters():
                    params_info.append(f"  - {param.name} ({param.type.value}): {param.description}")
                
                tool_info = f"**{tool.get_name()}**: {tool.get_description()}"
                if params_info:
                    tool_info += "\n" + "\n".join(params_info)
                tools_info.append(tool_info)
                
        return "\n\n".join(tools_info)
        

        
    async def chat(self, message: str, chat_history: Optional[List[Dict[str, str]]] = None) -> Dict[str, Any]:
        """Process a chat message using LangGraph."""
        try:
            logger.info(f"Starting chat with message: {message[:100]}...")
            
            # Convert chat history to messages
            messages = []
            if chat_history:
                for msg in chat_history:
                    if msg["role"] == "user":
                        messages.append(HumanMessage(content=msg["content"]))
                    elif msg["role"] == "assistant":
                        messages.append(AIMessage(content=msg["content"]))
                        
            # Add current message
            messages.append(HumanMessage(content=message))
            
            # Use the react agent directly
            result = await self.agent.ainvoke({"messages": messages}, {"recursion_limit": 6},)
            
            # Extract final response
            final_response = ""
            if "messages" in result and result["messages"]:
                last_message = result["messages"][-1]
                if hasattr(last_message, 'content'):
                    final_response = last_message.content
                elif isinstance(last_message, dict) and "content" in last_message:
                    final_response = last_message["content"]
                    
            return {
                "response": final_response,
                "intermediate_steps": [],
                "success": True,
                "error": None
            }
            
        except Exception as e:
            logger.error(f"LangGraph chat error: {str(e)}", exc_info=True)
            return {
                "response": f"抱歉，处理您的请求时出现错误: {str(e)}",
                "intermediate_steps": [],
                "success": False,
                "error": str(e)
            }

    async def chat_stream(self, message: str, chat_history: Optional[List[Dict[str, str]]] = None) -> AsyncGenerator[
        Dict[str, Any], None]:
        """Process a chat message using LangGraph with streaming."""
        try:
            logger.info(f"Starting streaming chat with message: {message[:100]}...")

            # Convert chat history to messages
            messages = []
            if chat_history:
                for msg in chat_history:
                    if msg["role"] == "user":
                        messages.append(HumanMessage(content=msg["content"]))
                    elif msg["role"] == "assistant":
                        messages.append(AIMessage(content=msg["content"]))

            # Add current message
            messages.append(HumanMessage(content=message))

            # Track state for streaming
            intermediate_steps = []
            final_response_started = False
            accumulated_response = ""
            final_ai_message = None

            # Stream the agent execution
            async for event in self.agent.astream({"messages": messages}):
                # Handle different event types from LangGraph
                print('event===', event)
                if isinstance(event, dict):
                    for node_name, node_output in event.items():
                        logger.info(f"Processing node: {node_name}, output type: {type(node_output)}")

                        # 处理 tools 节点
                        if "tools" in node_name.lower():
                            # 提取工具信息
                            tool_infos = []

                            if isinstance(node_output, dict) and "messages" in node_output:
                                messages_in_output = node_output["messages"]

                                for msg in messages_in_output:
                                    # 处理 ToolMessage 对象
                                    if hasattr(msg, 'name') and hasattr(msg, 'content'):
                                        tool_info = {
                                            "tool_name": msg.name,
                                            "tool_output": msg.content,
                                            "tool_call_id": getattr(msg, 'tool_call_id', ''),
                                            "status": "completed"
                                        }
                                        tool_infos.append(tool_info)
                                    elif isinstance(msg, dict):
                                        if 'name' in msg and 'content' in msg:
                                            tool_info = {
                                                "tool_name": msg['name'],
                                                "tool_output": msg['content'],
                                                "tool_call_id": msg.get('tool_call_id', ''),
                                                "status": "completed"
                                            }
                                            tool_infos.append(tool_info)

                            # 返回 tools_end 事件
                            for tool_info in tool_infos:
                                yield {
                                    "type": "tools_end",
                                    "content": f"工具 {tool_info['tool_name']} 执行完成",
                                    "tool_name": tool_info["tool_name"],
                                    "tool_output": tool_info["tool_output"],
                                    "node_name": node_name,
                                    "done": False
                                }
                                await asyncio.sleep(0.1)

                        # 处理 agent 节点
                        elif "agent" in node_name.lower():
                            if isinstance(node_output, dict) and "messages" in node_output:
                                messages_in_output = node_output["messages"]
                                if messages_in_output:
                                    last_msg = messages_in_output[-1]

                                    # 获取 finish_reason
                                    finish_reason = None
                                    if hasattr(last_msg, 'response_metadata'):
                                        finish_reason = last_msg.response_metadata.get('finish_reason')
                                    elif isinstance(last_msg, dict) and 'response_metadata' in last_msg:
                                        finish_reason = last_msg['response_metadata'].get('finish_reason')

                                    # 判断是否为 thinking 或 response
                                    if finish_reason == 'tool_calls':
                                        # thinking 状态
                                        thinking_content = "🤔 正在思考..."
                                        if hasattr(last_msg, 'content') and last_msg.content:
                                            thinking_content = f"🤔 思考: {last_msg.content[:200]}..."
                                        elif isinstance(last_msg, dict) and "content" in last_msg:
                                            thinking_content = f"🤔 思考: {last_msg['content'][:200]}..."

                                        yield {
                                            "type": "thinking",
                                            "content": thinking_content,
                                            "node_name": node_name,
                                            "raw_output": str(node_output)[:500] if node_output else "",
                                            "done": False
                                        }
                                        await asyncio.sleep(0.1)

                                    elif finish_reason == 'stop':
                                        # response 状态
                                        if hasattr(last_msg, 'content') and hasattr(last_msg,
                                                                                    '__class__') and 'AI' in last_msg.__class__.__name__:
                                            current_content = last_msg.content
                                            final_ai_message = last_msg

                                            if not final_response_started and current_content:
                                                final_response_started = True
                                                yield {
                                                    "type": "response_start",
                                                    "content": "",
                                                    "intermediate_steps": intermediate_steps,
                                                    "done": False
                                                }

                                            if current_content and len(current_content) > len(accumulated_response):
                                                new_content = current_content[len(accumulated_response):]

                                                for char in new_content:
                                                    accumulated_response += char
                                                    yield {
                                                        "type": "response",
                                                        "content": accumulated_response,
                                                        "intermediate_steps": intermediate_steps,
                                                        "done": False
                                                    }
                                                    await asyncio.sleep(0.03)

                                    else:
                                        # 其他 agent 状态
                                        yield {
                                            "type": "step",
                                            "content": f"📋 执行步骤: {node_name}",
                                            "node_name": node_name,
                                            "raw_output": str(node_output)[:500] if node_output else "",
                                            "done": False
                                        }
                                        await asyncio.sleep(0.1)

                        # 处理其他节点
                        else:
                            yield {
                                "type": "step",
                                "content": f"📋 执行步骤: {node_name}",
                                "node_name": node_name,
                                "raw_output": str(node_output)[:500] if node_output else "",
                                "done": False
                            }
                            await asyncio.sleep(0.1)

            # 最终完成事件
            yield {
                "type": "complete",
                "content": accumulated_response,
                "intermediate_steps": intermediate_steps,
                "done": True
            }

        except Exception as e:
            logger.error(f"Error in chat_stream: {str(e)}", exc_info=True)
            yield {
                "type": "error",
                "content": f"处理请求时出错: {str(e)}",
                "done": True
            }
                                        
            # 确保最终响应包含完整内容
            final_content = accumulated_response
            if not final_content and final_ai_message and hasattr(final_ai_message, 'content'):
                final_content = final_ai_message.content or ""
            
            # Final completion signal
            yield {
                "type": "response",
                "content": final_content,
                "intermediate_steps": intermediate_steps,
                "done": True
            }
            
        except Exception as e:
            logger.error(f"LangGraph chat stream error: {str(e)}", exc_info=True)
            yield {
                "type": "error",
                "content": f"抱歉，处理您的请求时出现错误: {str(e)}",
                "error": str(e),
                "done": True
            }
            
    def get_available_tools(self) -> List[Dict[str, Any]]:
        """Get list of available tools."""
        tools = []
        for tool in self.tools:
            tools.append({
                "name": tool.name,
                "description": tool.description,
                "parameters": [],
                "enabled": True
            })
        return tools
        
    def get_config(self) -> Dict[str, Any]:
        """Get current agent configuration."""
        return self.config.dict()
        
    def update_config(self, config: Dict[str, Any]):
        """Update agent configuration."""
        for key, value in config.items():
            if hasattr(self.config, key):
                setattr(self.config, key, value)
        
        # Recreate agent with new config
        self._create_agent()
        logger.info("Agent configuration updated")


# Global instance
_langgraph_agent_service: Optional[LangGraphAgentService] = None


def get_langgraph_agent_service(db_session=None) -> LangGraphAgentService:
    """Get or create LangGraph agent service instance."""
    global _langgraph_agent_service
    
    if _langgraph_agent_service is None:
        _langgraph_agent_service = LangGraphAgentService(db_session)
        logger.info("LangGraph Agent service initialized")
        
    return _langgraph_agent_service