import subprocess
import os
import xml.etree.ElementTree as ET
import re
import time
import sys
import json

try:
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    print("❌ 缺少依赖库，请先运行: pip install pillow")
    sys.exit(1)

# ==============================================================================
# 1. 全局配置
# ==============================================================================
class Config:
    ADB_PATH = "adb" 
    
    # 输出文件路径 (Agent 的输入来源)
    LOCAL_XML_PATH = "./tree/ui_dump.xml"
    JSON_PATH = "./tree/ui_elements.json"          # 最终的结构化数据
    DEBUG_JSON_PATH = "./tree/ui_elements_debug.json" # 调试用的完整节点数据
    RAW_SCREEN_PATH = "./tree/screen_raw.png"      # 原始截图
    LABELED_SCREEN_PATH = "./tree/screen_labeled.png" # 标注后的截图 (SoM)

    # 逻辑开关
    ENABLE_CHILD_TEXT_MERGE = True # 合并子元素文本
    EXPORT_ALL_VISIBLE_NODES = True # 导出所有可见节点，而不只 clickable + 纯文本
    ONLY_VISIBLE_TO_USER = True    # 仅导出 visible-to-user=true 的节点
    SHRINK_STEP_PX = 5             # 视觉内缩像素 (防遮挡)

# ==============================================================================
# 2. 几何与数学工具
# ==============================================================================
def parse_bounds_rect(bounds_str):
    """解析 bounds 字符串为列表 [x1, y1, x2, y2]"""
    if not bounds_str: return None
    try:
        coords = list(map(int, re.findall(r'\d+', bounds_str)))
        if len(coords) == 4: return coords 
    except: pass
    return None

def get_center(rect):
    """计算矩形几何中心"""
    return (rect[0] + rect[2]) // 2, (rect[1] + rect[3]) // 2

def get_area(rect):
    return (rect[2] - rect[0]) * (rect[3] - rect[1])

def resource_id_to_label(resource_id):
    """把 resource-id 转成更易读的调试标签"""
    if not resource_id:
        return ""
    tail = resource_id.split("/")[-1]
    tail = tail.replace("_", " ").replace("-", " ").strip()
    return tail

def is_inside(inner_rect, outer_rect):
    """判断 inner 是否完全在 outer 内部"""
    return (inner_rect[0] >= outer_rect[0] and 
            inner_rect[1] >= outer_rect[1] and 
            inner_rect[2] <= outer_rect[2] and 
            inner_rect[3] <= outer_rect[3])

def is_point_in_rect(point, rect):
    """判断点是否在矩形内"""
    x, y = point
    return rect[0] <= x <= rect[2] and rect[1] <= y <= rect[3]

def clip_rect(rect, screen_w, screen_h):
    """屏幕边缘裁剪 (处理只露出一半的控件)"""
    x1, y1, x2, y2 = rect
    new_x1 = max(0, x1)
    new_y1 = max(0, y1)
    new_x2 = min(screen_w, x2)
    new_y2 = min(screen_h, y2)
    if new_x2 <= new_x1 or new_y2 <= new_y1:
        return None
    return [new_x1, new_y1, new_x2, new_y2]

def find_safe_point(parent_rect, children_rects):
    """
    智能避让算法：在父控件中寻找一个不与子控件重叠的坐标
    """
    p_x1, p_y1, p_x2, p_y2 = parent_rect
    w, h = p_x2 - p_x1, p_y2 - p_y1
    
    # 搜索网格: 中心 -> 十字线 -> 四角
    search_grid = [
        (0.5, 0.5), (0.5, 0.25), (0.5, 0.75), 
        (0.25, 0.5), (0.75, 0.5),
        (0.2, 0.2), (0.8, 0.2), (0.2, 0.8), (0.8, 0.8)
    ]
    
    for rx, ry in search_grid:
        cx = int(p_x1 + w * rx)
        cy = int(p_y1 + h * ry)
        
        # 检查该点是否撞到了子控件
        hit = False
        for child in children_rects:
            if is_point_in_rect((cx, cy), child):
                hit = True
                break
        if not hit: return (cx, cy) # 找到安全点
            
    return get_center(parent_rect) # 兜底

