virtual_env_adapter.py

"""
Virtual Environment Adapter
Adapts VirtualAndroidEnv to AsyncEnv interface for compatibility with existing Mobile-Agent-v3 framework
"""

import numpy as np
from typing import Dict, Optional, Any
from agent_env.env import interface
from agent_env.env import representation_utils
from android_env.proto import adb_pb2
from virtual_env_gemini3 import VirtualAndroidEnv


class VirtualEnvAdapter(interface.AsyncEnv):
    """
    Virtual Environment Adapter，将Gemini3虚拟环境适配为AsyncEnv接口
    """
    
    def __init__(
        self,
        gemini_api_key: str,
        gemini_base_url: str = "https://openrouter.ai/api/v1",
        gemini_model: str = "google/gemini-2.0-flash-exp:free",
        initial_task: str = "",
        initial_image_path: str = None,
        resolution: tuple = (1080, 2400)
    ):
        """
        初始化Virtual Environment Adapter
        
        Args:
            gemini_api_key: Gemini API key
            gemini_base_url: API address
            gemini_model: Model name
            initial_task: Initial task description
            initial_image_path: 初始图像路径
            resolution: Screen resolution
        """
        self._virtual_env = VirtualAndroidEnv(
            gemini_api_key=gemini_api_key,
            gemini_base_url=gemini_base_url,
            gemini_model=gemini_model,
            initial_task=initial_task,
            initial_image_path=initial_image_path,
            resolution=resolution
        )
        
        self._current_screenshot = None
        self._view_hierarchy = None  # 虚拟环境不需要真实的view hierarchy
        self._interaction_cache = ""  # 交互缓存
        
        print(f"[VirtualEnvAdapter] Initialized")
    
    @property
    def controller(self):
        """Return controller (virtual environment uses mock controller)"""
        return VirtualController(self._virtual_env)
    
    def reset(self, go_home: bool = False) -> interface.State:
        """
        Reset environment
        
        Args:
            go_home: Whether to go home (virtual environment ignores this parameter)
            
        Returns:
            Initial state
        """
        print(f"\n[VirtualEnvAdapter] Resetting environment...")
        
        self._current_screenshot = self._virtual_env.reset()
        return self._create_state()
    
    def get_state(self, wait_to_stabilize: bool = False) -> interface.State:
        """
        获取当前Environment state
        
        Args:
            wait_to_stabilize: Whether to wait for UI to stabilize (virtual environment ignores this parameter)
        
        Returns:
            Current state（State对象）
        """
        state = self._virtual_env.get_current_state()
        self._current_screenshot = state.get('screenshot', self._current_screenshot)
        return self._create_state()
    
    def _create_state(self) -> interface.State:
        """
        Create State object
        
        Returns:
            interface.State对象
        """
        # Create mock forest (simplified XML tree)
        forest = self._get_mock_forest()
        
        # Create empty UI elements list (can be populated as needed)
        ui_elements = []
        
        return interface.State(
            pixels=self._current_screenshot,
            forest=forest,
            ui_elements=ui_elements,
            auxiliaries={
                'app_name': self._virtual_env.get_current_state().get('app_name', 'Unknown'),
                'action_count': self._virtual_env.get_current_state().get('action_count', 0)
            }
        )
    
    def _get_mock_forest(self) -> Any:
        """
        生成模拟的accessibility forest
        返回简化的XML结构
        """
        from xml.etree import ElementTree as ET
        
        state = self._virtual_env.get_current_state()
        app_name = state.get('app_name', 'Unknown')
        width, height = self._virtual_env.resolution
        
        # 创建简单的XML树
        root = ET.Element('hierarchy', attrib={'rotation': '0'})
        frame_node = ET.SubElement(root, 'node', attrib={
            'index': '0',
            'text': '',
            'resource-id': '',
            'class': 'android.widget.FrameLayout',
            'package': app_name,
            'content-desc': '',
            'checkable': 'false',
            'checked': 'false',
            'clickable': 'false',
            'enabled': 'true',
            'focusable': 'false',
            'focused': 'false',
            'scrollable': 'false',
            'long-clickable': 'false',
            'password': 'false',
            'selected': 'false',
            'bounds': f'[0,0][{width},{height}]'
        })
        
        # 添加内容节点
        content_node = ET.SubElement(frame_node, 'node', attrib={
            'index': '0',
            'text': 'Virtual Android Screen',
            'resource-id': 'android:id/content',
            'class': 'android.view.View',
            'package': app_name,
            'content-desc': 'Virtual screen generated by Gemini3',
            'checkable': 'false',
            'checked': 'false',
            'clickable': 'true',
            'enabled': 'true',
            'focusable': 'true',
            'focused': 'false',
            'scrollable': 'true',
            'long-clickable': 'false',
            'password': 'false',
            'selected': 'false',
            'bounds': f'[0,100][{width},{height-100}]'
        })
        
        return root
    
    def _get_mock_view_hierarchy(self) -> str:
        """
        生成模拟的view hierarchy
        由于虚拟环境没有真实的view hierarchy，这里返回一个简化版本
        """
        state = self._virtual_env.get_current_state()
        app_name = state.get('app_name', 'Unknown')
        
        # 简单的mock view hierarchy
        mock_vh = f"""<?xml version="1.0" encoding="UTF-8"?>
<hierarchy rotation="0">
  <node index="0" text="" resource-id="" class="android.widget.FrameLayout" package="{app_name}" 
        content-desc="" checkable="false" checked="false" clickable="false" enabled="true" 
        focusable="false" focused="false" scrollable="false" long-clickable="false" 
        password="false" selected="false" bounds="[0,0][1080,2400]">
    <node index="0" text="Virtual Android Screen" resource-id="android:id/content" 
          class="android.view.View" package="{app_name}" content-desc="Virtual screen generated by Gemini3" 
          checkable="false" checked="false" clickable="true" enabled="true" 
          focusable="true" focused="false" scrollable="true" long-clickable="false" 
          password="false" selected="false" bounds="[0,100][1080,2300]"/>
  </node>
</hierarchy>"""
        return mock_vh
    
    def execute_action(self, action) -> None:
        """
        Execute action并更新环境
        
        Args:
            action: JSONAction对象或Action字典
        """
        print(f"\n[VirtualEnvAdapter] Executing action: {action}")
        
        # 将action转换为字典格式
        if hasattr(action, 'action_type'):
            # JSONAction对象 - 直接使用Agent提供的description
            action_type = action.action_type
            
            # 优先使用Agent提供的description，否则生成默认描述
            if hasattr(action, 'description') and action.description:
                description = action.description
            else:
                # 后备方案：根据Action类型生成描述
                if action_type == 'click':
                    if hasattr(action, 'x') and hasattr(action, 'y'):
                        description = f"Click at position ({action.x}, {action.y})"
                    else:
                        description = "Click on screen"
                elif action_type == 'type':
                    text = action.text if hasattr(action, 'text') else ""
                    description = f"Type the text: '{text}'"
                elif action_type == 'scroll':
                    direction = action.direction if hasattr(action, 'direction') else "down"
                    description = f"Scroll {direction}"
                elif action_type == 'open_app':
                    app_name = action.app_name if hasattr(action, 'app_name') else "app"
                    description = f"Open {app_name} app"
                elif action_type == 'navigate_home':
                    description = "Press Home button to go to home screen"
                elif action_type == 'navigate_back':
                    description = "Press Back button to go back"
                else:
                    description = f"Perform {action_type} action"
            
            action_dict = {
                'action': action_type,
                'description': description
            }
            
            # 根据Action类型添加参数
            if hasattr(action, 'text') and action.text:
                action_dict['text'] = action.text
            if hasattr(action, 'x') and hasattr(action, 'y'):
                action_dict['coordinate'] = [action.x, action.y]
            if hasattr(action, 'direction') and action.direction:
                action_dict['direction'] = action.direction
            if hasattr(action, 'app_name') and action.app_name:
                action_dict['app_name'] = action.app_name
        else:
            # 已经是字典
            action_dict = action
        
        # Execute action
        new_screenshot, metadata = self._virtual_env.step(action_dict)
        self._current_screenshot = new_screenshot
    
    def execute_adb_call(self, call: list[str], timeout: float = 60.0):
        """
        模拟执行ADB调用（虚拟环境不支持真实ADB）
        
        Args:
            call: ADB命令列表
            timeout: 超时时间
            
        Returns:
            模拟的ADB响应
        """
        print(f"[VirtualEnvAdapter] Mock ADB call: {' '.join(call)}")
        
        # 返回模拟响应
        if 'whoami' in call:
            return type('obj', (object,), {'generic': 'shell'})()
        elif 'settings' in call and 'put' in call and 'global' in call:
            return type('obj', (object,), {'generic': ''})()
        else:
            return type('obj', (object,), {'generic': ''})()
    
    def close(self):
        """Close environment"""
        print(f"[VirtualEnvAdapter] Closing environment")
        # 关闭虚拟环境，保存轨迹文件
        if hasattr(self._virtual_env, 'close'):
            self._virtual_env.close()
        pass
    
    def ask_question(self, question: str, timeout_seconds: float = -1.0) -> str | None:
        """
        向用户提问（虚拟环境中返回None）
        
        Args:
            question: 问题
            timeout_seconds: 超时时间
            
        Returns:
            None（虚拟环境不支持用户交互）
        """
        print(f"[VirtualEnvAdapter] Ask question (not supported in virtual env): {question}")
        return None
    
    @property
    def foreground_activity_name(self) -> str:
        """返回前台App name"""
        state = self._virtual_env.get_current_state()
        return state.get('app_name', 'Unknown')
    
    @property
    def device_screen_size(self) -> tuple[int, int]:
        """返回设备屏幕尺寸"""
        return self._virtual_env.resolution
    
    @property
    def logical_screen_size(self) -> tuple[int, int]:
        """Return logical screen size"""
        return self._virtual_env.resolution
    
    @property
    def interaction_cache(self) -> str:
        """返回交互缓存"""
        return self._interaction_cache
    
    @interaction_cache.setter
    def interaction_cache(self, value: str) -> None:
        """设置交互缓存"""
        self._interaction_cache = value
    
    def hide_automation_ui(self) -> None:
        """隐藏自动化UI（虚拟环境不需要）"""
        pass
    
    @property
    def orientation(self) -> int:
        """
        返回屏幕方向
        
        Returns:
            0 (竖屏)
        """
        return 0
    
    @property
    def physical_frame_boundary(self) -> tuple[int, int, int, int]:
        """
        返回物理边框边界
        
        Returns:
            (0, 0, width, height)
        """
        width, height = self._virtual_env.resolution
        return (0, 0, width, height)
    
    def set_screenshot_path(self, filepath: str):
        """
        设置当前截图 file path（在Agent保存截图后调用）
        
        Args:
            filepath: 截图文件的完整路径
        """
        self._virtual_env.set_screenshot_path(filepath)
    
    def set_trajectory_dir(self, trajectory_dir: str):
        """
        设置轨迹文件夹路径
        
        Args:
            trajectory_dir: 轨迹文件夹的完整路径
        """
        self._virtual_env.set_trajectory_dir(trajectory_dir)


