virtual_env_gemini3.py

"""
Gemini3 Virtual Android Environment Implementation
Uses Gemini generative AI to simulate Android device interfaces
"""

import os
import time
import json
import base64
import asyncio
from io import BytesIO
from typing import Optional, Dict, List, Tuple, Any
from dataclasses import dataclass, field

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import requests


@dataclass
class EnvironmentState:
    """Environment state"""
    screenshot: np.ndarray
    app_name: str = "Home"
    screen_description: str = ""
    timestamp: float = field(default_factory=time.time)
    action_count: int = 0
    screenshot_path: str = ""  # Actual saved screenshot file path


class StateManager:
    """Manage virtualEnvironment state"""
    
    def __init__(self):
        self.current_state: Optional[EnvironmentState] = None
        self.previous_state: Optional[EnvironmentState] = None
        self.action_history: List[Dict] = []
        self.state_history: List[EnvironmentState] = []
        self.max_history = 10
    
    def update_state(self, new_screenshot: np.ndarray, action: Dict, app_name: str = None, screenshot_path: str = ""):
        """Update state"""
        self.previous_state = self.current_state
        
        self.current_state = EnvironmentState(
            screenshot=new_screenshot,
            app_name=app_name or (self.current_state.app_name if self.current_state else "Home"),
            action_count=len(self.action_history) + 1,
            screenshot_path=screenshot_path
        )
        
        self.action_history.append(action)
        self.state_history.append(self.current_state)
        
        # Keep history within reasonable range
        if len(self.state_history) > self.max_history:
            self.state_history.pop(0)
    
    def get_context_for_generation(self) -> Dict:
        """Get context needed for generation"""
        return {
            "current_screenshot": self.current_state.screenshot if self.current_state else None,
            "current_app": self.current_state.app_name if self.current_state else "Home",
            "recent_actions": self.action_history[-5:],
            "action_count": len(self.action_history)
        }
    
    def get_previous_screenshot(self) -> Optional[np.ndarray]:
        """Get previous screenshot"""
        return self.previous_state.screenshot if self.previous_state else None
    
    def get_previous_screenshot_path(self) -> str:
        """Get previous screenshot file path"""
        return self.previous_state.screenshot_path if self.previous_state else ""
    
    def set_current_screenshot_path(self, path: str):
        """Set current screenshot file path (for updating after saving)"""
        if self.current_state:
            self.current_state.screenshot_path = path


