"""Explicit-target screenshot, UI inspection and test runner (Python standard library).

phone = Android phone via ADB; watch-health = official paired-watch screenshot;
watch-import = image selected in the phone APK; inspector = explicitly supplied
Glyphix Inspector endpoint, whose device identity must be verified separately.
"""
import argparse
from datetime import datetime, timezone, timedelta
import hashlib
import json
from pathlib import Path
import re
import struct
import subprocess
import sys
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

PACKAGE = 'com.sydxky.rtos.bridge'
HEALTH = 'com.heytap.health'
PNG = b'\x89PNG\r\n\x1a\n'


def png_info(data):
    if len(data) < 33 or data[:8] != PNG or data[12:16] != b'IHDR':
        raise ValueError('Capture did not return a PNG')
    width, height = struct.unpack('>II', data[16:24])
    if not 1 <= width <= 8192 or not 1 <= height <= 8192:
        raise ValueError('Invalid capture dimensions')
    return {'width': width, 'height': height, 'sha256': hashlib.sha256(data).hexdigest()}


def save_png(path, data):
    info = png_info(data)
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(data)
    return dict(info, file=str(path.resolve()))


def jpeg_info(data):
    if not data.startswith(b'\xff\xd8') or not data.endswith(b'\xff\xd9'):
        raise ValueError('Capture did not return a complete JPEG')
    offset = 2
    while offset + 4 <= len(data):
        if data[offset] != 255:
            raise ValueError('Invalid JPEG marker')
        while offset < len(data) and data[offset] == 255:
            offset += 1
        marker = data[offset]
        offset += 1
        if marker in (0xDA, 0xD9):
            break
        size = int.from_bytes(data[offset:offset + 2], 'big')
        if size < 2 or offset + size > len(data):
            raise ValueError('Invalid JPEG segment')
        if marker in (0xC0, 0xC1, 0xC2, 0xC3, 0xC5, 0xC6, 0xC7, 0xC9, 0xCA, 0xCB, 0xCD, 0xCE, 0xCF):
            if size < 8:
                break
            height, width = struct.unpack('>HH', data[offset + 3:offset + 7])
            if not 1 <= width <= 8192 or not 1 <= height <= 8192:
                break
            return {'width': width, 'height': height, 'format': 'JPEG', 'sha256': hashlib.sha256(data).hexdigest()}
        offset += size
    raise ValueError('JPEG dimensions not found')


def save_watch_image(path, data):
    if data.startswith(PNG):
        info, suffixes = png_info(data), ('.png',)
    else:
        info, suffixes = jpeg_info(data), ('.jpg', '.jpeg')
    path = Path(path)
    if path.suffix.lower() not in suffixes:
        path = path.with_suffix(suffixes[0])
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(data)
    return dict(info, file=str(path.resolve()))


