"""
行为树核心组件
定义行为树的基本结构和执行逻辑
包含完整的节点系统、黑板系统和执行管理
"""

import time
import math
from enum import Enum
from abc import ABC, abstractmethod
from typing import Dict, List, Any, Optional, Callable, Union
from dataclasses import dataclass, field

class NodeStatus(Enum):
    """
    节点执行状态
    """
    SUCCESS = "success"
    FAILURE = "failure"
    RUNNING = "running"

class NodeMemory:
    """
    节点内存
    存储节点状态信息，支持节点间的记忆功能
    """
    
    def __init__(self):
        self.data: Dict[str, Any] = {}
        self.last_execution_time: float = 0.0
        self.execution_count: int = 0
        self.running_child_index: int = 0  # 用于控制节点跟踪当前执行的子节点
        
    def set(self, key: str, value: Any) -> None:
        """设置内存值"""
        self.data[key] = value
        
    def get(self, key: str, default: Any = None) -> Any:
        """获取内存值"""
        return self.data.get(key, default)
        
    def has(self, key: str) -> bool:
        """检查内存是否存在键"""
        return key in self.data
        
    def clear(self) -> None:
        """清空内存"""
        self.data.clear()
        
    def increment_execution_count(self) -> int:
        """增加执行计数并返回新计数"""
        self.execution_count += 1
        return self.execution_count

@dataclass
class BTNodeConfig:
    """
    节点配置
    """
    name: str = "BTNode"
    priority: int = 0  # 优先级，数字越大优先级越高
    max_execution_time: Optional[float] = None  # 最大执行时间
    can_be_interrupted: bool = True  # 是否可被中断

class BTNode(ABC):
    """
    行为树节点基类
    """
    
    def __init__(self, config: Union[BTNodeConfig, str] = "BTNode"):
        if isinstance(config, str):
            config = BTNodeConfig(name=config)
        self.config = config
        self.children: List['BTNode'] = []
        self.status = NodeStatus.RUNNING
        self.memory = NodeMemory()
        self.parent: Optional['BTNode'] = None
        self.last_execution_time: float = 0.0
        self.execution_count: int = 0
        
    @property
    def name(self) -> str:
        """获取节点名称"""
        return self.config.name
        
    def add_child(self, child: 'BTNode') -> None:
        """添加子节点"""
        child.parent = self
        self.children.append(child)
        
    def insert_child(self, index: int, child: 'BTNode') -> None:
        """在指定位置插入子节点"""
        child.parent = self
        self.children.insert(index, child)
        
    def remove_child(self, child: 'BTNode') -> bool:
        """移除子节点"""
        if child in self.children:
            child.parent = None
            self.children.remove(child)
            return True
        return False
            
    def get_children(self) -> List['BTNode']:
        """获取子节点列表"""
        return self.children.copy()
        
    def get_child_count(self) -> int:
        """获取子节点数量"""
        return len(self.children)
        
    def clear_children(self) -> None:
        """清空所有子节点"""
        for child in self.children:
            child.parent = None
        self.children.clear()
        
    def get_root(self) -> 'BTNode':
        """获取根节点"""
        current = self
        while current.parent:
            current = current.parent
        return current
        
    def get_path(self) -> List[int]:
        """获取从根节点到当前节点的路径索引列表"""
        path = []
        current = self
        while current.parent:
            index = current.parent.children.index(current)
            path.append(index)
            current = current.parent
        path.reverse()
        return path
        
    def find_child_by_name(self, name: str) -> Optional['BTNode']:
        """根据名称查找子节点"""
        for child in self.children:
            if child.name == name:
                return child
            # 递归搜索
            result = child.find_child_by_name(name)
            if result:
                return result
        return None
        
    def reset(self) -> None:
        """重置节点状态"""
        self.status = NodeStatus.RUNNING
        self.memory.clear()
        self.last_execution_time = 0.0
        self.execution_count = 0
        # 递归重置子节点
        for child in self.children:
            child.reset()
        
    @abstractmethod
    def execute(self, blackboard) -> NodeStatus:
        """
        执行节点逻辑
        :param blackboard: 黑板对象，用于数据共享
        :return: NodeStatus 执行状态
        """
        start_time = time.time()
        self.execution_count += 1
        
        # 检查是否超过最大执行时间
        if self.config.max_execution_time:
            if time.time() - start_time > self.config.max_execution_time:
                return NodeStatus.FAILURE
                
        self.last_execution_time = time.time() - start_time
        return NodeStatus.SUCCESS

    def on_enter(self, blackboard) -> None:
        """节点进入时调用"""
        pass
        
    def on_exit(self, blackboard, status: NodeStatus) -> None:
        """节点退出时调用"""
        pass
        
    def on_abort(self, blackboard) -> None:
        """节点被中断时调用"""
        pass