class Gemini3Generator:
    """Gemini3 image generation wrapper - uses OpenRouter API"""
    
    # Prompt templates
    INITIAL_PROMPT_TEMPLATE = """Generate Android home screen (1080x2400px):
- Show app icons grid
- Status bar and navigation bar
- Modern realistic UI style"""

    ACTION_PROMPT_TEMPLATE = """You are simulating an Android phone. The image above shows the CURRENT screen.

Action performed: {action_description}

Generate the NEXT screenshot that shows what appears AFTER this action is executed.

Requirements:
- Resolution: 1080 x 2400 pixels
- Based on the current screen shown above, show the natural result of the action
- Use realistic Android UI elements
- If clicking a button/option, show what that button/option leads to
- If opening an app, show that app's interface
- If typing text, show the text in the input field

Generate only the resulting screenshot."""

    
    def __init__(
        self, 
        api_key: str,
        base_url: str = "https://openrouter.ai/api/v1",
        model: str = "google/gemini-2.0-flash-exp:free",
        max_retry: int = 3,
        timeout: int = 60
    ):
        """
        Initialize Gemini3 generator
        
        Args:
            api_key: OpenRouter API key
            base_url: OpenRouter API address
            model: Model name
            max_retry: Maximum retry attempts
            timeout: Request timeout (seconds)
        """
        self.api_key = api_key
        self.base_url = base_url.rstrip('/')
        self.model = model
        self.max_retry = max_retry
        self.timeout = timeout
        self.generation_history: List[Dict] = []
        
        print(f"[Gemini3Generator] Initialized with model: {model}")
        print(f"[Gemini3Generator] API endpoint: {base_url}")
    
    def generate_screenshot(
        self, 
        prompt: str, 
        reference_image: Optional[np.ndarray] = None,
        task_description: str = "",
        reference_image_path: str = ""
    ) -> Tuple[np.ndarray, Dict]:
        """
        Generate Android screenshot
        
        Args:
            prompt: Generation prompt
            reference_image: Reference image (previous screenshot, for consistency)
            task_description: Task description
            reference_image_path: Reference image file path (for logging)
            
        Returns:
            (screenshot, metadata): Generated screenshot and metadata
        """
        start_time = time.time()
        
        try:
            print(f"\n[Gemini3Generator] === Generating Screenshot ===")
            print(f"[Gemini3Generator] Has reference image: {reference_image is not None}")
            if reference_image is not None:
                if reference_image_path:
                    print(f"[Gemini3Generator] Actual reference image file: {reference_image_path}")
                else:
                    print(f"[Gemini3Generator] [WARNING]  Reference image not saved (memory array)")
                print(f"[Gemini3Generator] Reference image shape: {reference_image.shape}")
                print(f"[Gemini3Generator] Reference image dtype: {reference_image.dtype}")
                print(f"[Gemini3Generator] Reference image min/max: [{reference_image.min()}, {reference_image.max()}]")
            print(f"[Gemini3Generator] Prompt length: {len(prompt)} chars")
            print(f"[Gemini3Generator] === Prompt Content ===")
            print(prompt)
            print(f"[Gemini3Generator] === End of Prompt ===")
            
            # Build API request
            messages = self._build_messages(prompt, reference_image)
            
            # Call OpenRouter API
            screenshot = self._call_api(messages)
            
            generation_time = time.time() - start_time
            
            # Record history
            metadata = {
                "prompt": prompt[:200] + "..." if len(prompt) > 200 else prompt,
                "generation_time": generation_time,
                "timestamp": time.time(),
                "success": True,
                "model": self.model,
                "has_reference": reference_image is not None
            }
            self.generation_history.append(metadata)
            
            print(f"[Gemini3Generator] [OK] Generation successful in {generation_time:.2f}s")
            print(f"[Gemini3Generator] Screenshot shape: {screenshot.shape}")
            
            return screenshot, metadata
            
        except Exception as e:
            generation_time = time.time() - start_time
            print(f"[Gemini3Generator] [FAILED] Generation failed: {e}")
            
            metadata = {
                "prompt": prompt[:200],
                "generation_time": generation_time,
                "timestamp": time.time(),
                "success": False,
                "error": str(e),
                "model": self.model
            }
            self.generation_history.append(metadata)
            
            # Throw exception directly, don't return placeholder
            raise Exception(f"Image generation failed: {e}")
    
    def _build_messages(self, prompt: str, reference_image: Optional[np.ndarray] = None) -> List[Dict]:
        """Build OpenRouter API message format"""
        content = []
        
        # Add reference image (if available)
        if reference_image is not None:
            img_base64 = self._numpy_to_base64(reference_image)
            content.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{img_base64}"
                }
            })
            content.append({
                "type": "text",
                "text": "Above is the CURRENT screenshot. "
            })
        
        # Add text prompt
        content.append({
            "type": "text",
            "text": prompt
        })
        
        messages = [
            {
                "role": "user",
                "content": content
            }
        ]
        
        return messages
    
    def _call_api(self, messages: List[Dict]) -> np.ndarray:
        """Call OpenRouter API"""
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://github.com/mobile-agent",
            "X-Title": "Mobile-Agent Virtual Environment"
        }
        
        payload = {
            "model": self.model,
            "messages": messages,
            "max_tokens": 2000,
            "temperature": 0.7,
        }
        
        # Simplified logging: only show prompt content (not detailed message structure)
        # print(f"[Gemini3Generator] === Sending Request ===")
        
        retry_count = 0
        last_error = None
        
        while retry_count < self.max_retry:
            try:
                print(f"\n[Gemini3Generator] API call attempt {retry_count + 1}/{self.max_retry}")
                
                response = requests.post(
                    f"{self.base_url}/chat/completions",
                    headers=headers,
                    json=payload,
                    timeout=self.timeout
                )
                
                if response.status_code == 200:
                    result = response.json()
                    message = result['choices'][0]['message']
                    content = message.get('content', '')
                    
                    # Simplified logging: only show key information
                    print(f"[Gemini3Generator] [OK] Response received (Status: {response.status_code})")
                    
                    # content may be list (containing image_url) or string
                    if isinstance(content, list):
                        for idx, item in enumerate(content):
                            if isinstance(item, dict):
                                if 'type' in item and item['type'] == 'image_url':
                                    image_url_data = item.get('image_url', {})
                                    if isinstance(image_url_data, dict) and 'url' in image_url_data:
                                        url = image_url_data['url']
                                        if url.startswith('data:image'):
                                            base64_data = url.split(',')[1]
                                            screenshot = self._base64_to_numpy(base64_data)
                                            print(f"[Gemini3Generator] [OK] Image decoded! Shape: {screenshot.shape}")
                                            return screenshot
                                        elif url.startswith('http'):
                                            img_response = requests.get(url, timeout=60)
                                            img = Image.open(BytesIO(img_response.content))
                                            screenshot = np.array(img)
                                            print(f"[Gemini3Generator] [OK] Image downloaded! Shape: {screenshot.shape}")
                                            return screenshot
                    
                    # Check images field (correct field for Gemini image generation model)
                    if 'images' in message and message['images']:
                        images = message['images']
                        first_image = images[0]
                        
                        if 'image_url' in first_image:
                            image_url_obj = first_image['image_url']
                            if isinstance(image_url_obj, dict) and 'url' in image_url_obj:
                                url = image_url_obj['url']
                                
                                if url.startswith('data:image'):
                                    base64_data = url.split(',', 1)[1]
                                    screenshot = self._base64_to_numpy(base64_data)
                                    print(f"[Gemini3Generator] [OK] Image decoded! Shape: {screenshot.shape}")
                                    return screenshot
                                elif url.startswith('http'):
                                    img_response = requests.get(url, timeout=60)
                                    img = Image.open(BytesIO(img_response.content))
                                    screenshot = np.array(img)
                                    print(f"[Gemini3Generator] [OK] Image downloaded! Shape: {screenshot.shape}")
                                    return screenshot
                    
                    # Check content field (alternative)
                    if isinstance(content, str) and content.startswith('data:image'):
                        screenshot = self._base64_to_numpy(content.split(',')[1])
                        print(f"[Gemini3Generator] [OK] Image decoded! Shape: {screenshot.shape}")
                        return screenshot
                    elif isinstance(content, dict) and 'url' in content:
                        img_response = requests.get(content['url'], timeout=30)
                        img = Image.open(BytesIO(img_response.content))
                        screenshot = np.array(img)
                        print(f"[Gemini3Generator] [OK] Image downloaded! Shape: {screenshot.shape}")
                        return screenshot
                    
                    # No image data found
                    print("[Gemini3Generator] [WARNING] No image data found in response!")
                    raise Exception(f"No image data found in API response. Model may not support image generation.")
                    
                else:
                    error_msg = f"API returned status {response.status_code}: {response.text}"
                    print(f"[Gemini3Generator] {error_msg}")
                    last_error = error_msg
                    
            except requests.exceptions.Timeout:
                last_error = "Request timeout"
                print(f"[Gemini3Generator] Timeout on attempt {retry_count + 1}")
                
            except Exception as e:
                last_error = str(e)
                print(f"[Gemini3Generator] Error on attempt {retry_count + 1}: {e}")
            
            retry_count += 1
            if retry_count < self.max_retry:
                wait_time = 2 ** retry_count  # Exponential backoff
                print(f"[Gemini3Generator] Retrying in {wait_time}s...")
                time.sleep(wait_time)
        
        raise Exception(f"Failed after {self.max_retry} attempts. Last error: {last_error}")
    
    def _numpy_to_base64(self, img: np.ndarray) -> str:
        """Convert numpy array to base64 string"""
        # Ensure uint8 format
        if img.dtype != np.uint8:
            img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8)
        
        pil_img = Image.fromarray(img)
        buffer = BytesIO()
        pil_img.save(buffer, format='PNG')
        img_base64 = base64.b64encode(buffer.getvalue()).decode()
        return img_base64
    
    def _base64_to_numpy(self, img_base64: str) -> np.ndarray:
        """Convert base64 string to numpy array"""
        img_data = base64.b64decode(img_base64)
        img = Image.open(BytesIO(img_data))
        return np.array(img)
    


