utils.py

import logging
import sys
import threading
import os
import requests
import subprocess
import glob
import fcntl
import json
from datetime import datetime
import time
from appium.webdriver.common.appiumby import AppiumBy
from selenium.common.exceptions import NoSuchElementException, TimeoutException


# --------------------
# Configuration loading
# --------------------

def load_config():
    """
    Load configuration file
    """
    config_file = 'config.json'

    # If config file doesn't exist, create a sample file
    if not os.path.exists(config_file):
        sample_config = {
            "wechat_sendkey": "Please fill in your SendKey here",
            "other_settings": {
                "debug": False,
                "max_retries": 3
            }
        }
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(sample_config, f, indent=4, ensure_ascii=False)
        print(f"Created config file template: {config_file}")
        print("Please edit the config file and fill in the correct key")
        return None

    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config
    except Exception as e:
        print(f"Failed to read config file: {e}")
        return None

def get_wechat_sendkey():
    """
    Get WeChat SendKey
    """
    # First try to get from environment variable
    key = os.getenv('WECHAT_SENDKEY')
    if key:
        return key

    # Then get from config file
    config = load_config()
    if config and config.get('wechat_sendkey') != "Please fill in your SendKey here":
        return config.get('wechat_sendkey')

    return None


def login_to_app(driver, app_name, logger):
    """
    Simulate login to app, automatically handle popups
    """
    logger.info(f"Logging in to app: {app_name}")
    
    def handle_popups(max_attempts=5):
        """
        Automatically handle various popups
        """
        popup_buttons = [
            # Allow type
            "允许", "Allow", "ALLOW", "允许一次", "Allow once", "允许所有时间", "Allow all the time",
            # Confirm type
            "确定", "OK", "确认", "Confirm", "好的", "Got it", "我知道了", "知道了", "Done",
            # Agree type
            "同意", "Agree", "接受", "Accept", "我同意", "I agree", "Got it", "Accept & Continue", "Agree & Continue",
            # Continue type
            "继续", "Continue", "下一步", "Next", "开始", "Start", "开始使用", "Get started", "CONTINUE",
            # Close type
            "关闭", "Close", "取消", "Cancel", "跳过", "Skip", "稍后", "Later", "不用了", "No thanks", "CLOSE"
            # Other common buttons
            "立即体验", "马上体验", "立即开始", "现在开始", "开启", "Enable", "Get started", "立即使用", "马上使用", "While using the app", "Finish", "Grant",
        ]
        
        for attempt in range(max_attempts):
            try:
                # Prioritize handling "Allow" and "OK" type buttons
                for button_text in popup_buttons:
                    try:
                        # Try multiple locating methods
                        selectors = [
                            f"//*[@text='{button_text}']",
                            f"//*[@content-desc='{button_text}']",
                            f"//*[contains(@text, '{button_text}')]",
                            f"//*[@resource-id and contains(@text, '{button_text}')]",
                            # Try clicking button's parent element
                            f"//*[@text='{button_text}']",
                            f"//*[@text='{button_text}']/ancestor::*[@clickable='true'][1]"
                        ]
                        
                        for selector in selectors:
                            try:
                                elements = driver.find_elements(AppiumBy.XPATH, selector)
                                for element in elements:
                                    try:
                                        if element.is_displayed() and element.is_enabled():
                                            element.click()
                                            logger.info(f"Successfully clicked popup button: {button_text}")
                                            clicked_any = True
                                            time.sleep(0.3)
                                    except Exception as e:
                                        logger.debug(f"Failed to click button: {e}")
                                        continue
                                        
                            except Exception:
                                continue
                        
                        if clicked_any:
                            time.sleep(1)
                            continue
                        else:
                            break
                    except Exception:
                        continue
                
                # If no text button found, try finding common dialog buttons via class
                try:
                    dialog_buttons = driver.find_elements(AppiumBy.XPATH,
                        "//android.widget.Button | //android.widget.TextView[@clickable='true']")

                    for button in dialog_buttons:
                        try:
                            button_text = button.get_attribute('text')
                            if button_text and any(keyword in button_text for keyword in
                                ['允许', 'Allow', '确定', 'OK', '同意', 'Agree', '继续', 'Continue']):
                                button.click()
                                logger.info(f"Found and clicked via class: {button_text}")
                                time.sleep(1)
                                return True
                        except Exception:
                            continue
                except Exception:
                    pass

                # Wait before retrying
                time.sleep(1)
                
            except Exception as e:
                logger.debug(f"Popup handling attempt {attempt+1} failed: {e}")

        return False

    def smart_click_element(xpath_or_id, element_desc="element"):
        """
        Smart click element with retry and popup handling
        """
        max_retries = 3
        for retry in range(max_retries):
            try:
                # First handle possible popups
                handle_popups(2)

                # Try locating element
                if xpath_or_id.startswith("//") or xpath_or_id.startswith("//*"):
                    element = driver.find_element(AppiumBy.XPATH, xpath_or_id)
                else:
                    element = driver.find_element(AppiumBy.ID, xpath_or_id)
                
                if element.is_displayed() and element.is_enabled():
                    element.click()
                    logger.info(f"Successfully clicked {element_desc}")
                    # Handle possible popups after clicking
                    time.sleep(1)
                    handle_popups(2)
                    return True
                else:
                    logger.warning(f"{element_desc} is not visible or clickable")

            except NoSuchElementException:
                logger.debug(f"Attempt {retry+1}: {element_desc} not found")
                # Retry after handling popups
                handle_popups(1)
                time.sleep(1)
            except Exception as e:
                logger.warning(f"Attempt {retry+1} clicking {element_desc} failed: {e}")
                time.sleep(1)
                
        return False
    
    def horizontal_swipe(times=1):
        # Get screen size
        screen_size = driver.get_window_size()
        width = screen_size['width']
        height = screen_size['height']

        # Set swipe parameters
        start_x = int(width * 0.8)
        end_x = int(width * 0.2)
        y = int(height * 0.5)
        duration = 300  # Swipe duration (milliseconds)
        
        for i in range(times):
            driver.swipe(start_x, y, end_x, y, duration)
            time.sleep(0.5)
    
    def vertical_swipe(times=1):
        # Get screen size
        screen_size = driver.get_window_size()
        width = screen_size['width']
        height = screen_size['height']

        # Set swipe parameters
        start_y = int(height * 0.8)
        end_y = int(height * 0.2)
        x = int(width * 0.5)
        duration = 300
        
        for i in range(times):
            driver.swipe(x, start_y, x, end_y, duration)
            time.sleep(0.5)
    
    # Start login process
    try:
        handle_popups(2)

        if app_name == "firefox":
            # smart_click_element("//*[@text='Agree and continue']", "Firefox同意按钮")
            smart_click_element("//*[@text='Not now']", "拒绝widget按钮")
            smart_click_element("//*[@text='Not now']", "拒绝encryption按钮")
            smart_click_element("//*[@text='Save and continue']", "theme按钮")
            smart_click_element("//*[@text='Save and continue']", "toolbar按钮")
                
        elif app_name == "Gadgetbridge":
            horizontal_swipe(4)
            smart_click_element("//*[@text='Go to the app']", "Gadgetbridge进入app按钮")
        
        elif app_name == "k9mail":
            # handle_popups(3)
            email_edittext = driver.find_element(
                AppiumBy.XPATH, 
                "//android.widget.EditText[.//android.widget.TextView[@text='Email address']]"
            )
            email_edittext.send_keys("testuser@qq.com")
            handle_popups(2)
            time.sleep(10)
            password_edittext = driver.find_element(
                AppiumBy.XPATH, 
                "//android.widget.EditText[.//android.widget.TextView[@text='Password']]"
            )
            password_edittext.send_keys("abc123456")
            handle_popups(2)
            time.sleep(5)
            handle_popups(2)
            name_edittext = driver.find_element(
                AppiumBy.XPATH, 
                "//android.widget.EditText[.//android.widget.TextView[@text='Your name*']]"
            )
            name_edittext.send_keys("Test User")
            handle_popups(2)
            time.sleep(2)
            handle_popups(2)
            time.sleep(5)
            handle_popups(1)
            time.sleep(1)
            handle_popups(1)
            time.sleep(1)
            smart_click_element("//android.widget.TextView[@text='Next']", "最后同意按钮")
            
        elif app_name == "myexpenses":
            horizontal_swipe(2)
            handle_popups(3)
        elif app_name == "redreader":
            smart_click_element("//*[@text='Accept']", "接受") 
            smart_click_element("//*[@text='Be anonymous']", "不登录")        
        elif app_name == "neurolab": 
            handle_popups(2)
            time.sleep(2)
            horizontal_swipe(2)
            smart_click_element("//android.widget.ImageButton", "√")
        elif app_name == "selfprivacy":
            for i in range(3):
                handle_popups(2)
                time.sleep(1)
            vertical_swipe(1)
            smart_click_element("//*[@content-desc='Skip to setup later']", "skip按钮")
        elif app_name == "money_manager_ex":
            smart_click_element("//*[@text='CLOSE']", "关闭欢迎页面")
            smart_click_element("//*[@content-desc='Navigate up']", "返回主页")
            smart_click_element("//*[@text='CREATE DATABASE']", "创建db")
            smart_click_element("//*[@text='SAVE']", "确定创建")
            smart_click_element("//*[@text='OPEN DATABASE']", "打开db")
            smart_click_element("//*[@text='your_data.mmb']", "选择db")
        elif app_name == "duckduckgo":
            smart_click_element("//*[@text=\"Let's do it!\"]", "开始使用")
            smart_click_element("//*[@text='Choose Your Browser']", "选择浏览器")
            handle_popups(2)
        elif app_name == "aurora.store":
            smart_click_element("//*[@text='Anonymous']", "匿名登录")
        elif app_name == "android.files":
            smart_click_element('//*[@resource-id="android:id/widget_frame"]', "允许访问文件权限")
            smart_click_element("//android.widget.ImageButton", "返回")
        elif app_name == "droidify":
            time.sleep(10)
        elif app_name == "book-story":
            handle_popups(4)
        elif app_name == "delta_chat":
            smart_click_element("//*[@text='CREATE NEW PROFILE']", "创建新用户")
            name_input = driver.find_element(AppiumBy.XPATH, '//*[@text="Your Name"]')
            name_input.send_keys("testuser")
            handle_popups(2)
        elif app_name == "spell4wiki":
            smart_click_element("//*[@resource-id=\"com.manimarank.spell4wiki:id/btnNext\"]", "继续")
            smart_click_element("//*[@resource-id=\"com.manimarank.spell4wiki:id/btnNext\"]", "继续")
            horizontal_swipe(4)
            handle_popups(3)
            smart_click_element("//*[@text='Skip Login']", "跳过登录")
        elif app_name == "expenses":
            horizontal_swipe(2)
            handle_popups(2)
        else:  # newpipe, Activity Manager, butterfly, News_Reader, souvenirs etc. - apps that don't need login
            logger.info(f"App {app_name} does not need login")
            # Still handle possible welcome pages and permission popups
            handle_popups()
            return True

        # After login, handle popups one final time
        time.sleep(2)
        handle_popups(1)

        logger.info(f"App {app_name} login process completed")
        return True

    except Exception as e:
        logger.error(f"Login failed: {e}")
        # Even if login failed, try handling popups
        try:
            handle_popups()
        except:
            pass
        return False