class DecoratorNode(BTNode):
    """
    装饰节点基类
    只能有一个子节点，用于修饰子节点的行为
    """
    
    def __init__(self, child: Optional['BTNode'] = None, config: Union[BTNodeConfig, str] = "DecoratorNode"):
        super().__init__(config)
        self.child = child
        if child:
            child.parent = self
            
    def set_child(self, child: 'BTNode') -> None:
        """设置子节点"""
        if self.child:
            self.child.parent = None
        self.child = child
        if child:
            child.parent = self
        
    def get_children(self) -> List['BTNode']:
        """获取子节点列表"""
        if self.child:
            return [self.child]
        return []
        
    def get_child_count(self) -> int:
        """获取子节点数量"""
        return 1 if self.child else 0
        
    def execute(self, blackboard) -> NodeStatus:
        """执行装饰节点逻辑"""
        start_time = time.time()
        self.execution_count += 1
        
        # 检查是否超过最大执行时间
        if self.config.max_execution_time:
            if time.time() - start_time > self.config.max_execution_time:
                return NodeStatus.FAILURE
                
        if not self.child:
            self.last_execution_time = time.time() - start_time
            return NodeStatus.FAILURE
            
        child_status = self.child.execute(blackboard)
        result = self._decorate(child_status)
        
        self.last_execution_time = time.time() - start_time
        return result
        
    @abstractmethod
    def _decorate(self, child_status: NodeStatus) -> NodeStatus:
        """
        装饰子节点的执行结果
        :param child_status: 子节点执行状态
        :return: 装饰后的执行状态
        """
        pass

class CompositeNode(BTNode):
    """
    组合节点基类
    有多个子节点，定义子节点的执行顺序和逻辑
    """
    
    def __init__(self, config: Union[BTNodeConfig, str] = "CompositeNode"):
        super().__init__(config)
        self.current_child_index = 0
        
    def execute(self, blackboard) -> NodeStatus:
        """执行组合节点逻辑"""
        start_time = time.time()
        self.execution_count += 1
        
        # 检查是否超过最大执行时间
        if self.config.max_execution_time:
            if time.time() - start_time > self.config.max_execution_time:
                return NodeStatus.FAILURE
                
        if not self.children:
            self.last_execution_time = time.time() - start_time
            return NodeStatus.FAILURE
            
        result = self._execute_child_logic(blackboard)
        
        self.last_execution_time = time.time() - start_time
        return result
        
    @abstractmethod
    def _execute_child_logic(self, blackboard) -> NodeStatus:
        """
        执行子节点逻辑
        :param blackboard: 黑板对象
        :return: 执行结果状态
        """
        pass

