#!/usr/bin/env python3
"""
ADB Ref System - 类似 X-OmniClaw 的 UI 元素引用系统
用法:
    python adb_ref.py snapshot          # 获取 UI 树，显示所有可交互元素
    python adb_ref.py tap e3            # 点击 ref=e3 的元素
    python adb_ref.py find "女兒殿下"   # 搜索包含文字的元素
    python adb_ref.py open "spotify:search:周杰伦"  # 打开 Deep Link
"""

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

# 配置
ADB_DEVICE = os.environ.get("ADB_DEVICE", "localhost:15556")
UI_XML_PATH = "/sdcard/ui.xml"
REF_CACHE_PATH = "/tmp/adb_refs.json"

def run_adb(cmd, timeout=30):
    """执行 ADB 命令"""
    full_cmd = f"adb -s {ADB_DEVICE} {cmd}"
    try:
        result = subprocess.run(full_cmd, shell=True, capture_output=True, text=True, timeout=timeout)
        return result.stdout + result.stderr, result.returncode
    except subprocess.TimeoutExpired:
        return "Timeout", 1

def parse_bounds(bounds_str):
    """解析 bounds="[x1,y1][x2,y2]" 返回 (x1, y1, x2, y2)"""
    match = re.findall(r'\[(\d+),(\d+)\]\[(\d+),(\d+)\]', bounds_str)
    if match:
        return tuple(map(int, match[0]))
    return None

def get_center(bounds):
    """计算中心点"""
    x1, y1, x2, y2 = bounds
    return (x1 + x2) // 2, (y1 + y2) // 2

def is_interactive(node):
    """判断元素是否可交互"""
    clickable = node.get('clickable') == 'true'
    focusable = node.get('focusable') == 'true'
    checkable = node.get('checkable') == 'true'
    has_text = bool(node.get('text', '').strip())
    has_desc = bool(node.get('content-desc', '').strip())
    
    # 可点击，或者有文字/描述的可聚焦元素
    return clickable or (focusable and (has_text or has_desc)) or checkable

def snapshot():
    """获取 UI 树并建立 ref 映射"""
    print("📱 Dumping UI tree...")
    out, code = run_adb(f"shell uiautomator dump {UI_XML_PATH}")
    if "dumped" not in out.lower() and code != 0:
        print(f"❌ Failed to dump: {out}")
        return None
    
    # 拉取 XML
    out, code = run_adb(f"pull {UI_XML_PATH} /tmp/ui.xml")
    if code != 0:
        print(f"❌ Failed to pull: {out}")
        return None
    
    # 解析 XML
    try:
        tree = ET.parse("/tmp/ui.xml")
        root = tree.getroot()
    except Exception as e:
        print(f"❌ Failed to parse XML: {e}")
        return None
    
    # 遍历所有节点，建立 ref 映射
    refs = {}
    ref_id = 1
    
    def traverse(node, depth=0):
        nonlocal ref_id
        
        bounds_str = node.get('bounds', '')
        bounds = parse_bounds(bounds_str)
        if not bounds:
            for child in node:
                traverse(child, depth + 1)
            return
        
        text = node.get('text', '').strip()
        desc = node.get('content-desc', '').strip()
        cls = node.get('class', '').split('.')[-1]
        pkg = node.get('package', '')
        clickable = node.get('clickable') == 'true'
        
        # 只记录可交互或有文字的元素
        if is_interactive(node) or text or desc:
            ref = f"e{ref_id}"
            center = get_center(bounds)
            refs[ref] = {
                'text': text,
                'desc': desc,
                'class': cls,
                'package': pkg,
                'bounds': bounds,
                'center': center,
                'clickable': clickable,
                'depth': depth
            }
            ref_id += 1
        
        for child in node:
            traverse(child, depth + 1)
    
    traverse(root)
    
    # 保存到缓存
    with open(REF_CACHE_PATH, 'w') as f:
        json.dump(refs, f, ensure_ascii=False, indent=2)
    
    # 显示结果
    print(f"\n✅ Found {len(refs)} interactive elements:\n")
    print(f"{'REF':<6} {'TEXT/DESC':<30} {'CLASS':<20} {'CENTER':<15} {'CLICK'}")
    print("-" * 90)
    
    for ref, info in refs.items():
        label = info['text'] or info['desc'] or '-'
        if len(label) > 28:
            label = label[:25] + "..."
        click = "✓" if info['clickable'] else ""
        center = f"({info['center'][0]}, {info['center'][1]})"
        print(f"{ref:<6} {label:<30} {info['class']:<20} {center:<15} {click}")
    
    return refs

def load_refs():
    """加载缓存的 ref 映射"""
    if not os.path.exists(REF_CACHE_PATH):
        print("❌ No refs cached. Run 'snapshot' first.")
        return None
    with open(REF_CACHE_PATH, 'r') as f:
        return json.load(f)

def tap(ref):
    """点击指定 ref 的元素"""
    refs = load_refs()
    if not refs:
        return False
    
    if ref not in refs:
        print(f"❌ Ref '{ref}' not found. Available: {', '.join(list(refs.keys())[:10])}...")
        return False
    
    info = refs[ref]
    x, y = info['center']
    label = info['text'] or info['desc'] or info['class']
    
    print(f"👆 Tapping {ref} ({label}) at ({x}, {y})...")
    out, code = run_adb(f"shell input tap {x} {y}")
    
    if code == 0:
        print(f"✅ Tapped!")
        return True
    else:
        print(f"❌ Failed: {out}")
        return False

def find(keyword):
    """搜索包含关键词的元素"""
    refs = load_refs()
    if not refs:
        return
    
    print(f"🔍 Searching for '{keyword}'...\n")
    found = []
    
    for ref, info in refs.items():
        text = (info['text'] + ' ' + info['desc']).lower()
        if keyword.lower() in text:
            found.append((ref, info))
    
    if not found:
        print(f"❌ No elements found containing '{keyword}'")
        return
    
    print(f"{'REF':<6} {'TEXT/DESC':<40} {'CENTER':<15}")
    print("-" * 65)
    for ref, info in found:
        label = info['text'] or info['desc']
        if len(label) > 38:
            label = label[:35] + "..."
        center = f"({info['center'][0]}, {info['center'][1]})"
        print(f"{ref:<6} {label:<40} {center:<15}")

def open_uri(uri):
    """打开 Deep Link"""
    print(f"🔗 Opening: {uri}")
    out, code = run_adb(f'shell am start -a android.intent.action.VIEW -d "{uri}"')
    if code == 0 and "Starting" in out:
        print("✅ Opened!")
        return True
    else:
        print(f"❌ Failed: {out}")
        return False

def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return
    
    cmd = sys.argv[1].lower()
    
    if cmd == "snapshot":
        snapshot()
    elif cmd == "tap" and len(sys.argv) >= 3:
        tap(sys.argv[2])
    elif cmd == "find" and len(sys.argv) >= 3:
        find(sys.argv[2])
    elif cmd == "open" and len(sys.argv) >= 3:
        open_uri(sys.argv[2])
    else:
        print(__doc__)

if __name__ == "__main__":
    main()