def send_to_wechat(content):
    """
    Send message to WeChat
    """
    key = get_wechat_sendkey()

    if not key:
        print("Warning: WECHAT_SENDKEY not configured, cannot send WeChat message")
        return False

    url = f"https://sctapi.ftqq.com/{key}.send"
    data = {
        "title": "Python crashed!",
        "desp": f"Error info follows:\n\n{content}"
    }

    try:
        response = requests.post(url, data=data, timeout=10)
        if response.status_code == 200:
            print(f"WeChat message sent successfully: {content[:50]}...")
            return True
        else:
            print(f"WeChat message send failed, status code: {response.status_code}")
            return False
    except Exception as e:
        print(f"Exception sending WeChat message: {e}")
        return False

class WeChatHandler(logging.Handler):
    """
    Custom log handler for sending error messages to WeChat
    """
    def __init__(self):
        super().__init__()
        self.setLevel(logging.ERROR)
        self.last_sent_time = {}
        self.min_interval = 300  # Minimum send interval (seconds), avoid frequent sending
        self.critical_keywords = [
            'Fatal',
            'failed to connect after',
            'maximum retries',
            'No available ports',
            'Failed to get page elements after maximum retries',
            'Driver启动失败',
            'Appium连接失败',
            '无法返回目标应用',
            'OutOfMemoryError',
            'SystemExit',
            'KeyboardInterrupt',
            'ConnectionRefusedError',
            'TimeoutError',
            'Critical'
        ]
        
    def emit(self, record):
        """
        Handle log record, send error message to WeChat
        """
        try:
            # Format error message
            log_entry = self.format(record)
            is_critical = any(keyword.lower() in log_entry.lower() for keyword in self.critical_keywords)

            if not is_critical and record.levelno < logging.CRITICAL:
                return  # Not critical error, skip

            # Prevent duplicate sending of same error (based on first 100 chars of message content)
            message_key = log_entry[:100]
            current_time = datetime.now().timestamp()

            if message_key in self.last_sent_time:
                if current_time - self.last_sent_time[message_key] < self.min_interval:
                    return  # Skip duplicate message

            self.last_sent_time[message_key] = current_time

            # Add context info
            context_info = f"""
Device info: {record.name if hasattr(record, 'name') else 'Unknown'}
Time: {datetime.fromtimestamp(current_time).strftime('%Y-%m-%d %H:%M:%S')}
Level: {record.levelname}

Error details:
{log_entry}
            """.strip()

            # Send message asynchronously to avoid blocking main program
            threading.Thread(
                target=send_to_wechat,
                args=(context_info,),
                daemon=True
            ).start()

        except Exception as e:
            print(f"WeChatHandler processing error: {e}")