# ==============================================================================
# 3. ADB 基础操作
# ==============================================================================
def run_cmd(cmd):
    try:
        if Config.ADB_PATH != "adb" and cmd.startswith("adb"):
            cmd = cmd.replace("adb", f'"{Config.ADB_PATH}"', 1)
        res = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8")
        return res.stdout.strip()
    except Exception: return None

def get_devices():
    out = run_cmd("adb devices")
    if not out: return []
    devs = []
    for line in out.split('\n'):
        if "\t" in line and "List of" not in line:
            parts = line.split("\t")
            devs.append({"serial": parts[0].strip(), "status": parts[1].strip()})
    return devs

def get_screen_size(serial):
    prefix = f"adb -s {serial}" if serial else "adb"
    out = run_cmd(f"{prefix} shell wm size")
    if out:
        m = re.search(r'(\d+)x(\d+)', out)
        if m: return int(m.group(1)), int(m.group(2))
    return 1080, 2400

def select_device():
    devs = [d for d in get_devices() if d['status'] == 'device']
    if not devs: return None
    if len(devs) == 1: return devs[0]['serial']
    # 只有多台设备时才交互，否则自动选
    print("多台设备在线，请选择:")
    for i, d in enumerate(devs): print(f"{i+1}. {d['serial']}")
    try:
        idx = int(input("序号: ")) - 1
        return devs[idx]['serial'] if 0 <= idx < len(devs) else None
    except: return None

