#!/usr/bin/env python3
import json
import os
import re
import subprocess
import sys
import threading
import time
import urllib.request

import websocket

CHROME = os.environ.get("CHROME", "/opt/chrome-linux64/chrome")
PORT = int(os.environ.get("CDP_PORT", "9225"))
URL = sys.argv[1]
PROFILE = sys.argv[2] if len(sys.argv) > 2 else "/tmp/cdp-run-profile"
EVENT_LOG = []
DEBUG_OBJECT_ADDRESSES = []


def renderer_chrome_bases():
    bases = set()
    profile_marker = ("--user-data-dir=" + os.path.abspath(PROFILE)).encode()
    for entry in os.listdir("/proc"):
        if not entry.isdigit():
            continue
        try:
            with open(f"/proc/{entry}/cmdline", "rb") as command_file:
                command = command_file.read()
            if b"--type=renderer" not in command or profile_marker not in command:
                continue
            with open(f"/proc/{entry}/maps", "r", encoding="utf-8") as maps:
                for line in maps:
                    fields = line.split()
                    if (len(fields) >= 6 and fields[2] == "00000000" and
                            fields[-1] == CHROME):
                        bases.add(int(fields[0].split("-", 1)[0], 16))
                        break
        except (FileNotFoundError, PermissionError, ProcessLookupError):
            continue
    return sorted(bases)


def request(ws, ident, method, params=None):
    message = {"id": ident, "method": method}
    if params is not None:
        message["params"] = params
    ws.send(json.dumps(message))
    while True:
        response = json.loads(ws.recv())
        if response.get("id") == ident:
            if "error" in response:
                raise RuntimeError(response["error"])
            return response.get("result", {})
        method_name = response.get("method", "")
        if method_name in (
            "Network.loadingFailed", "Runtime.exceptionThrown",
            "Log.entryAdded", "Inspector.targetCrashed"):
            EVENT_LOG.append(response)


debug_chrome = bool(os.environ.get("DEBUG_CHROME"))
chrome_output_path = os.environ.get("CHROME_OUTPUT_FILE")
chrome_output_file = (
    open(chrome_output_path, "w", encoding="utf-8")
    if chrome_output_path else None
)
chrome = subprocess.Popen(
    [
        CHROME,
        "--headless=new",
        "--no-sandbox",
        "--disable-crashpad",
        "--disable-breakpad",
        "--disable-crash-reporter",
        "--noerrdialogs",
        "--no-first-run",
        "--ozone-platform=headless",
        "--ozone-override-screen-size=800,600",
        "--use-angle=swiftshader-webgl",
        "--remote-allow-origins=*",
        f"--remote-debugging-port={PORT}",
        f"--user-data-dir={PROFILE}",
        *(["--js-flags=" + os.environ["CHROME_JS_FLAGS"]]
          if os.environ.get("CHROME_JS_FLAGS") else []),
        URL,
    ],
    stdout=(chrome_output_file if chrome_output_file else
            (subprocess.PIPE if debug_chrome else subprocess.DEVNULL)),
    stderr=(subprocess.STDOUT if (debug_chrome or chrome_output_file) else
            subprocess.DEVNULL),
    text=debug_chrome,
    bufsize=1 if debug_chrome else -1,
)


def pump_chrome_output():
    for line in chrome.stdout:
        print(line, end="", flush=True)
        match = re.match(r"^0x([0-9a-fA-F]+) <String\[", line)
        if match:
            DEBUG_OBJECT_ADDRESSES.append(int(match.group(1), 16))


if debug_chrome:
    threading.Thread(target=pump_chrome_output, daemon=True).start()