class BehaviorTree:
    """
    行为树类
    """
    
    def __init__(self, name: str, root_node: BTNode):
        self.name = name
        self.root_node = root_node
        self.blackboard = None
        self.is_running = False
        self.start_time = 0.0
        self.last_execution_time = 0.0
        self.execution_count = 0
        self.status = NodeStatus.RUNNING
        self.custom_functions: Dict[str, Callable] = {}
        
    def set_blackboard(self, blackboard) -> None:
        """设置黑板"""
        self.blackboard = blackboard
        
    def run(self) -> NodeStatus:
        """运行行为树"""
        if not self.root_node or not self.blackboard:
            return NodeStatus.FAILURE
            
        self.is_running = True
        self.start_time = time.time()
        self.execution_count += 1
        
        try:
            self.status = self.root_node.execute(self.blackboard)
        except Exception as e:
            print(f"行为树执行出错: {e}")
            self.status = NodeStatus.FAILURE
            import traceback
            traceback.print_exc()
            
        self.last_execution_time = time.time() - self.start_time
        self.is_running = False
        
        return self.status
        
    def execute(self) -> NodeStatus:
        """执行行为树（同run方法）"""
        return self.run()
        
    def reset(self) -> None:
        """重置行为树"""
        if self.root_node:
            self.root_node.reset()
        self.is_running = False
        self.status = NodeStatus.RUNNING
        self.execution_count = 0
        
    def add_custom_function(self, name: str, func: Callable) -> None:
        """添加自定义函数"""
        self.custom_functions[name] = func
        
    def get_custom_function(self, name: str) -> Optional[Callable]:
        """获取自定义函数"""
        return self.custom_functions.get(name)
        
    def get_node_by_path(self, path: List[int]) -> Optional[BTNode]:
        """根据路径获取节点"""
        current = self.root_node
        for index in path:
            if 0 <= index < len(current.children):
                current = current.children[index]
            else:
                return None
        return current
        
    def get_statistics(self) -> Dict[str, Any]:
        """获取行为树统计信息"""
        def collect_node_stats(node: BTNode, stats: Dict[str, Any]) -> None:
            node_type = type(node).__name__
            if node_type not in stats['node_types']:
                stats['node_types'][node_type] = {'count': 0, 'total_executions': 0, 'avg_execution_time': 0}
            
            stats['node_types'][node_type]['count'] += 1
            stats['node_types'][node_type]['total_executions'] += node.execution_count
            stats['total_nodes'] += 1
            stats['total_executions'] += node.execution_count
            
            for child in node.children:
                collect_node_stats(child, stats)
        
        stats = {
            'name': self.name,
            'status': self.status.value,
            'is_running': self.is_running,
            'execution_count': self.execution_count,
            'last_execution_time': self.last_execution_time,
            'total_nodes': 0,
            'total_executions': 0,
            'node_types': {}
        }
        
        if self.root_node:
            collect_node_stats(self.root_node, stats)
            
        return stats

class BehaviorTreeManager:
    """
    行为树管理器
    用于管理多个行为树的生命周期
    """
    
    def __init__(self):
        self.trees: Dict[str, BehaviorTree] = {}
        self.active_trees: List[str] = []
        self.update_interval = 0.016  # 默认每帧更新（假设60FPS）
        self.last_update_time = 0.0
        
    def add_tree(self, name: str, tree: BehaviorTree) -> bool:
        """添加行为树"""
        if name in self.trees:
            return False
        self.trees[name] = tree
        return True
        
    def remove_tree(self, name: str) -> bool:
        """移除行为树"""
        if name in self.trees:
            del self.trees[name]
            if name in self.active_trees:
                self.active_trees.remove(name)
            return True
        return False
        
    def get_tree(self, name: str) -> Optional[BehaviorTree]:
        """获取行为树"""
        return self.trees.get(name)
        
    def activate_tree(self, name: str) -> bool:
        """激活行为树"""
        if name in self.trees and name not in self.active_trees:
            self.active_trees.append(name)
            return True
        return False
        
    def deactivate_tree(self, name: str) -> bool:
        """停用行为树"""
        if name in self.active_trees:
            self.active_trees.remove(name)
            return True
        return False
        
    def update_all(self, dt: float = None) -> None:
        """更新所有激活的行为树"""
        current_time = time.time()
        if dt is None:
            dt = current_time - self.last_update_time
        self.last_update_time = current_time
        
        # 只更新激活的树
        for tree_name in self.active_trees.copy():
            tree = self.trees.get(tree_name)
            if tree and not tree.is_running:
                tree.run()
                
    def get_tree_stats(self, name: str) -> Optional[Dict[str, Any]]:
        """获取特定树的统计信息"""
        tree = self.get_tree(name)
        if tree:
            return tree.get_statistics()
        return None
        
    def get_all_stats(self) -> Dict[str, Dict[str, Any]]:
        """获取所有树的统计信息"""
        stats = {}
        for name, tree in self.trees.items():
            stats[name] = tree.get_statistics()
        return stats
        
    def reset_all(self) -> None:
        """重置所有树"""
        for tree in self.trees.values():
            tree.reset()