#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
统一元素管理器 - 合并元素列表功能

功能：
1. 统一元素列表接口
2. 自动平台检测和适配
3. 智能过滤和排序
"""

import xml.etree.ElementTree as ET
import re
from typing import List, Dict, Optional


class ElementManager:
    """统一元素管理器"""
    
    def __init__(self, mobile_client):
        self.client = mobile_client
    
    def _is_ios(self) -> bool:
        """判断当前是否为 iOS 平台"""
        return getattr(self.client, 'platform', 'android') == 'ios'
    
    def _get_ios_client(self):
        """获取 iOS 客户端"""
        if hasattr(self.client, '_ios_client') and self.client._ios_client:
            return self.client._ios_client
        if hasattr(self.client, 'wda') and self.client.wda:
            return self.client.wda
        return None
    
    def list_elements(self, max_elements: int = 100, filter_interactive: bool = True, 
                       include_empty_clickable: bool = True,
                       include_system_ui: bool = False) -> List[Dict]:
        """统一元素列表接口
        
        Args:
            max_elements: 最大返回元素数量
            filter_interactive: 是否只返回可交互元素
            include_empty_clickable: 是否包含无text但clickable的元素（如图片按钮）
            include_system_ui: 是否包含系统UI元素（状态栏、通知栏等），默认False
        
        Returns:
            元素列表，每个元素包含:
            - text: 文本内容
            - bounds: 边界坐标 [x1,y1][x2,y2]
            - clickable: 是否可点击
            - hint: 位置提示（仅对无text的clickable元素）
        """
        try:
            if self._is_ios():
                return self._list_elements_ios(max_elements, filter_interactive)
            else:
                return self._list_elements_android(max_elements, filter_interactive, 
                                                    include_empty_clickable, include_system_ui)
        except Exception as e:
            return [{"error": f"❌ 获取元素列表失败: {e}"}]
    
    def _list_elements_android(self, max_elements: int, filter_interactive: bool, 
                                 include_empty_clickable: bool = True,
                                 include_system_ui: bool = False) -> List[Dict]:
        """Android元素列表实现
        
        增强功能：返回所有clickable元素，包括无text的图片按钮
        默认过滤系统UI元素，只返回App内元素
        """
        try:
            xml_string = self.client.u2.dump_hierarchy(compressed=False)
            if not xml_string:
                return []
            
            root = ET.fromstring(xml_string)
            elements = []
            all_elements_for_hint = []  # 用于计算hint
            
            def parse_node(node, depth=0):
                if depth > 30:  # 增加深度限制，支持更深层嵌套的UI
                    return
                
                text = node.get('text', '').strip()
                resource_id = node.get('resource-id', '')
                class_name = node.get('class', '')
                content_desc = node.get('content-desc', '').strip()
                bounds_str = node.get('bounds', '')
                clickable = node.get('clickable', 'false') == 'true'
                focusable = node.get('focusable', 'false') == 'true'
                enabled = node.get('enabled', 'true') == 'true'
                
                x1, y1, x2, y2 = 0, 0, 0, 0
                if bounds_str:
                    match = re.match(r'\[(\d+),(\d+)\]\[(\d+),(\d+)\]', bounds_str)
                    if match:
                        x1, y1, x2, y2 = map(int, match.groups())
                
                width, height = x2 - x1, y2 - y1
                
                # 过滤太小的元素
                if width < 10 or height < 10:
                    for child in node:
                        parse_node(child, depth + 1)
                    return
                
                # 彻底排除系统UI元素（状态栏、通知栏等）- 整个子树都跳过
                if 'com.android.systemui' in resource_id:
                    return  # 不遍历子节点
                
                # 排除华为浮动任务等
                if 'com.huawei' in resource_id:
                    return
                
                # 过滤全屏容器元素（通常是布局容器，不可交互）
                # 但仍然需要遍历它们的子节点
                if width > 700 and height > 1400 and not clickable and not text:
                    # 不加入elements列表，但继续遍历子节点
                    for child in node:
                        parse_node(child, depth + 1)
                    return
                
                class_name_simple = class_name.split('.')[-1] if class_name else ''
                
                element = {
                    'text': text,
                    'resource-id': resource_id,
                    'class': class_name_simple,
                    'content-desc': content_desc,
                    'bounds': bounds_str,
                    'x1': x1, 'y1': y1, 'x2': x2, 'y2': y2,
                    'x': x1, 'y': y1,
                    'width': width,
                    'height': height,
                    'clickable': clickable,
                    'focusable': focusable,
                    'enabled': enabled,
                    'depth': depth
                }
                
                # 收集所有有text的元素用于hint计算
                if text or content_desc:
                    all_elements_for_hint.append(element)
                
                # 标记是否为系统UI元素
                is_system_ui = (
                    'com.android.systemui' in resource_id or 
                    'com.huawei' in resource_id or
                    (content_desc and any(kw in content_desc for kw in ['通知', '信号', '充电', '护眼', '静音', 'WLAN', 'EDGE']))
                )
                element['_is_system_ui'] = is_system_ui
                
                if filter_interactive:
                    # 只保留：clickable元素 或 有text的元素
                    if clickable or text:
                        elements.append(element)
                else:
                    if clickable or text:
                        elements.append(element)
                
                for child in node:
                    parse_node(child, depth + 1)
            
            parse_node(root)
            
            # 为无text的clickable元素添加hint（基于相邻元素）
            for elem in elements:
                if elem.get('clickable') and not elem.get('text') and not elem.get('content-desc'):
                    hint = self._generate_hint(elem, all_elements_for_hint)
                    if hint:
                        elem['hint'] = hint
            
            # 最终过滤：排除系统UI和无意义元素
            filtered = []
            for e in elements:
                # 过滤系统UI（除非明确要求包含）
                if not include_system_ui and e.get('_is_system_ui', False):
                    continue
                
                # 移除内部标记
                e.pop('_is_system_ui', None)
                filtered.append(e)
            
            # 排序：clickable优先，然后按Y坐标（从上到下），再按X坐标（从左到右）
            filtered.sort(key=lambda e: (
                not e.get('clickable', False),  # clickable优先
                e.get('y1', 0),                 # Y坐标从上到下
                e.get('x1', 0),                 # X坐标从左到右
            ))
            
            return filtered[:max_elements]
            
        except Exception as e:
            return [{"error": f"❌ Android元素解析失败: {e}"}]
    
    def _generate_hint(self, target: Dict, all_elements: List[Dict]) -> str:
        """为无text元素生成位置提示hint
        
        策略：找到上下左右相邻的有text元素（包括clickable元素的text），生成方位提示
        优先级：左侧 > 上方 > 右侧 > 下方
        """
        target_x1, target_y1 = target.get('x1', 0), target.get('y1', 0)
        target_x2, target_y2 = target.get('x2', 0), target.get('y2', 0)
        target_x_center = (target_x1 + target_x2) // 2
        target_y_center = (target_y1 + target_y2) // 2
        
        left_candidates = []
        right_candidates = []
        top_candidates = []
        bottom_candidates = []
        
        for elem in all_elements:
            if elem is target:
                continue
            
            # 获取元素文本（优先text，其次content-desc）
            elem_text = elem.get('text', '') or elem.get('content-desc', '')
            if not elem_text:
                continue
            
            elem_x1, elem_y1 = elem.get('x1', 0), elem.get('y1', 0)
            elem_x2, elem_y2 = elem.get('x2', 0), elem.get('y2', 0)
            elem_x_center = (elem_x1 + elem_x2) // 2
            elem_y_center = (elem_y1 + elem_y2) // 2
            
            # 同一行判断：Y坐标有重叠，或Y中心差距小于元素高度的一半
            y_overlap = not (elem_y2 < target_y1 or elem_y1 > target_y2)
            height_threshold = max(50, (target_y2 - target_y1) // 2)
            same_row = y_overlap or abs(elem_y_center - target_y_center) < height_threshold
            
            # 同一列判断
            x_overlap = not (elem_x2 < target_x1 or elem_x1 > target_x2)
            width_threshold = max(50, (target_x2 - target_x1) // 2)
            same_col = x_overlap or abs(elem_x_center - target_x_center) < width_threshold
            
            if same_row:
                # 左侧元素（元素右边界在target左边界附近或左边）
                if elem_x_center < target_x_center and elem_x2 <= target_x1 + 50:
                    dist = abs(target_x1 - elem_x2)
                    left_candidates.append((dist, elem_text, "右侧"))
                # 右侧元素
                elif elem_x_center > target_x_center and elem_x1 >= target_x2 - 50:
                    dist = abs(elem_x1 - target_x2)
                    right_candidates.append((dist, elem_text, "左侧"))
            
            if same_col:
                # 上方元素
                if elem_y_center < target_y_center and elem_y2 <= target_y1 + 50:
                    dist = abs(target_y1 - elem_y2)
                    top_candidates.append((dist, elem_text, "下方"))
                # 下方元素
                elif elem_y_center > target_y_center and elem_y1 >= target_y2 - 50:
                    dist = abs(elem_y1 - target_y2)
                    bottom_candidates.append((dist, elem_text, "上方"))
        
        # 优先级：左侧 > 上方 > 右侧 > 下方（取最近的）
        for candidates in [left_candidates, top_candidates, right_candidates, bottom_candidates]:
            if candidates:
                candidates.sort(key=lambda x: x[0])
                return f"{candidates[0][1]}{candidates[0][2]}"
        
        return ""
    
    def find_nearby(self, text: str, direction: str = "right") -> Dict:
        """查找指定文字附近的clickable元素
        
        Args:
            text: 参照元素的文字内容
            direction: 方向 - "left", "right", "above", "below"
        
        Returns:
            找到的元素信息，包含bounds
        """
        try:
            # 获取所有元素（包括非clickable的文字元素）
            elements = self.list_elements(max_elements=300, filter_interactive=False)
            
            # 找到参照元素
            ref_element = None
            for e in elements:
                if e.get('text') == text or text in e.get('text', ''):
                    ref_element = e
                    break
            
            if not ref_element:
                return {"success": False, "message": f"未找到文字: {text}"}
            
            ref_x1, ref_y1 = ref_element.get('x1', 0), ref_element.get('y1', 0)
            ref_x2, ref_y2 = ref_element.get('x2', 0), ref_element.get('y2', 0)
            ref_x_center = (ref_x1 + ref_x2) // 2
            ref_y_center = (ref_y1 + ref_y2) // 2
            
            # 根据方向筛选候选元素
            candidates = []
            
            for e in elements:
                if e is ref_element:
                    continue
                if not e.get('clickable'):
                    continue
                
                e_x1, e_y1 = e.get('x1', 0), e.get('y1', 0)
                e_x2, e_y2 = e.get('x2', 0), e.get('y2', 0)
                e_x_center = (e_x1 + e_x2) // 2
                e_y_center = (e_y1 + e_y2) // 2
                
                # 检查是否在同一行/列
                y_overlap = not (e_y2 < ref_y1 or e_y1 > ref_y2)
                x_overlap = not (e_x2 < ref_x1 or e_x1 > ref_x2)
                
                if direction == "right":
                    # 同一行，且在右侧
                    if y_overlap and e_x1 >= ref_x2 - 20:
                        dist = e_x1 - ref_x2
                        candidates.append((dist, e))
                elif direction == "left":
                    # 同一行，且在左侧
                    if y_overlap and e_x2 <= ref_x1 + 20:
                        dist = ref_x1 - e_x2
                        candidates.append((dist, e))
                elif direction == "below":
                    # 同一列，且在下方
                    if x_overlap and e_y1 >= ref_y2 - 20:
                        dist = e_y1 - ref_y2
                        candidates.append((dist, e))
                elif direction == "above":
                    # 同一列，且在上方
                    if x_overlap and e_y2 <= ref_y1 + 20:
                        dist = ref_y1 - e_y2
                        candidates.append((dist, e))
            
            if not candidates:
                return {"success": False, "message": f"在 {text} 的 {direction} 方向未找到clickable元素"}
            
            # 取最近的
            candidates.sort(key=lambda x: x[0])
            found = candidates[0][1]
            
            return {
                "success": True,
                "bounds": found.get('bounds'),
                "x1": found.get('x1'),
                "y1": found.get('y1'),
                "x2": found.get('x2'),
                "y2": found.get('y2'),
                "text": found.get('text', ''),
                "hint": f"{text}的{direction}侧",
                "resource-id": found.get('resource-id', ''),
                "class": found.get('class', '')
            }
            
        except Exception as e:
            return {"success": False, "message": f"查找失败: {e}"}
    
    def _list_elements_ios(self, max_elements: int, filter_interactive: bool) -> List[Dict]:
        """iOS元素列表实现"""
        try:
            ios_client = self._get_ios_client()
            if not ios_client:
                return [{"error": "❌ iOS客户端未初始化"}]
            
            # 获取iOS页面源码
            elements = []
            
            try:
                # 方法1：使用WDA的source
                source = ios_client.wda.source()
                if source:
                    root = ET.fromstring(source)
                    elements = self._parse_ios_elements(root, max_elements, filter_interactive)
            except Exception:
                pass
            
            # 如果方法1失败，尝试方法2：使用find_elements
            if not elements:
                try:
                    # 获取所有可交互元素
                    interactive_elements = []
                    
                    # 查找按钮
                    buttons = ios_client.wda.find_elements_by_class_name("XCUIElementTypeButton")
                    for btn in buttons:
                        if len(interactive_elements) < max_elements:
                            elements.append(self._convert_ios_element(btn, "Button"))
                    
                    # 查找文本框
                    text_fields = ios_client.wda.find_elements_by_class_name("XCUIElementTypeTextField")
                    for tf in text_fields:
                        if len(elements) < max_elements:
                            elements.append(self._convert_ios_element(tf, "TextField"))
                    
                    # 查找其他可交互元素
                    other_elements = ios_client.wda.find_elements_by_class_name("XCUIElementTypeCell")
                    for elem in other_elements:
                        if len(elements) < max_elements:
                            elements.append(self._convert_ios_element(elem, "Cell"))
                except Exception:
                    pass
            
            return elements[:max_elements]
            
        except Exception as e:
            return [{"error": f"❌ iOS元素解析失败: {e}"}]
    
    def _parse_ios_elements(self, root, max_elements: int, filter_interactive: bool) -> List[Dict]:
        """解析iOS XML元素"""
        elements = []
        
        def parse_ios_node(node, depth=0):
            if depth > 20 or len(elements) >= max_elements:
                return
            
            # 提取iOS元素属性
            elem_type = node.get('type', '')
            name = node.get('name', '')
            label = node.get('label', '')
            value = node.get('value', '')
            visible = node.get('visible', 'false') == 'true'
            enabled = node.get('enabled', 'true') == 'true'
            
            # iOS特有的bounds格式
            bounds_str = node.get('bounds', '')
            x, y, width, height = 0, 0, 0, 0
            if bounds_str:
                match = re.match(r'\[(\d+),(\d+)\]\[(\d+),(\d+)\]', bounds_str)
                if match:
                    x1, y1, x2, y2 = map(int, match.groups())
                    x, y = x1, y1
                    width, height = x2 - x1, y2 - y1
            
            # 判断是否可交互
            interactive_types = [
                'XCUIElementTypeButton', 'XCUIElementTypeTextField', 'XCUIElementTypeSecureTextField',
                'XCUIElementTypeSwitch', 'XCUIElementTypeSlider', 'XCUIElementTypePicker',
                'XCUIElementTypeCell', 'XCUIElementTypeCollectionViewCell'
            ]
            
            is_interactive = elem_type in interactive_types
            
            # 过滤条件
            if filter_interactive and not is_interactive:
                # 递归处理子节点
                for child in node:
                    parse_ios_node(child, depth + 1)
                return
            
            # 构建元素信息
            element = {
                'text': name or label or value or '',
                'type': elem_type,
                'name': name,
                'label': label,
                'value': value,
                'bounds': bounds_str,
                'x': x,
                'y': y,
                'width': width,
                'height': height,
                'clickable': is_interactive,
                'enabled': enabled,
                'visible': visible,
                'depth': depth
            }
            
            elements.append(element)
            
            # 递归处理子节点
            for child in node:
                parse_ios_node(child, depth + 1)
        
        parse_ios_node(root)
        
        # 智能排序
        elements.sort(key=lambda e: (
            not e.get('clickable', False),
            not bool(e.get('text', '')),
            e.get('depth', 0)
        ))
        
        return elements
    
    def _convert_ios_element(self, element, elem_type: str) -> Dict:
        """转换iOS WDA元素为统一格式"""
        try:
            # 获取元素属性
            bounds = element.bounds
            rect = bounds.rect
            
            return {
                'text': element.name or element.label or element.value or '',
                'type': elem_type,
                'name': element.name,
                'label': element.label,
                'value': element.value,
                'bounds': f"[{int(rect.x)},{int(rect.y)}][{int(rect.x + rect.width)},{int(rect.y + rect.height)}]",
                'x': int(rect.x),
                'y': int(rect.y),
                'width': int(rect.width),
                'height': int(rect.height),
                'clickable': True,
                'enabled': element.is_enabled,
                'visible': element.is_displayed,
                'depth': 0
            }
        except Exception as e:
            return {
                'error': f"转换iOS元素失败: {e}",
                'type': elem_type
            }
    
    def find_element_by_text(self, text: str, exact: bool = False) -> Optional[Dict]:
        """根据文本查找元素"""
        elements = self.list_elements(filter_interactive=False)
        
        for elem in elements:
            if 'error' in elem:
                continue
            
            elem_text = elem.get('text', '').lower()
            search_text = text.lower()
            
            if exact:
                if elem_text == search_text:
                    return elem
            else:
                if search_text in elem_text:
                    return elem
        
        return None
    
    def find_element_by_id(self, resource_id: str) -> Optional[Dict]:
        """根据resource-id查找元素"""
        elements = self.list_elements(filter_interactive=False)
        
        for elem in elements:
            if 'error' in elem:
                continue
            
            elem_id = elem.get('resource-id', '')
            if elem_id == resource_id or elem_id.endswith(f":id/{resource_id}"):
                return elem
        
        return None
    
    def get_clickable_elements(self) -> List[Dict]:
        """获取所有可点击元素"""
        elements = self.list_elements(filter_interactive=True)
        return [elem for elem in elements if 'error' not in elem and elem.get('clickable', False)]
    
    def get_elements_by_type(self, elem_type: str) -> List[Dict]:
        """根据类型获取元素"""
        elements = self.list_elements(filter_interactive=False)
        
        return [elem for elem in elements 
                if 'error' not in elem and 
                (elem.get('type', '') == elem_type or elem.get('class', '') == elem_type)]