# ==============================================================================
# 4. 数据解析核心 (Parsing Layer)
# ==============================================================================
def parse_ui_tree(device_serial):
    """解析 XML，执行合并、裁剪、避让，返回清洗后的列表"""
    prefix = f"adb -s {device_serial}"
    screen_w, screen_h = get_screen_size(device_serial)
    
    # 0. 清除旧缓存
    run_cmd(f"{prefix} shell rm -f /sdcard/window_dump.xml")
    
    # 1. Dump XML
    run_cmd(f"{prefix} shell uiautomator dump /sdcard/window_dump.xml")
    run_cmd(f"{prefix} pull /sdcard/window_dump.xml {Config.LOCAL_XML_PATH}")
    
    if not os.path.exists(Config.LOCAL_XML_PATH): return []

    tree = ET.parse(Config.LOCAL_XML_PATH)
    root = tree.getroot()
    
    all_clickables = []
    all_non_clickables = []
    non_clickable_texts = []
    
    # 2. 初步提取与屏幕裁剪
    for node in root.iter():
        text = node.attrib.get("text", "")
        desc = node.attrib.get("content-desc", "")
        content = text if text else desc
        resource_id = node.attrib.get("resource-id", "")
        visible_to_user = node.attrib.get("visible-to-user", "true") != "false"
        enabled = node.attrib.get("enabled", "true") == "true"
        if Config.ONLY_VISIBLE_TO_USER and not visible_to_user:
            continue
        
        raw_rect = parse_bounds_rect(node.attrib.get("bounds", ""))
        if not raw_rect: continue

        # 裁剪：确保只处理屏幕内的元素
        visible_rect = clip_rect(raw_rect, screen_w, screen_h)
        if not visible_rect: continue
        
        item = {
            "id": id(node),
            "type": node.attrib.get("class", ""),
            "content": content,
            "text": text,
            "desc": desc,
            "resource_id": resource_id,
            "resource_label": resource_id_to_label(resource_id),
            "package": node.attrib.get("package", ""),
            "enabled": enabled,
            "visible_to_user": visible_to_user,
            "rect": visible_rect,
            "area": get_area(visible_rect),
            "center": get_center(visible_rect),
            "temp_children_text": []
        }
        
        if node.attrib.get("clickable") == "true":
            all_clickables.append(item)
        else:
            all_non_clickables.append(item)
            if content:
                non_clickable_texts.append(item)

    # 3. 文本合并 (子元素文本归属给父按钮)
    if Config.ENABLE_CHILD_TEXT_MERGE:
        for txt in non_clickable_texts:
            # 找到包含该文本的所有空白父按钮
            parents = [btn for btn in all_clickables if not btn["content"] and is_inside(txt["rect"], btn["rect"])]
            if parents:
                # 归属给面积最小(最内层)的父按钮
                best = sorted(parents, key=lambda x: x["area"])[0]
                best["temp_children_text"].append(txt["content"])
        
        # 应用合并
        for btn in all_clickables:
            if not btn["content"] and btn["temp_children_text"]:
                btn["content"] = " ".join(btn["temp_children_text"])

    # 4. 智能避让 (点击穿透处理)
    for parent in all_clickables:
        # 找到当前 parent 内部的所有子点击区域
        children = [c['rect'] for c in all_clickables if c['id'] != parent['id'] and is_inside(c['rect'], parent['rect'])]
        if children:
            # 检查当前中心点是否安全，不安全则寻找新点
            if any(is_point_in_rect(parent['center'], c) for c in children):
                parent['center'] = find_safe_point(parent['rect'], children)

    # 5. 选择需要导出的非 clickable 节点
    # 全量模式：导出所有可见非 clickable 节点
    # 兼容模式：仅导出不在 clickable 内部的纯文本节点
    if Config.EXPORT_ALL_VISIBLE_NODES:
        export_non_clickables = all_non_clickables
    else:
        export_non_clickables = []
        for txt in non_clickable_texts:
            inside_clickable = any(is_inside(txt["rect"], btn["rect"]) for btn in all_clickables)
            if not inside_clickable:
                export_non_clickables.append(txt)

    # 6. 生成最终数据结构 (格式化为 Agent 需要的 JSON)
    final_data = []
    export_items = []
    for item in all_clickables:
        export_items.append({
            **item,
            "clickable": True,
        })
    for item in export_non_clickables:
        export_items.append({
            **item,
            "clickable": False,
        })

    for i, item in enumerate(export_items):
        final_data.append({
            "id": i,                        # 映射给 Agent 的短 ID
            "context": item["content"],     # 文本内容
            "center": item["center"],       # 安全点击坐标 (x, y)
            "type": item["type"],           # 控件类型 (辅助信息)
            "clickable": item["clickable"], # 是否可点击
            "bbox": item["rect"],           # 边界框 (辅助信息)
            "text": item["text"],           # 原始 text
            "desc": item["desc"],           # 原始 content-desc
            "resource_id": item["resource_id"],
            "resource_label": item["resource_label"],
            "package": item["package"],
            "enabled": item["enabled"],
            "visible_to_user": item["visible_to_user"],
            
            # 以下字段仅用于内部画图计算层级，保存后可忽略
            "_area": item["area"],
            "_raw_id": item["id"]
        })
        
    return final_data

