#!/usr/bin/env python3
import argparse
import json
import os
import socket
import subprocess
import time
import urllib.parse

import websocket

ADB_HOST = os.environ.get("ADB_SERVER_HOST", "127.0.0.1")
ADB_PORT = int(os.environ.get("ADB_SERVER_PORT", "15037"))
SOCKET_NAME = os.environ.get("CHROME_DEVTOOLS_SOCKET", "chrome_devtools_remote")


def adb_serial():
    configured = os.environ.get("ANDROID_SERIAL")
    if configured:
        return configured
    output = subprocess.run(
        ["adb", "devices"], check=True, capture_output=True, text=True
    ).stdout
    devices = [
        line.split()[0]
        for line in output.splitlines()[1:]
        if len(line.split()) >= 2 and line.split()[1] == "device"
    ]
    if len(devices) != 1:
        raise RuntimeError(
            "Set ANDROID_SERIAL when the ADB server does not expose exactly one device."
        )
    return devices[0]


ADB_SERIAL = adb_serial()


def _request(sock, service):
    payload = service.encode()
    sock.sendall(f"{len(payload):04x}".encode() + payload)
    status = sock.recv(4)
    if status != b"OKAY":
        if status == b"FAIL":
            length = int(sock.recv(4), 16)
            raise RuntimeError(sock.recv(length).decode())
        raise RuntimeError(f"ADB service {service!r} returned {status!r}")


def adb_socket(service):
    sock = socket.create_connection((ADB_HOST, ADB_PORT), timeout=10)
    _request(sock, f"host:transport:{ADB_SERIAL}")
    _request(sock, service)
    return sock


def targets():
    sock = adb_socket(f"localabstract:{SOCKET_NAME}")
    sock.sendall(b"GET /json HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
    response = b""
    while b"\r\n\r\n" not in response:
        response += sock.recv(65536)
    headers, body = response.split(b"\r\n\r\n", 1)
    content_length = None
    for line in headers.split(b"\r\n"):
        if line.lower().startswith(b"content-length:"):
            content_length = int(line.split(b":", 1)[1])
            break
    if content_length is None:
        raise RuntimeError("DevTools response has no Content-Length")
    while len(body) < content_length:
        body += sock.recv(65536)
    sock.close()
    return json.loads(body[:content_length])


class CDP:
    def __init__(self, target):
        path = urllib.parse.urlparse(target["webSocketDebuggerUrl"]).path
        sock = adb_socket(f"localabstract:{SOCKET_NAME}")
        self.ws = websocket.create_connection(
            f"ws://localhost{path}", socket=sock, timeout=15, suppress_origin=True
        )
        self.next_id = 1

    def close(self):
        self.ws.close()

    def command(self, method, params=None):
        command_id = self.next_id
        self.next_id += 1
        self.ws.send(json.dumps({"id": command_id, "method": method, "params": params or {}}))
        while True:
            message = json.loads(self.ws.recv())
            if message.get("id") == command_id:
                if "error" in message:
                    raise RuntimeError(message["error"])
                return message.get("result", {})

    def evaluate(self, expression, await_promise=True):
        result = self.command(
            "Runtime.evaluate",
            {
                "expression": expression,
                "returnByValue": True,
                "awaitPromise": await_promise,
                "userGesture": True,
            },
        )
        value = result.get("result", {})
        if value.get("subtype") == "error":
            raise RuntimeError(value.get("description", value))
        return value.get("value")


def find_target(fragment):
    matches = [target for target in targets() if fragment in target.get("url", "")]
    if not matches:
        raise RuntimeError(f"No target containing {fragment!r}: {[t.get('url') for t in targets()]}")
    return matches[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--list", action="store_true")
    parser.add_argument("--url-contains")
    parser.add_argument("--eval")
    parser.add_argument("--navigate")
    args = parser.parse_args()
    if args.list:
        print(json.dumps(targets(), indent=2))
        return
    target = find_target(args.url_contains)
    cdp = CDP(target)
    try:
        if args.navigate:
            print(json.dumps(cdp.command("Page.navigate", {"url": args.navigate}), indent=2))
            time.sleep(2)
        if args.eval:
            print(json.dumps(cdp.evaluate(args.eval), indent=2, ensure_ascii=False))
    finally:
        cdp.close()


if __name__ == "__main__":
    main()