class VirtualController:
    """
    虚拟控制器，模拟真实设备的控制器接口
    """
    
    def __init__(self, virtual_env: VirtualAndroidEnv):
        self._virtual_env = virtual_env
        self._last_action_result = None
    
    def get_screenshot(self) -> np.ndarray:
        """获取当前截图"""
        state = self._virtual_env.get_current_state()
        return state.get('screenshot')
    
    def execute_action(self, action: Dict[str, Any]) -> Dict[str, Any]:
        """
        Execute action
        
        Args:
            action: Action字典
            
        Returns:
            执行结果
        """
        new_screenshot, metadata = self._virtual_env.step(action)
        self._last_action_result = {
            'screenshot': new_screenshot,
            'metadata': metadata,
            'success': metadata.get('success', True)
        }
        return self._last_action_result
    
    def tap(self, x: int, y: int):
        """点击操作"""
        action = {
            'action': 'click',
            'coordinate': [x, y],
            'description': f'Tap at ({x}, {y})'
        }
        return self.execute_action(action)
    
    def swipe(self, x1: int, y1: int, x2: int, y2: int, duration: int = 300):
        """滑动操作"""
        action = {
            'action': 'swipe',
            'start': [x1, y1],
            'end': [x2, y2],
            'duration': duration,
            'description': f'Swipe from ({x1},{y1}) to ({x2},{y2})'
        }
        return self.execute_action(action)
    
    def input_text(self, text: str):
        """Input text"""
        action = {
            'action': 'input',
            'text': text,
            'description': f'Input text: {text}'
        }
        return self.execute_action(action)
    
    def press_key(self, keycode: str):
        """按键操作"""
        action = {
            'action': 'press_key',
            'key': keycode,
            'description': f'Press key: {keycode}'
        }
        return self.execute_action(action)
    
    def navigate_home(self):
        """返回主屏幕"""
        action = {
            'action': 'navigate_home',
            'description': 'Navigate to home screen'
        }
        return self.execute_action(action)
    
    def navigate_back(self):
        """返回上一屏"""
        action = {
            'action': 'navigate_back',
            'description': 'Navigate back'
        }
        return self.execute_action(action)
    
    def execute_adb_call(self, call, timeout: float = 60.0):
        """
        模拟执行ADB调用（虚拟环境不支持真实ADB）
        虚拟环境不需要真实的ADB操作，所有检查都返回成功
        
        Args:
            call: ADB命令
            timeout: 超时时间
            
        Returns:
            模拟的ADB响应
        """
        # 创建成功的ADB响应对象
        response = adb_pb2.AdbResponse()
        response.status = adb_pb2.AdbResponse.Status.OK
        
        # 根据命令类型返回适当的输出
        call_str = str(call)
        
        # 文件/目录检查 - 都返回存在
        if 'test -' in call_str or 'if [' in call_str:
            response.generic.output = b'Exists'
        # 其他命令返回空输出表示成功
        else:
            response.generic.output = b''
        
        return response