class ThreadSafeLogger:
    def __init__(self, name="AndroidTestRunner"):
        self.name = name
        self.loggers = {}
        self.lock = threading.Lock()
        self.wechat_handler_added = False

        # Ensure log directory exists
        self._ensure_log_directory()

    def _ensure_log_directory(self):
        """Ensure log directory exists"""
        log_dir = 'logs'
        if not os.path.exists(log_dir):
            os.makedirs(log_dir)
            print(f"Created log directory: {log_dir}")

    def _add_wechat_handler_once(self):
        """
        Globally add WeChat handler once, avoid duplicate addition
        """
        if not self.wechat_handler_added:
            root_logger = logging.getLogger()

            # Check if WeChat handler already exists
            has_wechat_handler = any(
                isinstance(handler, WeChatHandler)
                for handler in root_logger.handlers
            )
            
            if not has_wechat_handler:
                wechat_handler = WeChatHandler()
                formatter = logging.Formatter(
                    '%(asctime)s - %(name)s - [%(levelname)s] - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S'
                )
                wechat_handler.setFormatter(formatter)
                root_logger.addHandler(wechat_handler)
                
            self.wechat_handler_added = True
        
    def get_logger(self, device_name=None, app_name=None):
        """
        Get thread-safe logger
        """
        # Ensure WeChat handler is added only once
        self._add_wechat_handler_once()

        thread_id = threading.get_ident()

        # Create independent logger for each device or thread
        if device_name and app_name:
            logger_name = f"{self.name}_{device_name}_{app_name}"
        elif device_name:
            logger_name = f"{self.name}_{device_name}"
        else:
            logger_name = f"{self.name}_{thread_id}"

        if logger_name not in self.loggers:
            with self.lock:
                if logger_name not in self.loggers:
                    logger = logging.getLogger(logger_name)
                    logger.setLevel(logging.INFO)

                    # Prevent duplicate handler addition
                    if not logger.handlers:
                        # Console handler
                        console_handler = logging.StreamHandler(sys.stdout)
                        console_handler.setLevel(logging.INFO)

                        # Ensure log directory exists
                        self._ensure_log_directory()

                        # File handler - simplified filename
                        log_filename = f'logs/{logger_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'

                        try:
                            file_handler = logging.FileHandler(log_filename, encoding='utf-8')
                            file_handler.setLevel(logging.DEBUG)
                        except Exception as e:
                            print(f"Failed to create file handler: {e}")
                            # If file handler creation failed, only use console handler
                            file_handler = None

                        # Formatter
                        formatter = logging.Formatter(
                            '%(asctime)s - %(name)s - [%(levelname)s] - %(message)s',
                            datefmt='%Y-%m-%d %H:%M:%S'
                        )

                        console_handler.setFormatter(formatter)
                        logger.addHandler(console_handler)

                        if file_handler:
                            file_handler.setFormatter(formatter)
                            logger.addHandler(file_handler)

                        # Set logger propagation to ensure error messages pass to root logger's WeChat handler
                        logger.propagate = True

                    self.loggers[logger_name] = logger

        return self.loggers[logger_name]