class VirtualAndroidEnv:
    """Virtual Android environment based on Gemini3"""
    
    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 = "",
        resolution: Tuple[int, int] = (1080, 2400),
        initial_image_path: str = None
    ):
        """
        Initialize virtual Android environment
        
        Args:
            gemini_api_key: Gemini API key（OpenRouter）
            gemini_base_url: API address
            gemini_model: Model name
            initial_task: Initial task description
            resolution: Screen resolution (width, height)
            initial_image_path: Initial image path, skip generation if provided
        """
        self.task_description = initial_task
        self.resolution = resolution
        
        # Initial image path (skip initial generation if provided)
        if initial_image_path is None:
            # Default to start_Android.png
            script_dir = os.path.dirname(os.path.abspath(__file__))
            self.initial_image_path = os.path.join(script_dir, "start_Android.png")
        else:
            # Ensure absolute path
            if os.path.isabs(initial_image_path):
                self.initial_image_path = initial_image_path
            else:
                # Convert relative path to absolute path
                self.initial_image_path = os.path.abspath(initial_image_path)
        
        # Initialize generator and state manager
        self.generator = Gemini3Generator(
            api_key=gemini_api_key,
            base_url=gemini_base_url,
            model=gemini_model
        )
        self.state_manager = StateManager()
        
        # Pending state (update after Agent saves file)
        self._pending_screenshot = None
        self._pending_action = None
        self._pending_app = None
        
        # Trajectory folder path (for reading last screenshot)
        self.trajectory_dir = None
        self.current_step = 0  # Current step count
        
        # Real-time trajectory file path
        self.trajectory_file = None
        self.trajectory_file_handle = None
        
        print(f"[VirtualAndroidEnv] Initialized")
        print(f"[VirtualAndroidEnv] Task: {initial_task}")
        print(f"[VirtualAndroidEnv] Resolution: {resolution}")
        print(f"[VirtualAndroidEnv] Initial image: {self.initial_image_path}")
    
    def reset(self, task_description: Optional[str] = None) -> np.ndarray:
        """
        Reset environment to initial state
        
        Args:
            task_description: Task description, use initial task if not provided
            
        Returns:
            initial_screenshot: Initial screenshot
        """
        if task_description:
            self.task_description = task_description
        
        print(f"\n[VirtualAndroidEnv] Resetting environment...")
        print(f"[VirtualAndroidEnv] Task: {self.task_description}")
        
        # Load initial image directly, no generation
        if os.path.exists(self.initial_image_path):
            print(f"[VirtualAndroidEnv] Loading initial image from: {self.initial_image_path}")
            initial_img = Image.open(self.initial_image_path)
            
            # Convert to RGB format (remove alpha channel if exists)
            if initial_img.mode == 'RGBA':
                # Create white background
                background = Image.new('RGB', initial_img.size, (255, 255, 255))
                background.paste(initial_img, mask=initial_img.split()[3])  # Use alpha channel as mask
                initial_img = background
            elif initial_img.mode != 'RGB':
                initial_img = initial_img.convert('RGB')
            
            initial_screenshot = np.array(initial_img)
            print(f"[VirtualAndroidEnv] [OK] Loaded initial image, shape: {initial_screenshot.shape}")
        else:
            print(f"[VirtualAndroidEnv] [WARNING]  Initial image not found: {self.initial_image_path}")
            print(f"[VirtualAndroidEnv] Generating initial screenshot instead...")
            
            # Fallback: generateInitial screenshot
            prompt = self.generator.INITIAL_PROMPT_TEMPLATE
            
            initial_screenshot, metadata = self.generator.generate_screenshot(
                prompt=prompt,
                reference_image=None,
                task_description=self.task_description
            )
        
        # Submit initial state directly, use initial image file path
        initial_path = self.initial_image_path if os.path.exists(self.initial_image_path) else ""
        self.state_manager.update_state(
            new_screenshot=initial_screenshot,
            action={"action": "reset", "description": "Initialize environment"},
            app_name="Home",
            screenshot_path=initial_path  # Use initial image file path
        )
        
        # Clear pending (if any)
        self._pending_screenshot = None
        self._pending_action = None
        self._pending_app = None
        
        print(f"[VirtualAndroidEnv] Environment reset complete")
        if initial_path:
            print(f"[VirtualAndroidEnv] Initial screenshot path: {initial_path}")
        
        return initial_screenshot
    
    def _get_last_screenshot_from_trajectory(self) -> Tuple[Optional[np.ndarray], str]:
        """
        Read last screenshot from trajectory folder
        
        Returns:
            (screenshot, filepath): Screenshot array and file path, return if not found (None, "")
        """
        if not self.trajectory_dir or not os.path.exists(self.trajectory_dir):
            return None, ""
        
        # Find all screenshot_*.png files
        import glob
        import re
        pattern = os.path.join(self.trajectory_dir, "screenshot_*.png")
        screenshot_files = glob.glob(pattern)
        
        # Sort by numeric index in filename (e.g., screenshot_0.png, screenshot_1.png, ...)
        # Extract number from filename and sort numerically
        def extract_number(filepath):
            match = re.search(r'screenshot_(\d+)\.png', os.path.basename(filepath))
            return int(match.group(1)) if match else -1
        
        screenshot_files = sorted(screenshot_files, key=extract_number)
        
        if not screenshot_files:
            # No screenshots found, try using initial image
            if os.path.exists(self.initial_image_path):
                print(f"[VirtualAndroidEnv] No screenshots in trajectory, using initial image")
                img = Image.open(self.initial_image_path)
                if img.mode == 'RGBA':
                    background = Image.new('RGB', img.size, (255, 255, 255))
                    background.paste(img, mask=img.split()[3])
                    img = background
                elif img.mode != 'RGB':
                    img = img.convert('RGB')
                return np.array(img), self.initial_image_path
            return None, ""
        
        # Get last screenshot
        last_screenshot_path = screenshot_files[-1]
        print(f"[VirtualAndroidEnv] Loading last screenshot: {last_screenshot_path}")
        
        try:
            img = Image.open(last_screenshot_path)
            if img.mode == 'RGBA':
                background = Image.new('RGB', img.size, (255, 255, 255))
                background.paste(img, mask=img.split()[3])
                img = background
            elif img.mode != 'RGB':
                img = img.convert('RGB')
            screenshot = np.array(img)
            return screenshot, last_screenshot_path
        except Exception as e:
            print(f"[VirtualAndroidEnv] [FAILED] Failed to load screenshot: {e}")
            return None, ""
    
    def set_trajectory_dir(self, trajectory_dir: str):
        """
        Set trajectory directory (where Agent saves screenshots)
        
        Args:
            trajectory_dir: Trajectory directory path
        """
        self.trajectory_dir = trajectory_dir
        print(f"[VirtualAndroidEnv] Trajectory directory set to: {trajectory_dir}")
        
        # Create real-time trajectory file
        if trajectory_dir and os.path.exists(trajectory_dir):
            self.trajectory_file = os.path.join(trajectory_dir, "traj.jsonl")
            # Open file for appending
            self.trajectory_file_handle = open(self.trajectory_file, 'a', encoding='utf-8')
            print(f"[VirtualAndroidEnv] Real-time trajectory file: {self.trajectory_file}")
    
    def step(self, action: Dict) -> Tuple[np.ndarray, Dict]:
        """
        Execute an action, get new screenshot
        
        Args:
            action: Action dict containing type and parameters
            
        Returns:
            (screenshot, metadata): New screenshot and metadata
        """
        print(f"\n[VirtualAndroidEnv] Executing action: {action.get('action', 'unknown')}")
        print(f"[VirtualAndroidEnv] Action description: {action.get('description', 'N/A')}")
        
        # Build action prompt
        prompt = self._build_action_prompt(action)
        
        # Read last screenshot from trajectory folder as reference
        reference_image, reference_image_path = self._get_last_screenshot_from_trajectory()
        
        if reference_image is None and self.state_manager.current_state:
            print(f"[VirtualAndroidEnv] [WARNING]  No previous screenshot, using current screenshot as reference")
            reference_image = self.state_manager.current_state.screenshot
            reference_image_path = self.state_manager.current_state.screenshot_path
        
        if reference_image is not None:
            print(f"[VirtualAndroidEnv] [OK] Using previous screenshot as reference")
            if reference_image_path:
                print(f"[VirtualAndroidEnv]   Actual reference image: {reference_image_path}")
            else:
                print(f"[VirtualAndroidEnv]   [WARNING]  Reference image not saved to file (numpy array in memory)")
            print(f"[VirtualAndroidEnv]   Will generate next step screenshot")
        else:
            print(f"[VirtualAndroidEnv] [WARNING]  No reference image available")
            reference_image_path = ""
        
        # Generate new screenshot
        new_screenshot, metadata = self.generator.generate_screenshot(
            prompt=prompt,
            reference_image=reference_image,
            task_description=self.task_description,
            reference_image_path=reference_image_path
        )
        
        # Infer new app name
        new_app = self._infer_app_name(action)
        
        # Cache new screenshot and metadata, update state after Agent saves file
        self._pending_screenshot = new_screenshot
        self._pending_action = action
        self._pending_app = new_app
        
        metadata['app_name'] = new_app
        metadata['step_number'] = len(self.state_manager.action_history) + 1  # +1Because not yet actually updated
        
        # Save trajectory step in real-time
        # 注释掉：轨迹保存已由mobile_agent_v3.py统一处理（Manager-Operator-Reflector格式）
        # self._save_trajectory_step(action, metadata)
        
        return new_screenshot, metadata
    
    def _save_trajectory_step(self, action: Dict, metadata: Dict):
        """
        Save trajectory step in real-time到jsonl文件
        
        Args:
            action: Executed action
            metadata: Metadata
        """
        if not self.trajectory_file_handle:
            return
        
        try:
            step_number = metadata.get('step_number', 0)
            
            # Build trajectory step record
            step_record = {
                "step": step_number,
                "action": action.get('action', 'unknown'),
                "action_type": action.get('action', 'unknown'),
                "description": action.get('description', ''),
                "app_name": metadata.get('app_name', 'Unknown'),
                "timestamp": time.time()
            }
            
            # Add other action parameters
            for key, value in action.items():
                if key not in ['action', 'description']:
                    step_record[key] = value
            
            # Write to jsonl file
            self.trajectory_file_handle.write(json.dumps(step_record, ensure_ascii=False) + '\n')
            self.trajectory_file_handle.flush()  # Flush to disk immediately
            
        except Exception as e:
            print(f"[VirtualAndroidEnv] [WARNING] Failed to save trajectory step: {e}")
    
    def _build_action_prompt(self, action: Dict) -> str:
        """Build prompt for action execution"""
        # 直接使用action description
        action_description = action.get('description', 'Perform an action')
        
        # 使用简化的模板
        prompt = self.generator.ACTION_PROMPT_TEMPLATE.format(
            action_description=action_description
        )
        
        return prompt
    
    def _infer_app_name(self, action: Dict) -> str:
        """根据Action推断当前App name"""
        action_type = action.get("action", "")
        
        if action_type == "open_app":
            return action.get("text", "Unknown App")
        
        elif action_type == "system_button":
            button = action.get("button", "")
            if button.lower() == "home":
                return "Home"
        
        # 否则保持当前应用
        if self.state_manager.current_state:
            return self.state_manager.current_state.app_name
        
        return "Unknown"
    
    def get_current_state(self) -> Dict:
        """
        获取当前Environment state描述
        
        Returns:
            状态信息字典
        """
        # 如果有待处理的截图，返回其信息
        if hasattr(self, '_pending_screenshot') and self._pending_screenshot is not None:
            return {
                "app_name": self._pending_app or "Unknown",
                "action_count": len(self.state_manager.action_history),
                "action_history": self.state_manager.action_history[-5:] if self.state_manager.action_history else [],
                "task": self.task_description,
                "has_screenshot": True,
                "screenshot": self._pending_screenshot
            }
        
        # 否则返回已确认的状态
        context = self.state_manager.get_context_for_generation()
        
        return {
            "app_name": context['current_app'],
            "action_count": context['action_count'],
            "action_history": context['recent_actions'],
            "task": self.task_description,
            "has_screenshot": context['current_screenshot'] is not None
        }
    
    def save_screenshot(self, filepath: str, screenshot: Optional[np.ndarray] = None):
        """Save screenshot to file"""
        if screenshot is None:
            # 优先使用 pending 状态的截图
            if hasattr(self, '_pending_screenshot') and self._pending_screenshot is not None:
                screenshot = self._pending_screenshot
            elif self.state_manager.current_state:
                screenshot = self.state_manager.current_state.screenshot
            else:
                print("[VirtualAndroidEnv] No screenshot to save")
                return
        
        img = Image.fromarray(screenshot.astype('uint8'))
        img.save(filepath)
        print(f"[VirtualAndroidEnv] Screenshot saved to: {filepath}")
        
        # 如果有待处理的状态，现在更新它
        if hasattr(self, '_pending_screenshot') and self._pending_screenshot is not None:
            self.set_screenshot_path(filepath)
        else:
            # 否则只更新路径
            self.state_manager.set_current_screenshot_path(filepath)
    
    def set_screenshot_path(self, filepath: str):
        """
        设置当前截图 file path（用于外部保存截图后更新）
        此时才真正Update state，确保 screenshot_path 已知
        """
        # 如果有待更新的状态，现在进行更新
        if hasattr(self, '_pending_screenshot') and self._pending_screenshot is not None:
            self.state_manager.update_state(
                new_screenshot=self._pending_screenshot,
                action=self._pending_action,
                app_name=self._pending_app,
                screenshot_path=filepath
            )
            print(f"[VirtualAndroidEnv] [OK] State updated with screenshot: {filepath}")
            # 清除待处理数据
            self._pending_screenshot = None
            self._pending_action = None
            self._pending_app = None
        else:
            # 兼容性：如果没有待处理的状态，只更新路径
            self.state_manager.set_current_screenshot_path(filepath)
            print(f"[VirtualAndroidEnv] Screenshot path updated to: {filepath}")
    
    def get_generation_history(self) -> List[Dict]:
        """获取生成历史记录"""
        return self.generator.generation_history
    
    def close(self):
        """Close environment，释放资源"""
        if self.trajectory_file_handle:
            try:
                self.trajectory_file_handle.close()
                print(f"[VirtualAndroidEnv] [OK] Trajectory file closed: {self.trajectory_file}")
            except Exception as e:
                print(f"[VirtualAndroidEnv] [WARNING] Error closing trajectory file: {e}")
            finally:
                self.trajectory_file_handle = None
    
    def __del__(self):
        """析构函数，确保文件句柄被关闭"""
        self.close()


