"""
组件抽象层：Agent 基础抽象类

定义不同类型 Agent 的抽象基类，支持基于 LangGraph 的多种实现模式
"""

from abc import ABC, abstractmethod
from typing import Dict, Any, List, Optional, Union, TypeVar, Generic
from langchain_core.language_models import BaseChatModel
from langchain_core.tools import BaseTool
from langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph
from langgraph.types import Command
from langchain_google_genai import ChatGoogleGenerativeAI, HarmCategory, HarmBlockThreshold
from langchain_openai import ChatOpenAI
from src.utils.logger import get_logger
from langgraph.prebuilt import create_react_agent

# 状态类型变量
StateType = TypeVar('StateType', bound=Dict[str, Any])


class BaseComponent(ABC):
    """原子组件基类
    
    所有原子组件的最基础抽象，定义单一职责的组件接口
    """
    
    def __init__(self, name: str, description: str = ""):
        """初始化组件
        
        Args:
            name: 组件名称，唯一标识
            description: 组件描述
        """
        self.name = name
        self.description = description
        self._initialized = False
    
    @abstractmethod
    def initialize(self) -> None:
        """初始化组件，子类必须实现"""
        pass
    
    @property
    def is_initialized(self) -> bool:
        """检查组件是否已初始化"""
        return self._initialized
    
    def __str__(self) -> str:
        return f"{self.__class__.__name__}(name='{self.name}')"
    
    def __repr__(self) -> str:
        return self.__str__()


class BaseAgent(BaseComponent, Generic[StateType]):
    """Agent 基类
    
    定义所有 Agent 的通用接口和行为
    """
    
    def __init__(
        self, 
        name: str, 
        description: str = "",
        model_provider: str = "google_genai",
        tools: Optional[List[BaseTool]] = None
    ):
        """初始化 Agent
        
        Args:
            name: Agent 名称
            description: Agent 描述
            model_provider: 模型提供商 ("openai" 或 "google_genai")
            tools: Agent 可用的工具列表
        """
        super().__init__(name, description)
        self.model_provider = model_provider
        self.tools = tools or []
        self.model: Optional[BaseChatModel] = None
    
    @abstractmethod
    def process(self, state: StateType, config: Optional[RunnableConfig] = None) -> Union[StateType, Command]:
        """处理状态
        
        Args:
            state: 输入状态
            config: 运行时配置
            
        Returns:
            处理后的状态或命令对象
        """
        pass
    
    @abstractmethod
    def _initialize_model(self) -> BaseChatModel:
        """初始化模型，子类必须实现"""
        pass
    
    def initialize(self) -> None:
        """初始化 Agent"""
        if not self._initialized:
            self.model = self._initialize_model()
            if self.tools:
                self.model = self.model.bind_tools(self.tools)
            self._initialized = True
    
    def add_tool(self, tool: BaseTool) -> None:
        """添加工具
        
        Args:
            tool: 要添加的工具
        """
        if tool not in self.tools:
            self.tools.append(tool)
            # 如果已初始化，需要重新绑定工具
            if self._initialized and self.model:
                self.model = self.model.bind_tools(self.tools)


class BaseLLMAgent(BaseAgent[StateType]):
    """基础 LLM Agent
    
    直接基于 LLM 的简单智能体，适用于单轮对话和简单任务
    """
    
    def __init__(
        self,
        name: str,
        description: str = "",
        model_provider: str = "google_genai",
        system_prompt: str = "",
        temperature: float = 0.1,
        max_tokens: int = 10240,
        tools: Optional[List[BaseTool]] = None
    ):
        super().__init__(name, description, model_provider, tools)
        self.system_prompt = system_prompt
        self.temperature = temperature
        self.max_tokens = max_tokens
    
    def _initialize_model(self) -> BaseChatModel:
        """初始化 LLM 模型"""
        import os

        if self.model_provider == "google_genai":
            # 使用直接的ChatGoogleGenerativeAI来避免认证问题

            api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
            if not api_key:
                raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY not found in environment")

            return ChatGoogleGenerativeAI(
                model="gemini-2.5-flash",
                google_api_key=api_key,
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                # 调整安全设置，避免正常内容被误拦截
                safety_settings={
                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH
                }
            )
        else:
            # OpenAI模型

            return ChatOpenAI(
                model="gpt-4o",
                temperature=self.temperature,
                max_tokens=self.max_tokens
            )
    
    def process(self, state: StateType, config: Optional[RunnableConfig] = None) -> StateType:
        """处理状态 - LLM 直接处理"""
        if not self.is_initialized:
            self.initialize()

        # 构建消息
        messages = self._build_messages(state)

        # 调用模型（添加错误处理）
        try:
            response = self.model.invoke(messages)

            # 检查响应有效性
            if not response:
                raise ValueError("LLM returned None response")

            if hasattr(response, 'content'):
                if not response.content or response.content.strip() == "":
                    raise ValueError(f"LLM returned empty content {response}")

        except Exception as e:
            # 记录错误并返回带错误标记的状态
            logger = get_logger(self.name)
            logger.error(f"LLM调用失败: {e}", exc_info=True)

            # 创建一个空响应对象
            class EmptyResponse:
                content = ""
            response = EmptyResponse()

        # 更新状态
        return self._update_state(state, response)
    
    @abstractmethod
    def _build_messages(self, state: StateType) -> List[Dict[str, Any]]:
        """构建消息，子类实现具体逻辑"""
        """
        文本消息和图片消息可以交织，以下是一个例子
        message_content = [
            {"type": "text", "text": text_prompt},  # 文本类型消息
            ...
            {"type": "image_url", "image_url": {"url": image_url}}, # 图片类型消息
            ...
            {"type": "image_url", "image_url": {"url": image_url}}, # 图片类型消息
            {"type": "text", "text": text_prompt},  # 文本类型消息
            ...
        ]
        """

    
    @abstractmethod
    def _update_state(self, state: StateType, response: Any) -> StateType:
        """更新状态，子类实现具体逻辑"""
        pass