# Global log manager
log_manager = ThreadSafeLogger()

def get_logger(device_name=None, app_name=None):
    """Convenience function to get logger"""
    return log_manager.get_logger(device_name, app_name)

# --------------------
# ADB and device info tools
# --------------------

def run_adb_command(command):
    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    return result.stdout.strip()

def get_device_name():
    command = "adb devices"
    res = []
    result = run_adb_command(command).split('\n')
    for line in result:
        if "List" not in line and line.strip():
            device_info = line.split()
            if len(device_info) >= 2 and device_info[1] == "device":
                res.append(device_info[0])
    
    logger = get_logger()
    logger.info(f"Found {len(res)} available devices: {res}")
    return res

def get_android_version(device_name=None):
    if not device_name:
        command = "adb shell getprop ro.build.version.release"
    else:
        command = f"adb -s {device_name} shell getprop ro.build.version.release"
    result = run_adb_command(command)

    if device_name:
        device_logger = get_logger(device_name)
        device_logger.info(f"Device Android version: {result}")

    return result

# --------------------
# Run record and unified numbering tools
# --------------------

def _ensure_app_dir(app: str):
    app_dir = os.path.join('result', app)
    os.makedirs(app_dir, exist_ok=True)
    return app_dir

def get_next_run_index(app: str) -> int:
    """Get a globally incrementing run number by app dimension (unique across devices/threads)."""
    app_dir = _ensure_app_dir(app)
    counter_path = os.path.join(app_dir, '.run_counter')
    # Initialize counter file
    if not os.path.exists(counter_path):
        with open(counter_path, 'w', encoding='utf-8') as f:
            f.write('0\n')
    # Lock file, safely increment
    with open(counter_path, 'r+', encoding='utf-8') as f:
        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
        try:
            content = f.read().strip() or '0'
            try:
                cur = int(content)
            except ValueError:
                cur = 0
            nxt = cur + 1
            f.seek(0)
            f.write(str(nxt) + '\n')
            f.truncate()
        finally:
            fcntl.flock(f.fileno(), fcntl.LOCK_UN)
    return nxt