try:
    targets = None
    for _ in range(100):
        try:
            with urllib.request.urlopen(
                f"http://127.0.0.1:{PORT}/json", timeout=1
            ) as response:
                targets = json.load(response)
            break
        except Exception:
            time.sleep(0.1)
    if targets is None:
        raise RuntimeError("Chrome debugging endpoint did not start")
    target = next(item for item in targets if item["type"] == "page")
    ws = websocket.create_connection(
        target["webSocketDebuggerUrl"],
        timeout=int(os.environ.get("CDP_SOCKET_TIMEOUT", "300")))
    try:
        ident = 1
        request(ws, ident, "Network.enable")
        ident += 1
        request(ws, ident, "Runtime.enable")
        ident += 1
        request(ws, ident, "Log.enable")
        ident += 1
        title = None
        previous_title = None
        stop_titles = set(os.environ.get(
            "STOP_TITLES", "PASS,FAIL,ERROR").split(","))
        stop_title_prefixes = tuple(filter(None, os.environ.get(
            "STOP_TITLE_PREFIXES", "ERROR ").split(",")))
        inject_at_title = os.environ.get("INJECT_AT_TITLE")
        inject_expression = os.environ.get("INJECT_EXPRESSION")
        inject_script_paths = os.environ.get("INJECT_SCRIPT_PATHS")
        if inject_script_paths:
            injected_sources = []
            for inject_path in inject_script_paths.split(":"):
                with open(inject_path, "r", encoding="utf-8") as script_file:
                    injected_sources.append(script_file.read())
            inject_expression = "\n".join(injected_sources)
        injected = False
        set_debug_address_at = os.environ.get(
            "SET_LAST_DEBUG_ADDRESS_AT_TITLE")
        debug_address_set = False
        for _ in range(int(os.environ.get("MAX_POLLS", "3600"))):
            try:
                result = request(
                    ws,
                    ident,
                    "Runtime.evaluate",
                    {"expression": "document.title", "returnByValue": True},
                )
            except RuntimeError as error:
                if "Cannot find default execution context" not in str(error):
                    raise
                ident += 1
                time.sleep(0.1)
                continue
            ident += 1
            title = result["result"].get("value")
            if title != previous_title:
                print("stage_title=" + str(title), flush=True)
                previous_title = title
            if (not injected and inject_at_title and inject_expression and
                    title == inject_at_title):
                injected = True
                if os.environ.get("PRINT_RENDERER_BASE"):
                    print("renderer_chrome_bases=" + ",".join(
                        "0x" + format(base, "x")
                        for base in renderer_chrome_bases()), flush=True)
                inject_socket_timeout = int(os.environ.get(
                    "INJECT_SOCKET_TIMEOUT", "0"))
                if inject_socket_timeout > 0:
                    ws.settimeout(inject_socket_timeout)
                inject_result = request(
                    ws,
                    ident,
                    "Runtime.evaluate",
                    {"expression": inject_expression,
                     "returnByValue": True,
                     "awaitPromise": True},
                )
                ident += 1
                print("inject_result=" + json.dumps(
                    inject_result, separators=(",", ":")), flush=True)
                if inject_socket_timeout > 0:
                    ws.settimeout(int(os.environ.get(
                        "CDP_SOCKET_TIMEOUT", "300")))
                if (os.environ.get("SET_LAST_DEBUG_ADDRESS_AFTER_INJECT") and
                        DEBUG_OBJECT_ADDRESSES):
                    compressed_address = (
                        (DEBUG_OBJECT_ADDRESSES[-1] & 0xffffffff) - 1
                    ) & 0xffffffff
                    post_inject_result = request(
                        ws,
                        ident,
                        "Runtime.evaluate",
                        {"expression": (
                            "regexpPreparedWorkerInfo.address=" +
                            str(compressed_address)),
                         "returnByValue": True},
                    )
                    ident += 1
                    print("post_inject_debug_address=0x" +
                          format(compressed_address, "x"), flush=True)
                    print("post_inject_debug_result=" + json.dumps(
                        post_inject_result, separators=(",", ":")),
                          flush=True)
            if (not debug_address_set and set_debug_address_at and
                    title == set_debug_address_at and DEBUG_OBJECT_ADDRESSES):
                debug_address_set = True
                compressed_address = (
                    (DEBUG_OBJECT_ADDRESSES[-1] & 0xffffffff) - 1
                ) & 0xffffffff
                if os.environ.get("DEBUG_ADDRESS_ASSIGN_ONLY"):
                    debug_expression = (
                        "regexpPreparedWorkerInfo.address=" +
                        str(compressed_address))
                else:
                    debug_expression = (
                        "typeof continuePostRootRace==='function'?" +
                        "continuePostRootRace(" + str(compressed_address) +
                        "):regexpPreparedWorkerInfo.address=" +
                        str(compressed_address))
                debug_result = request(
                    ws,
                    ident,
                    "Runtime.evaluate",
                    {"expression": debug_expression,
                     "returnByValue": True},
                )
                ident += 1
                print("debug_object_address=0x" +
                      format(compressed_address, "x"), flush=True)
                print("debug_address_result=" + json.dumps(
                    debug_result, separators=(",", ":")), flush=True)
            if (title in stop_titles or
                    any(str(title).startswith(prefix)
                        for prefix in stop_title_prefixes)):
                stop_delay_ms = int(os.environ.get("STOP_DELAY_MS", "0"))
                if stop_delay_ms > 0:
                    time.sleep(stop_delay_ms / 1000)
                break
            time.sleep(0.1)
        body = request(
            ws,
            ident,
            "Runtime.evaluate",
            {"expression": "document.body.innerText", "returnByValue": True},
        )["result"].get("value", "")
        print("title=" + str(title))
        print(body)
        for event in EVENT_LOG:
            print("cdp_event=" + json.dumps(event, separators=(",", ":")))
        cache_name = os.environ.get("ANALYZE_CACHE")
        if cache_name:
            ident += 1
            origin = URL.split("/", 3)[:3]
            security_origin = "/".join(origin)
            names = request(
                ws,
                ident,
                "CacheStorage.requestCacheNames",
                {"securityOrigin": security_origin},
            )
            ident += 1
            cache = next(
                item for item in names["caches"]
                if item["cacheName"] == cache_name
            )
            entries = request(
                ws,
                ident,
                "CacheStorage.requestEntries",
                {"cacheId": cache["cacheId"], "skipCount": 0,
                 "pageSize": 5000},
            )["cacheDataEntries"]
            by_round = {}
            for item in entries:
                parts = item["requestURL"].rstrip("/").split("/")
                if len(parts) >= 2 and parts[-1] in ("A.js", "B.js"):
                    by_round.setdefault(int(parts[-2]), {})[
                        parts[-1][0]
                    ] = item["responseTime"]
            equal = [
                round_no for round_no, values in by_round.items()
                if values.get("A") == values.get("B")
            ]
            print("analyzed_entry_count=" + str(len(entries)))
            print("equal_response_rounds=" + ",".join(map(str, equal)))
            print("sample_response_times=" + json.dumps(
                dict(sorted(by_round.items())[:20]), separators=(",", ":")))
            target_cache_name = os.environ.get(
                "TARGET_CACHE", "code-cache-collision-target-v1")
            target_cache = next((
                item for item in names["caches"]
                if item["cacheName"] == target_cache_name
            ), None)
            if target_cache:
                ident += 1
                target_entries = request(
                    ws,
                    ident,
                    "CacheStorage.requestEntries",
                    {"cacheId": target_cache["cacheId"], "skipCount": 0,
                     "pageSize": 5000},
                )["cacheDataEntries"]
                target_times = {
                    int(item["requestURL"].rsplit("/", 1)[1][:-3]):
                    item["responseTime"] for item in target_entries
                }
                equals_a = sum(
                    target_times.get(round_no) == values.get("A")
                    for round_no, values in by_round.items()
                )
                equals_b = sum(
                    target_times.get(round_no) == values.get("B")
                    for round_no, values in by_round.items()
                )
                print("target_entry_count=" + str(len(target_entries)))
                print("target_time_equals_source_A=" + str(equals_a))
                print("target_time_equals_source_B=" + str(equals_b))
    finally:
        try:
            ws.send(json.dumps({"id": 999999, "method": "Browser.close"}))
        except Exception:
            pass
        try:
            chrome.wait(timeout=10)
        except subprocess.TimeoutExpired:
            pass
        ws.close()
finally:
    if chrome_output_file:
        chrome_output_file.close()
    if chrome.poll() is None:
        chrome.terminate()
        try:
            chrome.wait(timeout=5)
        except subprocess.TimeoutExpired:
            chrome.kill()