# ==============================================================================
# 5. 视觉处理核心 (Visual Layer)
# ==============================================================================
def draw_labels(device_serial, ui_elements):
    """截图并绘制 SoM 标记 (仅绘制有文本的节点，带层级内缩)"""
    prefix = f"adb -s {device_serial}" if device_serial else "adb"
    run_cmd(f"{prefix} shell screencap -p /sdcard/screen.png")
    run_cmd(f"{prefix} pull /sdcard/screen.png {Config.RAW_SCREEN_PATH}")
    
    if not os.path.exists(Config.RAW_SCREEN_PATH): return None

    image = Image.open(Config.RAW_SCREEN_PATH).convert("RGB")
    draw = ImageDraw.Draw(image)
    try: font = ImageFont.truetype("arial.ttf", 24)
    except: font = ImageFont.load_default()

    drawable_items = [item for item in ui_elements if str(item.get("context", "")).strip()]
    if not drawable_items:
        image.save(Config.LABELED_SCREEN_PATH)
        return Config.LABELED_SCREEN_PATH

    # 计算每个元素的层级 (用于视觉内缩)
    # 复杂度 O(N^2)，但在 UI 元素数量级下很快
    for item in drawable_items:
        level = 0
        for other in drawable_items:
            if item['id'] == other['id']: continue
            # 如果 item 在 other 内部，层级+1
            if is_inside(item['bbox'], other['bbox']):
                # 处理相同大小的重叠情况
                if other['_area'] > item['_area']: level += 1
                elif other['_area'] == item['_area'] and item['id'] > other['id']: level += 1
        item['_shrink_level'] = level

    # 绘制
    for item in drawable_items:
        rect = item['bbox']
        idx = item['id']
        level = item.get('_shrink_level', 0)
        
        # 计算内缩
        gap = min(level * Config.SHRINK_STEP_PX, min(rect[2]-rect[0], rect[3]-rect[1])//2 - 2)
        gap = max(0, gap)
        
        vis_rect = [rect[0]+gap, rect[1]+gap, rect[2]-gap, rect[3]-gap]
        
        # 画框
        draw.rectangle(vis_rect, outline="red", width=3)
        
        # 画标签
        tag = str(idx)
        bbox = font.getbbox(tag)
        w, h = bbox[2]-bbox[0], bbox[3]-bbox[1]
        
        tx, ty = vis_rect[0], vis_rect[1] - h - 4
        if ty < 0: ty = vis_rect[1] # 顶端边界处理
        
        draw.rectangle([tx, ty, tx+w+8, ty+h+4], fill="red", outline="red")
        draw.text((tx+4, ty), tag, fill="white", font=font)

    image.save(Config.LABELED_SCREEN_PATH)
    return Config.LABELED_SCREEN_PATH

# ==============================================================================
# 主程序入口
# ==============================================================================
if __name__ == "__main__":
    device = select_device()
    if not device:
        print("❌ 未连接设备")
        sys.exit(1)

    os.makedirs("./tree", exist_ok=True)

    print(f"🚀 开始处理设备: {device}")
    start_time = time.time()

    # 1. 获取并清洗数据
    ui_data = parse_ui_tree(device)
    
    if ui_data:
        # 2. 生成视觉反馈图
        draw_labels(device, ui_data)
        
        # 3. 清理不需要的内部字段，只保留对外导出字段
        final_json_data = []
        debug_json_data = []
        for item in ui_data:
            final_json_data.append({
                "id": item["id"],
                "context": item["context"], # 文本描述
                "center": item["center"],   # 点击坐标
                "type": item["type"],       # 控件类型
                "clickable": item["clickable"], # 是否可点击
                # "bbox": item["bbox"]      # 可选：如果 Agent 需要知道大小
            })
            debug_json_data.append({
                "id": item["id"],
                "context": item["context"],
                "text": item["text"],
                "desc": item["desc"],
                "resource_id": item["resource_id"],
                "resource_label": item["resource_label"],
                "center": item["center"],
                "bbox": item["bbox"],
                "type": item["type"],
                "package": item["package"],
                "clickable": item["clickable"],
                "enabled": item["enabled"],
                "visible_to_user": item["visible_to_user"],
            })

        # 4. 保存 JSON
        with open(Config.JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(final_json_data, f, ensure_ascii=False, indent=2)
        with open(Config.DEBUG_JSON_PATH, "w", encoding="utf-8") as f:
            json.dump(debug_json_data, f, ensure_ascii=False, indent=2)

        print(f"✅ 处理完成 (耗时 {time.time()-start_time:.2f}s)")
        print(f"📂 数据已保存: {Config.JSON_PATH}")
        print(f"🧪 调试数据已保存: {Config.DEBUG_JSON_PATH}")
        print(f"🖼️  图片已保存: {Config.LABELED_SCREEN_PATH}")
    else:
        print("⚠️ 未发现可见 UI 元素或解析失败")