def count_existing_runs(app: str, algo: str) -> int:
    """Count number of runs produced for an app by algorithm (count by bug_report files)."""
    pattern = os.path.join('result', app, '*', f'{algo}_bug_report.json')
    return len(glob.glob(pattern))

def list_missing_run_indices(app: str, algo: str, target_rounds: int):
    """Return list of run indices in [1..target_rounds] missing {al}_bug_report.json."""
    missing = []
    for idx in range(1, target_rounds + 1):
        p = os.path.join('result', app, str(idx), f'{algo}_bug_report.json')
        if not os.path.exists(p):
            missing.append(idx)
    return missing

def build_missing_tasks(al_list, app_names, target_rounds: int):
    """Build missing task list: for each (app, algo), fill up to target_rounds runs, with specific run indices."""
    tasks = []
    for app in app_names:
        for al in al_list:
            miss_idxs = list_missing_run_indices(app, al, target_rounds)
            tasks.extend([(app, al, idx) for idx in miss_idxs])
    return tasks

def count_valid_activities_files(app: str, algo: str) -> int:
    """Count number of valid activities.json files produced for an app by algorithm (file exists and readable)."""
    pattern = os.path.join('result', app, '*', f'{algo}_activities.json')
    files = glob.glob(pattern)
    valid_count = 0
    logger = get_logger()

    for file_path in files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                # File is valid if it exists and is valid JSON dict format
                if isinstance(data, dict):
                    valid_count += 1
                    logger.debug(f"Found valid file: {file_path}, contains {len(data)} activities")
        except Exception as e:
            logger.warning(f"Error reading file {file_path}: {e}")
            continue

    return valid_count