class Phone:
    scope = 'android_phone_ui'

    def __init__(self, serial, adb='adb'):
        if not serial:
            raise ValueError('--serial is required for an ADB target')
        self.adb = [adb, '-s', serial]

    def run(self, *args, binary=False, timeout=35):
        r = subprocess.run(self.adb + list(args), capture_output=True, timeout=timeout)
        if r.returncode:
            raise RuntimeError(r.stderr.decode('utf-8', 'replace').strip() or 'ADB operation failed')
        return r.stdout if binary else r.stdout.decode('utf-8', 'replace').strip()

    def unlocked(self):
        policy = self.run('shell', 'dumpsys', 'window', 'policy')
        if re.search(r'\bshowing=true|mIsShowing=true|mInputRestricted=true', policy):
            raise RuntimeError('Unlock the phone locally, then retry. UI automation does not unlock credentials.')

    def tree(self):
        self.unlocked()
        self.run('shell', 'uiautomator', 'dump', '/data/local/tmp/rtos-ui.xml')
        root = ET.fromstring(self.run('shell', 'cat', '/data/local/tmp/rtos-ui.xml'))
        for node in root.iter('node'):
            if node.get('password') == 'true':
                node.set('text', '[redacted]')
                node.set('content-desc', '[redacted]')
        return ET.tostring(root, encoding='unicode')

    def screenshot(self, path):
        self.unlocked()
        return save_png(path, self.run('exec-out', 'screencap', '-p', binary=True))

    @staticmethod
    def find(tree, selector):
        matches = []
        if selector[0] not in ('id', 'text', 'description'):
            raise ValueError('Phone UI selectors support id, text or description')
        for node in ET.fromstring(tree).iter('node'):
            key, value = selector
            attribute = {'id': 'resource-id', 'text': 'text', 'description': 'content-desc'}[key]
            bounds = list(map(int, re.findall(r'-?\d+', node.get('bounds', ''))))
            visible = len(bounds) == 4 and bounds[2] > bounds[0] and bounds[3] > bounds[1] and node.get('visible-to-user') != 'false'
            if node.get(attribute) == value and visible:
                matches.append(node)
        if len(matches) > 1:
            raise ValueError('UI selector is ambiguous; use a unique resource id or visible text')
        return matches[0] if matches else None

    def click(self, selector):
        node = self.find(self.tree(), selector)
        if node is None or node.get('enabled') != 'true':
            raise ValueError('Enabled UI element not found')
        x1, y1, x2, y2 = map(int, re.findall(r'-?\d+', node.get('bounds', '')))
        self.run('shell', 'input', 'tap', str((x1 + x2) // 2), str((y1 + y2) // 2))

    def scroll(self, selector, dx, dy):
        node = self.find(self.tree(), selector)
        if node is None:
            raise ValueError('UI element not found')
        x1, y1, x2, y2 = map(int, re.findall(r'-?\d+', node.get('bounds', '')))
        x, y = (x1 + x2) // 2, (y1 + y2) // 2
        self.run('shell', 'input', 'swipe', str(x), str(y),
                 str(max(x1 + 1, min(x2 - 1, x + dx))),
                 str(max(y1 + 1, min(y2 - 1, y + dy))), '500')

    def exists(self, selector):
        return self.find(self.tree(), selector) is not None


class Inspector:
    scope = 'glyphix_inspector_endpoint_identity_not_attested'

    def __init__(self, url):
        parsed = urllib.parse.urlsplit(url or '')
        if parsed.scheme != 'http' or parsed.hostname not in ('127.0.0.1', 'localhost') or parsed.username or parsed.password:
            raise ValueError('--inspector must be an explicit loopback HTTP URL; forward a device endpoint locally first')
        self.base = url.rstrip('/')
        self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def request(self, path, params=None, method='GET', binary=False):
        suffix = '?' + urllib.parse.urlencode(params) if params else ''
        req = urllib.request.Request(self.base + path + suffix, method=method)
        with self.opener.open(req, timeout=15) as response:
            data = response.read(8 * 1024 * 1024 + 1)
        if len(data) > 8 * 1024 * 1024:
            raise ValueError('Inspector response exceeds 8 MiB')
        if binary:
            return data
        try:
            return json.loads(data)
        except (ValueError, UnicodeDecodeError):
            return data.decode('utf-8', 'replace')

    def snapshot(self, image=None):
        try:
            tree = self.request('/snapshot.json')
            if image:
                name = tree['frames'][-1]['snapshot']
                if not isinstance(name, str) or not re.fullmatch(r'[A-Za-z0-9_.-]+', name):
                    raise ValueError('Invalid Inspector snapshot name')
                result = save_png(image, self.request('/snapshot/' + name, binary=True))
                Path(image).with_suffix('.tree.json').write_text(json.dumps(tree, ensure_ascii=False, indent=2), encoding='utf-8')
                return result
            return tree
        finally:
            # Snapshot pauses the runtime. Always resume, including image/JSON errors.
            self.request('/control', {'pause': 'false'}, 'POST')

    def tree(self):
        return self.snapshot()

    def screenshot(self, path):
        return self.snapshot(path)

    @staticmethod
    def locator(selector):
        key, value = selector
        if key not in ('id', 'text', 'tag', 'class') or not value or len(value) > 200:
            raise ValueError('Use id, text, tag or class with 1-200 characters')
        if key == 'text':
            # Escape regex syntax as literal text; never concatenate executable user code.
            pattern = '^' + re.sub(r'([\\.^$*+?{}\[\]|()])', r'\\\1', value) + '$'
            expression = 'getHasText(new RegExp(' + json.dumps(pattern, ensure_ascii=True) + '))'
        else:
            expression = {'id': 'getById', 'tag': 'getByTag', 'class': 'getByClass'}[key] + '(' + json.dumps(value, ensure_ascii=True) + ')'
        return "require('@system.test').default." + expression

    def action(self, selector, method):
        # Public Locator methods return promises. Poll only our tagged completion event.
        tag = 'RTOS_UI_' + __import__('uuid').uuid4().hex
        script = '(' + self.locator(selector) + ').' + method + ".then(function(){console.log('" + tag + "_PASS')},function(){console.log('" + tag + "_FAIL')})"
        self.request('/fn/eval', {'lang': 'js', 'script': script}, 'POST')
        deadline = time.monotonic() + 15
        while time.monotonic() < deadline:
            for row in self.request('/logs'):
                if row.get('tag') == 'eval':
                    continue
                body = row.get('body', '')
                if tag + '_FAIL' in body:
                    raise RuntimeError('Glyphix UI action failed')
                if tag + '_PASS' in body:
                    return
            time.sleep(.15)
        raise TimeoutError('No Glyphix UI completion marker; inspect the endpoint and target')

    def click(self, selector):
        self.action(selector, 'click()')

    def scroll(self, selector, dx, dy):
        self.action(selector, 'scroll(' + str(int(dx)) + ',' + str(int(dy)) + ',0.5)')

    def exists(self, selector):
        try:
            self.action(selector, 'wait()')
            return True
        except (RuntimeError, TimeoutError):
            return False


def imported_watch_capture(phone, path):
    metadata = json.loads(phone.run('shell', 'run-as', PACKAGE, 'cat', 'files/ui/watch-capture.json'))
    data = phone.run('exec-out', 'run-as', PACKAGE, 'cat', 'files/ui/watch-capture.png', binary=True)
    if png_info(data)['sha256'] != metadata['sha256']:
        raise ValueError('Imported capture integrity mismatch')
    result = save_png(path, data)
    return dict(result, provenance=metadata)


def health_screenshot(phone, path):
    """Navigate the official UI. Do not call its non-exported activities or private services."""
    phone.unlocked()
    # Health may reuse an id or add numbered filenames. Track content as well
    # as modification time; public images outside its directory are ignored.
    def captures():
        rows = phone.run('shell', 'content query --uri content://media/external/images/media --projection _id:_data:width:height:date_modified')
        images = {}
        for row in rows.splitlines():
            match = re.search(r'_id=(\d+), _data=((?:/storage/emulated/0|/sdcard)/Pictures/com\.heytap\.health/screenshots/watch_screen_shot(?: \(\d+\))?\.(?:jpeg|jpg|png)), width=(\d+), height=(\d+), date_modified=(\d+)', row)
            if not match:
                continue
            image_id, remote, width, height, modified = match.groups()
            if not all(200 <= int(value) <= 1024 for value in (width, height)):
                continue
            # The content read subcommand is rejected on this P30PRO provider;
            # ADB can read the exact public file without changing permissions.
            try:
                data = phone.run('exec-out', 'cat', remote, binary=True)
                info = png_info(data) if data.startswith(PNG) else jpeg_info(data)
            except (ValueError, RuntimeError):
                continue  # File may be mid-write while Health saves it.
            images[image_id] = (modified, info['sha256'], data)
        return images

    before = captures()
    phone.run('shell', 'am', 'start', '-n', HEALTH + '/.oobe.LaunchActivity')
    device_selected = False
    for _ in range(12):
        tree = phone.tree()
        button = ('id', HEALTH + ':id/device_settings_tv_screen_shot')
        if phone.find(tree, button) is not None:
            phone.click(button)
            break
        for selector in (('text', '暂不'), ('text', '手表截屏'), ('text', '更多功能')):
            if phone.find(tree, selector) is not None:
                phone.click(selector)
                break
        else:
            selector = ('id', HEALTH + ':id/app_navigation_manager')
            if not device_selected and phone.find(tree, selector) is not None:
                phone.click(selector)
                device_selected = True
            elif device_selected:
                # Use current display dimensions, avoiding the system navigation edge.
                sizes = re.findall(r'(\d+)x(\d+)', phone.run('shell', 'wm', 'size'))
                w, h = map(int, sizes[-1])
                phone.run('shell', 'input', 'swipe', str(w // 2), str(h * 4 // 5), str(w // 2), str(h // 3), '500')
            else:
                raise RuntimeError('Open the paired device page in the health app, then retry')
        time.sleep(.6)
    else:
        raise RuntimeError('Official watch screenshot control not found; no screenshot was captured')
    deadline = time.monotonic() + 50
    while time.monotonic() < deadline:
        for image_id, (modified, sha, data) in captures().items():
            if image_id not in before or (modified, sha) != before[image_id][:2]:
                return dict(save_watch_image(path, data), provenance='official_health_ui_watch_screenshot')
        time.sleep(1)
    raise TimeoutError('No new watch screenshot image was exported. In the APK, select the official screenshot and use watch-import.')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--target', required=True, choices=['phone', 'watch-health', 'watch-import', 'inspector'])
    parser.add_argument('--serial')
    parser.add_argument('--adb', default='adb')
    parser.add_argument('--inspector', help='Explicit locally forwarded Glyphix Inspector endpoint')
    sub = parser.add_subparsers(dest='command', required=True)
    sub.add_parser('capabilities')
    sub.add_parser('screenshot').add_argument('--output', type=Path, required=True)
    sub.add_parser('tree').add_argument('--output', type=Path, required=True)
    for name in ('click', 'scroll', 'assert-visible'):
        p = sub.add_parser(name)
        p.add_argument('--by', required=True, choices=['id', 'text', 'description', 'tag', 'class'])
        p.add_argument('--value', required=True)
        if name == 'scroll':
            p.add_argument('--dx', type=int, default=0)
            p.add_argument('--dy', type=int, required=True)
    p = sub.add_parser('run')
    p.add_argument('plan', type=Path)
    p.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    result = {'time': datetime.now(timezone(timedelta(hours=8))).isoformat(), 'target': args.target}
    if args.command == 'capabilities':
        result.update(screenshot_implemented=True, ui_tree_implemented=args.target in ('phone', 'inspector'),
                      control_implemented=args.target in ('phone', 'inspector'),
                      availability='conditional_on_selected_backend_and_target_connection',
                      note='Inspector requires a reachable runtime endpoint; Wear Engine 263 alone does not expose this endpoint.')
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return
    backend = Inspector(args.inspector) if args.target == 'inspector' else Phone(args.serial, args.adb)
    if args.target in ('watch-health', 'watch-import'):
        if args.command != 'screenshot':
            raise ValueError('This watch backend supports screenshot only. General watch control requires a confirmed Inspector endpoint.')
        result.update((health_screenshot if args.target == 'watch-health' else imported_watch_capture)(backend, args.output))
    elif args.command == 'screenshot':
        result.update(backend.screenshot(args.output), evidence_scope=backend.scope)
    elif args.command == 'tree':
        tree = backend.tree()
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(tree if isinstance(tree, str) else json.dumps(tree, ensure_ascii=False, indent=2), encoding='utf-8')
        result.update(file=str(args.output.resolve()), evidence_scope=backend.scope)
    elif args.command == 'run':
        plan = json.loads(args.plan.read_text(encoding='utf-8'))
        steps = plan.get('steps', [])
        if not 1 <= len(steps) <= 50:
            raise ValueError('Test plan must contain 1-50 steps')
        args.output.mkdir(parents=True, exist_ok=True)
        result.update(evidence_scope=backend.scope, steps=[], passed=True)
        for index, step in enumerate(steps):
            entry = {'index': index, 'action': step.get('action'), 'passed': False}
            try:
                selector = (step.get('by'), step.get('value'))
                if step['action'] == 'click': backend.click(selector)
                elif step['action'] == 'scroll': backend.scroll(selector, max(-2000, min(2000, int(step.get('dx', 0)))), max(-2000, min(2000, int(step.get('dy', 0)))))
                elif step['action'] == 'assert-visible':
                    if not backend.exists(selector): raise AssertionError('UI element not visible')
                elif step['action'] == 'screenshot': entry['capture'] = backend.screenshot(args.output / (str(index) + '.png'))
                else: raise ValueError('Unsupported test action')
                entry['passed'] = True
            except Exception as error:
                entry['error'] = type(error).__name__
                result['passed'] = False
            result['steps'].append(entry)
            if not entry['passed']: break
        (args.output / 'report.json').write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding='utf-8')
    else:
        selector = (args.by, args.value)
        if args.command == 'click': backend.click(selector)
        elif args.command == 'scroll': backend.scroll(selector, args.dx, args.dy)
        elif not backend.exists(selector): raise AssertionError('UI element not visible')
        result.update(passed=True, evidence_scope=backend.scope)
    print(json.dumps(result, ensure_ascii=False, indent=2))
    if result.get('passed') is False: return 2
    return 0


if __name__ == '__main__':
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')
    try:
        sys.exit(main())
    except (OSError, ValueError, RuntimeError, AssertionError, subprocess.SubprocessError) as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)