"""Capture public screenshots by operating the actual embedded keyboard UI."""
import argparse
from pathlib import Path
import re
import subprocess
import time
import xml.etree.ElementTree as ET


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--serial', required=True)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)

    def adb(*command):
        return subprocess.check_output(['adb', '-s', args.serial, *command], timeout=30)

    def tap(value, description=False):
        adb('shell', 'uiautomator', 'dump', '/sdcard/crab-capture.xml')
        tree = ET.fromstring(adb('exec-out', 'cat', '/sdcard/crab-capture.xml'))
        for node in tree.iter('node'):
            if node.get('content-desc' if description else 'text') == value:
                x1, y1, x2, y2 = map(int, re.findall(r'\d+', node.get('bounds')))
                adb('shell', 'input', 'tap', str((x1 + x2)//2), str((y1 + y2)//2))
                time.sleep(.22)
                return
        raise RuntimeError('Missing UI control: ' + value)

    def shot(name):
        time.sleep(2)
        (args.output / (name + '.png')).write_bytes(adb('exec-out', 'screencap', '-p'))

    adb('shell', 'input', 'keyevent', '224')
    adb('shell', 'wm', 'dismiss-keyguard')
    adb('shell', 'am', 'force-stop', 'com.sydxky.crabinput')
    adb('shell', 'am', 'start', '-n', 'com.sydxky.crabinput/.MainActivity')
    time.sleep(3)
    shot('android-home')
    tap('demo_text', True)
    tap('26键')
    shot('android-qwerty')
    for letter in 'shi':
        tap(letter)
    tap('更多候选', True)
    shot('android-candidates-16')
    tap('返回键盘', True)
    # Begin a clean composition by closing and reopening the real editor.
    tap('收起键盘', True)
    tap('demo_text', True)
    tap('14键')
    shot('android-fourteen')
    for key in ['BN', 'UI', 'GH', 'AS', 'OP']:
        tap(key)
    tap('更多候选', True)
    shot('android-phrases')
    tap('你好')
    tap('收起键盘', True)
    shot('android-result')
    adb('shell', 'am', 'start', '-n', 'com.sydxky.crabinput.host/.HostActivity')
    time.sleep(2)
    tap('fixture_editor', True)
    tap('26键')
    for letter in 'nihao':
        tap(letter)
    tap('你好')
    tap('收起键盘', True)
    shot('android-host')
    adb('shell', 'rm', '-f', '/sdcard/crab-capture.xml')
    print('Captured 7 real-device UI states; no system IME changes.')


if __name__ == '__main__':
    main()