class BaseReactAgent(BaseAgent[StateType]):
    """基础 ReAct Agent
    
    基于 LangGraph 的 ReAct 模式智能体，支持工具调用和推理
    """
    
    def __init__(
        self,
        name: str,
        description: str = "",
        model_provider: str = "google_genai",
        prompt: Union[str, callable] = "",
        tools: Optional[List[BaseTool]] = None
    ):
        super().__init__(name, description, model_provider, tools)
        self.prompt = prompt
        self.react_agent = None
    
    def _initialize_model(self) -> BaseChatModel:
        """初始化模型"""
        import os

        if self.model_provider == "google_genai":
            # 使用直接的ChatGoogleGenerativeAI来避免认证问题

            api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
            if not api_key:
                raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY not found in environment")

            return ChatGoogleGenerativeAI(
                model="gemini-2.5-flash",
                google_api_key=api_key,
                temperature=0.1,
                max_tokens=10240,
                # 调整安全设置，避免正常内容被误拦截
                safety_settings={
                    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
                    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
                    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
                    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH
                }
            )
        else:

            return ChatOpenAI(
                model="gpt-4o",
                temperature=0.1,
                max_tokens=10240
            )
    
    def initialize(self) -> None:
        """初始化 ReAct Agent"""
        if not self._initialized:
            super().initialize()
            # 创建 ReAct Agent
            
            self.react_agent = create_react_agent(
                self.model,
                tools=self.tools,
                prompt=self.prompt if self.prompt else None
            )
            self._initialized = True
    
    def process(self, state: StateType, config: Optional[RunnableConfig] = None) -> StateType:
        """处理状态 - 使用 ReAct Agent"""
        if not self.is_initialized:
            self.initialize()
        
        # 转换状态为消息格式
        messages = self._extract_messages(state)
        
        # 调用 ReAct Agent
        response = self.react_agent.invoke({"messages": messages}, config=config)
        
        # 更新状态
        return self._merge_response(state, response)
    
    @abstractmethod
    def _extract_messages(self, state: StateType) -> List[Dict[str, Any]]:
        """从状态中提取消息，子类实现"""
        pass
    
    @abstractmethod
    def _merge_response(self, state: StateType, response: Dict[str, Any]) -> StateType:
        """合并响应到状态，子类实现"""
        pass


class BaseGraphAgent(BaseAgent[StateType]):
    """基础 Graph Agent
    
    基于 StateGraph 的复杂智能体，支持多节点工作流
    """
    
    def __init__(
        self,
        name: str,
        description: str = "",
        state_class: type = dict
    ):
        super().__init__(name, description)
        self.state_class = state_class
        self.graph: Optional[StateGraph] = None
        self.compiled_graph = None
    
    def _initialize_model(self) -> BaseChatModel:
        """Graph Agent 不需要单一模型"""
        return None
    
    def initialize(self) -> None:
        """初始化 Graph Agent"""
        if not self._initialized:
            self.graph = self._build_graph()
            self.compiled_graph = self.graph.compile()
            self._initialized = True
    
    @abstractmethod
    def _build_graph(self) -> StateGraph:
        """构建图结构，子类必须实现"""
        pass
    
    def process(self, state: StateType, config: Optional[RunnableConfig] = None) -> StateType:
        """处理状态 - 使用编译的图"""
        if not self.is_initialized:
            self.initialize()
        
        result = self.compiled_graph.invoke(state, config=config)
        return result
    
    def add_node(self, name: str, node_func: callable) -> None:
        """添加节点到图中"""
        if self.graph is None:
            self.graph = StateGraph(self.state_class)
        self.graph.add_node(name, node_func)
        
        # 重新编译
        if self._initialized:
            self.compiled_graph = self.graph.compile()
    
    def add_edge(self, from_node: str, to_node: str) -> None:
        """添加边到图中"""
        if self.graph is None:
            raise ValueError("Graph not initialized. Call initialize() first.")
        
        self.graph.add_edge(from_node, to_node)
        
        # 重新编译
        if self._initialized:
            self.compiled_graph = self.graph.compile()
    
    def add_conditional_edges(
        self, 
        from_node: str, 
        condition_func: callable, 
        condition_map: Dict[str, str]
    ) -> None:
        """添加条件边到图中"""
        if self.graph is None:
            raise ValueError("Graph not initialized. Call initialize() first.")
        
        self.graph.add_conditional_edges(from_node, condition_func, condition_map)
        
        # 重新编译
        if self._initialized:
            self.compiled_graph = self.graph.compile()