if __name__ == "__main__":
    # 测试代码
    print("="*60)
    print("Testing VirtualAndroidEnv")
    print("="*60)
    
    # Configuration (set OPENROUTER_API_KEY environment variable)
    API_KEY = os.getenv("OPENROUTER_API_KEY")
    if not API_KEY:
        raise ValueError("OPENROUTER_API_KEY environment variable must be set")
    BASE_URL = "https://api.dou.chat/v1"
    MODEL = "google/gemini-3-pro-image-preview"
    
    # 创建虚拟环境
    env = VirtualAndroidEnv(
        gemini_api_key=API_KEY,
        gemini_base_url=BASE_URL,
        gemini_model=MODEL,
        initial_task="打开设置并关闭WiFi"
    )
    
    # Reset environment
    initial_screenshot = env.reset()
    print(f"\nInitial screenshot shape: {initial_screenshot.shape}")
    
    # 保存Initial screenshot
    os.makedirs("./test_outputs", exist_ok=True)
    env.save_screenshot("./test_outputs/step_0_initial.png")
    
    # 执行一些测试Action
    test_actions = [
        {
            "action": "open_app",
            "text": "settings",
            "description": "打开设置应用"
        },
        {
            "action": "click",
            "coordinate": [540, 600],
            "description": "点击WiFi选项"
        },
        {
            "action": "click",
            "coordinate": [950, 200],
            "description": "点击WiFi开关"
        }
    ]
    
    for i, action in enumerate(test_actions):
        print(f"\n{'='*60}")
        print(f"Step {i+1}: {action['description']}")
        print(f"{'='*60}")
        
        screenshot, metadata = env.step(action)
        
        print(f"New screenshot shape: {screenshot.shape}")
        print(f"Generation time: {metadata.get('generation_time', 0):.2f}s")
        print(f"Current app: {metadata.get('app_name', 'Unknown')}")
        
        # 保存截图
        env.save_screenshot(f"./test_outputs/step_{i+1}_{action['action']}.png")
    
    # 显示状态
    print(f"\n{'='*60}")
    print("Current State:")
    print(f"{'='*60}")
    state = env.get_current_state()
    for key, value in state.items():
        if key != "action_history":
            print(f"{key}: {value}")
    
    print("\n[OK] Test completed!")