# Convenience function: directly send urgent message to WeChat
def send_urgent_message(title, content):
    """
    Send urgent message to WeChat
    """
    key = get_wechat_sendkey()
    url = f"https://sctapi.ftqq.com/{key}.send"
    data = {
        "title": title,
        "desp": content
    }
    
    try:
        response = requests.post(url, data=data, timeout=10)
        if response.status_code == 200:
            print(f"Urgent message sent successfully: {title}")
            return True
        else:
            print(f"Urgent message send failed, status code: {response.status_code}")
            return False
    except Exception as e:
        print(f"Exception sending urgent message: {e}")
        return False

# Test function
def test_wechat_logging():
    """
    Test WeChat logging function
    """
    logger = get_logger("TestDevice", "TestApp")

    # Test different log levels
    logger.info("This is a test info message")
    logger.warning("This is a test warning")
    logger.error("This is a test error - should be sent to WeChat")

    # Test urgent message
    send_urgent_message("Test title", "This is an urgent test message")

def handle_app_permissions(driver, logger, max_attempts=10):
    """
    Function specifically for handling app permission requests
    """
    permission_keywords = [
        "允许", "Allow", "ALLOW", "允许一次", "Allow once", 
        "允许所有时间", "Allow all the time", "始终允许", "Always allow"
    ]
    
    for attempt in range(max_attempts):
        try:
            for keyword in permission_keywords:
                try:
                    element = driver.find_element(AppiumBy.XPATH, f"//*[@text='{keyword}']")
                    if element.is_displayed():
                        element.click()
                        logger.info(f"Granted permission: {keyword}")
                        time.sleep(1)
                        return True
                except NoSuchElementException:
                    continue

            # Check if there are still permission dialogs
            try:
                driver.find_element(AppiumBy.XPATH, "//*[contains(@text, '权限') or contains(@text, 'permission')]")
                time.sleep(1)  # Wait for permission dialog to fully load
                continue
            except NoSuchElementException:
                break  # No more permission dialogs

        except Exception as e:
            logger.debug(f"Permission handling attempt {attempt+1} failed: {e}")
            break

    return False

def skip_welcome_screens(driver, logger, max_attempts=5):
    """
    Skip welcome screens and guide pages
    """
    skip_keywords = [
        "跳过", "Skip", "SKIP", "下一步", "Next", "NEXT",
        "开始使用", "Get started", "开始", "Start", "继续", "Continue",
        "立即体验", "马上开始", "现在开始"
    ]
    
    for attempt in range(max_attempts):
        try:
            for keyword in skip_keywords:
                try:
                    element = driver.find_element(AppiumBy.XPATH, f"//*[@text='{keyword}']")
                    if element.is_displayed() and element.is_enabled():
                        element.click()
                        logger.info(f"Skipped welcome page: {keyword}")
                        time.sleep(2)
                        break
                except NoSuchElementException:
                    continue
            else:
                break  # No skip button found
        except Exception as e:
            logger.debug(f"Failed to skip welcome page: {e}")
            break

    return True