"""
Generalized CDP flow runner.

Drives an ALREADY-RUNNING debug Chrome (--remote-debugging-port=9222) through an
ordered list of steps: type text into inputs found by XPath, click elements by
XPath, press keys, wait, wait-for-element, navigate. Typing and clicking use
real CDP Input events (Input.insertText / Input.dispatchMouseEvent) so React
and friends see genuine user input, with a JS fallback for zero-size elements.

Steps are given on the command line (order preserved) and/or in a JSON flow
file. Any number of each step type is allowed.

RUN (from WSL, using Windows Python so localhost:9222 is reachable natively):
    bash scripts/cdp-flow/run.sh --tab higgsfield \
        --type "//textarea[@name='prompt']" "a red fox" \
        --click "//button[.//span[text()='Generate']]" \
        --waitfor "//div[contains(@class,'result')]"

    bash scripts/cdp-flow/run.sh --flow scripts/cdp-flow/flows/login.json --var pw=secret

Flow file format — fully self-contained, no other flags needed. "tab" picks an
open tab by URL/title substring; "new_tab": "<url>" opens a tab in the REAL
profile; "temp_context" spawns an ISOLATED temp browser context (own cookies /
storage, nothing shared — see docs/TEMP_CHROME_CONTEXTS.md), opens its window
at "url", and runs the steps in there. Temp contexts are tracked in
scripts/temp-profile/contexts.json, so `bash scripts/temp-profile/run.sh
list|dispose|sweep` and the 24h TTL all apply. "dispose": true tears the
context down after a SUCCESSFUL flow (a failed flow always keeps it open for
inspection). "vars" are defaults for the {name} placeholders, overridable
per-run with --var; steps run in order:
    {
      "temp_context": {"url": "https://example.com/login", "label": "login-test",
                       "dispose": false},
      "timeout": 15,
      "delay": 0.3,
      "vars": {"user": "me@example.com"},
      "steps": [
        {"do": "waitfor", "xpath": "//input[@name='email']"},
        {"do": "type",    "xpath": "//input[@name='email']", "text": "{user}"},
        {"do": "type",    "xpath": "//input[@name='password']", "text": "{pw}"},
        {"do": "click",   "xpath": "//button[@type='submit']"},
        {"do": "press",   "key": "Enter"},
        {"do": "wait",    "seconds": 2}
      ]
    }

Notes:
  - "type" replaces the field's existing content (select-all first).
  - XPaths are evaluated in the top document only (no iframe descent).
  - Pick the tab with --tab <url-or-title substring> or open one with --new-tab.

Deps (Windows Python):  websocket-client
"""

import argparse
import json
import os
import re
import sys
import time
import urllib.request

import websocket  # websocket-client

for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8")
    except Exception:
        pass

CDP_HTTP = os.environ.get("CDP_HTTP", "http://localhost:9222")

KEYS = {
    "Enter":      {"key": "Enter", "code": "Enter", "windowsVirtualKeyCode": 13, "text": "\r"},
    "Tab":        {"key": "Tab", "code": "Tab", "windowsVirtualKeyCode": 9},
    "Escape":     {"key": "Escape", "code": "Escape", "windowsVirtualKeyCode": 27},
    "Backspace":  {"key": "Backspace", "code": "Backspace", "windowsVirtualKeyCode": 8},
    "Delete":     {"key": "Delete", "code": "Delete", "windowsVirtualKeyCode": 46},
    "Space":      {"key": " ", "code": "Space", "windowsVirtualKeyCode": 32, "text": " "},
    "ArrowUp":    {"key": "ArrowUp", "code": "ArrowUp", "windowsVirtualKeyCode": 38},
    "ArrowDown":  {"key": "ArrowDown", "code": "ArrowDown", "windowsVirtualKeyCode": 40},
    "ArrowLeft":  {"key": "ArrowLeft", "code": "ArrowLeft", "windowsVirtualKeyCode": 37},
    "ArrowRight": {"key": "ArrowRight", "code": "ArrowRight", "windowsVirtualKeyCode": 39},
}


class CDP:
    """Minimal synchronous CDP client on the browser websocket, with sessions."""

    def __init__(self):
        with urllib.request.urlopen(f"{CDP_HTTP}/json/version", timeout=5) as r:
            ws_url = json.load(r)["webSocketDebuggerUrl"]
        self.ws = websocket.create_connection(ws_url, max_size=None, timeout=30)
        self._id = 0

    def call(self, method, params=None, session_id=None, timeout=30):
        self._id += 1
        msg = {"id": self._id, "method": method, "params": params or {}}
        if session_id:
            msg["sessionId"] = session_id
        self.ws.send(json.dumps(msg))
        deadline = time.time() + timeout
        while time.time() < deadline:
            self.ws.settimeout(max(0.1, deadline - time.time()))
            try:
                reply = json.loads(self.ws.recv())
            except websocket.WebSocketTimeoutException:
                break
            if reply.get("id") == self._id:
                if "error" in reply:
                    raise RuntimeError(f"CDP {method}: {reply['error']}")
                return reply.get("result", {})
        raise TimeoutError(f"Timed out on {method}")

    def close(self):
        try:
            self.ws.close()
        except Exception:
            pass


class Page:
    """One attached page tab: evaluate JS + dispatch real input events."""

    def __init__(self, cdp, target_id):
        self.cdp = cdp
        self.session = cdp.call("Target.attachToTarget",
                                {"targetId": target_id, "flatten": True})["sessionId"]

    def eval(self, expr):
        res = self.cdp.call("Runtime.evaluate",
                            {"expression": expr, "returnByValue": True},
                            session_id=self.session)
        exc = res.get("exceptionDetails")
        if exc:
            raise RuntimeError(f"page JS error: {exc.get('text')} "
                               f"{(exc.get('exception') or {}).get('description', '')}")
        return res.get("result", {}).get("value")

    def input(self, method, params):
        self.cdp.call(method, params, session_id=self.session)


def js_find(xpath):
    return ("document.evaluate(%s, document, null, "
            "XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue"
            % json.dumps(xpath))


def find_visible(page, xpath):
    """Locate + scroll to the element; returns its center or a JS-click marker."""
    return page.eval(
        "(() => {"
        f"  const el = {js_find(xpath)};"
        "  if (!el) return {err: 'not found'};"
        "  el.scrollIntoView({block: 'center', inline: 'center'});"
        "  const r = el.getBoundingClientRect();"
        "  return {x: r.left + r.width / 2, y: r.top + r.height / 2,"
        "          zero: r.width === 0 || r.height === 0};"
        "})()")


def step_type(page, xpath, text):
    got = page.eval(
        "(() => {"
        f"  const el = {js_find(xpath)};"
        "  if (!el) return {err: 'not found'};"
        "  el.scrollIntoView({block: 'center', inline: 'center'});"
        "  el.focus();"
        "  try {"
        "    if (el.select) el.select();"
        "    else {"
        "      const rng = document.createRange(); rng.selectNodeContents(el);"
        "      const sel = getSelection(); sel.removeAllRanges(); sel.addRange(rng);"
        "    }"
        "  } catch (e) {}"
        "  return {ok: true};"
        "})()")
    if got.get("err"):
        raise RuntimeError(f"type: element {got['err']}: {xpath}")
    page.input("Input.insertText", {"text": text})


def step_click(page, xpath):
    got = find_visible(page, xpath)
    if got.get("err"):
        raise RuntimeError(f"click: element {got['err']}: {xpath}")
    if got.get("zero"):
        page.eval(f"(() => {{ const el = {js_find(xpath)}; el && el.click(); }})()")
        return "js-click (zero-size element)"
    x, y = got["x"], got["y"]
    page.input("Input.dispatchMouseEvent", {"type": "mouseMoved", "x": x, "y": y})
    for kind in ("mousePressed", "mouseReleased"):
        page.input("Input.dispatchMouseEvent", {
            "type": kind, "x": x, "y": y,
            "button": "left", "buttons": 1, "clickCount": 1,
        })
    return f"clicked at ({x:.0f}, {y:.0f})"


def step_press(page, key):
    spec = KEYS.get(key)
    if spec is None:
        if len(key) != 1:
            raise RuntimeError(f"press: unknown key '{key}' (known: {', '.join(KEYS)})")
        spec = {"key": key, "text": key}
    down = {"type": "keyDown", **spec}
    if "text" not in down:
        down["type"] = "rawKeyDown"
    page.input("Input.dispatchKeyEvent", down)
    page.input("Input.dispatchKeyEvent",
               {"type": "keyUp", **{k: v for k, v in spec.items() if k != "text"}})


def step_waitfor(page, xpath, timeout):
    deadline = time.time() + timeout
    while time.time() < deadline:
        got = page.eval(
            "(() => {"
            f"  const el = {js_find(xpath)};"
            "  if (!el) return false;"
            "  const r = el.getBoundingClientRect();"
            "  return r.width > 0 && r.height > 0;"
            "})()")
        if got:
            return
        time.sleep(0.25)
    raise TimeoutError(f"waitfor: element never appeared within {timeout:g}s: {xpath}")


def step_goto(page, url, timeout):
    page.cdp.call("Page.navigate", {"url": url}, session_id=page.session)
    time.sleep(0.5)
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            if page.eval("document.readyState") == "complete":
                return
        except Exception:
            pass  # document swapping out mid-navigation
        time.sleep(0.25)
    raise TimeoutError(f"goto: page never finished loading within {timeout:g}s: {url}")


# Same tracking file as scripts/temp-profile/temp_context.py, so contexts we
# spawn here show up in its list/dispose and get swept by its 24h TTL.
TEMP_STATE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                               "..", "temp-profile", "contexts.json")


def _load_temp_state():
    try:
        with open(TEMP_STATE_PATH, encoding="utf-8") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {}


def track_context(cid, label):
    state = _load_temp_state()
    state[cid] = {"created": time.time(), "label": label}
    with open(TEMP_STATE_PATH, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2)


def untrack_context(cid):
    state = _load_temp_state()
    if state.pop(cid, None) is not None:
        with open(TEMP_STATE_PATH, "w", encoding="utf-8") as f:
            json.dump(state, f, indent=2)


def dispose_context(cdp, cid):
    """Close the context's tabs one-by-one BEFORE disposing the context.

    Disposing a context that still has a live window has crashed the user's
    entire Chrome (see docs/TEMP_CHROME_CONTEXTS.md); never skip this path.
    """
    for t in cdp.call("Target.getTargets").get("targetInfos", []):
        if t.get("browserContextId") == cid and t.get("type") == "page":
            cdp.call("Target.closeTarget", {"targetId": t["targetId"]})
    time.sleep(0.5)
    cdp.call("Target.disposeBrowserContext", {"browserContextId": cid})
    untrack_context(cid)


def pick_tab(match):
    """Match an exact target-id prefix first, else a URL/title substring."""
    with urllib.request.urlopen(f"{CDP_HTTP}/json/list", timeout=10) as r:
        tabs = [t for t in json.load(r) if t.get("type") == "page"]

    def listing(ts):
        return "\n".join(f"  {t['id'][:8]}  {t['title'][:50]!r}  {t['url'][:80]}"
                         for t in ts)

    id_hits = [t for t in tabs if t["id"].lower().startswith(match.lower())]
    if len(id_hits) == 1:
        return id_hits[0]["id"]

    hits = [t for t in tabs
            if match.lower() in t.get("url", "").lower()
            or match.lower() in t.get("title", "").lower()]
    if not hits:
        raise RuntimeError(f"no open tab matches '{match}'. Open tabs:\n{listing(tabs)}")
    if len(hits) > 1:
        raise RuntimeError(f"'{match}' is ambiguous — rerun with one of these tab ids "
                           f"as --tab/\"tab\":\n{listing(hits)}")
    return hits[0]["id"]


class OrderedStep(argparse.Action):
    """Append (step-kind, values) to a shared list so CLI order is preserved."""

    def __call__(self, parser, ns, values, option_string=None):
        if getattr(ns, "steps", None) is None:
            ns.steps = []
        ns.steps.append((self.dest, values))


def load_flow(path, cli_vars):
    """Returns (config, steps). CLI --var values override the file's "vars"."""
    with open(path, encoding="utf-8") as f:
        flow = json.load(f)

    variables = {**flow.get("vars", {}), **cli_vars}

    def sub(s):
        # Replace {name} only for defined vars; leave every other brace alone
        # (XPaths, data: URLs, JS snippets may contain literal braces).
        if not isinstance(s, str):
            return s
        return re.sub(r"\{(\w+)\}",
                      lambda m: str(variables.get(m.group(1), m.group(0))), s)

    steps = []
    for i, raw in enumerate(flow.get("steps", []), 1):
        do = raw.get("do")
        if do == "type":
            steps.append(("type", [sub(raw["xpath"]), sub(raw["text"]),
                                   bool(raw.get("secret"))]))
        elif do == "click":
            steps.append(("click", sub(raw["xpath"])))
        elif do == "press":
            steps.append(("press", raw["key"]))
        elif do == "wait":
            steps.append(("wait", str(raw["seconds"])))
        elif do == "waitfor":
            steps.append(("waitfor", sub(raw["xpath"])))
        elif do == "goto":
            steps.append(("goto", sub(raw["url"])))
        else:
            raise RuntimeError(f"flow step {i}: unknown 'do': {do!r}")

    config = {k: flow[k] for k in ("tab", "new_tab", "temp_context", "timeout", "delay")
              if k in flow}
    tc = config.get("temp_context")
    if isinstance(tc, str):  # shorthand: just the URL
        config["temp_context"] = {"url": tc}
    if tc is not None and "url" in config["temp_context"]:
        config["temp_context"]["url"] = sub(config["temp_context"]["url"])
    return config, steps


def main():
    p = argparse.ArgumentParser(
        description="Run an ordered type/click flow against debug Chrome over CDP",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Steps run in the order given on the command line.")
    tab = p.add_mutually_exclusive_group()
    tab.add_argument("--tab", help="use the open tab whose URL/title contains this")
    tab.add_argument("--new-tab", metavar="URL", help="open a new tab at URL and use it")
    tab.add_argument("--temp-context", metavar="URL",
                     help="spawn an isolated temp browser context, open URL in it, "
                          "and run the flow there (see docs/TEMP_CHROME_CONTEXTS.md)")
    p.add_argument("--label", default="cdp-flow",
                   help="label for the temp context in temp-profile list (default cdp-flow)")
    p.add_argument("--dispose", action="store_true",
                   help="tear down the temp context after a successful flow")
    p.add_argument("--flow", help="JSON flow file (runs before any CLI steps)")
    p.add_argument("--var", action="append", default=[], metavar="NAME=VALUE",
                   help="substitute {NAME} placeholders in the flow file")
    p.add_argument("--timeout", type=float,
                   help="per-waitfor/goto timeout in seconds (default 15)")
    p.add_argument("--delay", type=float,
                   help="pause between steps in seconds (default 0.3)")
    p.add_argument("--type", dest="type", nargs=2, metavar=("XPATH", "TEXT"),
                   action=OrderedStep, help="focus XPATH and type TEXT (replaces content)")
    p.add_argument("--click", dest="click", metavar="XPATH",
                   action=OrderedStep, help="click the element at XPATH")
    p.add_argument("--press", dest="press", metavar="KEY",
                   action=OrderedStep, help="press a key (Enter, Tab, Escape, ...)")
    p.add_argument("--wait", dest="wait", metavar="SECONDS",
                   action=OrderedStep, help="sleep for SECONDS")
    p.add_argument("--waitfor", dest="waitfor", metavar="XPATH",
                   action=OrderedStep, help="wait until XPATH exists and is visible")
    p.add_argument("--goto", dest="goto", metavar="URL",
                   action=OrderedStep, help="navigate the tab to URL")
    p.add_argument("--emit", metavar="FILE",
                   help="write {tab_id, context_id} JSON here after the tab is "
                        "attached, so a later flow can reuse the same window "
                        "(--tab <tab_id>)")
    args = p.parse_args()

    variables = {}
    for v in args.var:
        name, _, value = v.partition("=")
        variables[name] = value

    config, steps = {}, []
    if args.flow:
        config, steps = load_flow(args.flow, variables)
    steps += getattr(args, "steps", None) or []
    if not steps:
        p.error("no steps given (use --type/--click/... or --flow)")

    # CLI flags override the flow file's config.
    temp = None
    if args.temp_context:
        temp = {"url": args.temp_context}
    elif not (args.tab or args.new_tab) and config.get("temp_context"):
        temp = dict(config["temp_context"])
    if temp is not None:
        if not temp.get("url"):
            p.error("temp_context needs a \"url\" to open")
        temp.setdefault("label", args.label)
        if args.dispose:
            temp["dispose"] = True
    tab = None if temp else (args.tab or (None if args.new_tab else config.get("tab")))
    new_tab = None if (temp or tab) else (args.new_tab or config.get("new_tab"))
    timeout = args.timeout if args.timeout is not None else config.get("timeout", 15)
    delay = args.delay if args.delay is not None else config.get("delay", 0.3)
    if not tab and not new_tab and not temp:
        p.error("pick a tab: --tab/--new-tab/--temp-context, "
                "or \"tab\"/\"new_tab\"/\"temp_context\" in the flow file")

    try:
        cdp = CDP()
    except Exception as e:
        print(f"cannot reach debug Chrome at {CDP_HTTP}: {e}", file=sys.stderr)
        sys.exit(1)

    spawned_ctx = None  # a temp context THIS run created (and may dispose)
    ctx_id = None       # context of the tab in use, for --emit (spawned or attached)
    try:
        if temp:
            spawned_ctx = cdp.call("Target.createBrowserContext")["browserContextId"]
            ctx_id = spawned_ctx
            track_context(spawned_ctx, temp["label"])
            print(f"temp context {spawned_ctx[:8]} ({temp['label']}) -> {temp['url']}")
            target_id = cdp.call("Target.createTarget", {
                "url": temp["url"], "browserContextId": spawned_ctx, "newWindow": True,
            })["targetId"]
        elif new_tab:
            target_id = cdp.call("Target.createTarget", {"url": new_tab})["targetId"]
        else:
            target_id = pick_tab(tab)
            print(f"using tab {target_id[:8]} (matched '{tab}')")
        cdp.call("Target.activateTarget", {"targetId": target_id})
        page = Page(cdp, target_id)
        if temp or new_tab:
            step_goto(page, temp["url"] if temp else new_tab, timeout)  # initial load

        if args.emit:
            if ctx_id is None:  # attached to an existing tab: look up its context
                for t in cdp.call("Target.getTargets").get("targetInfos", []):
                    if t.get("targetId") == target_id:
                        ctx_id = t.get("browserContextId")
                        break
            with open(args.emit, "w", encoding="utf-8") as f:
                json.dump({"tab_id": target_id, "context_id": ctx_id}, f)
            print(f"emitted tab/context ids to {args.emit}")

        for i, (kind, val) in enumerate(steps, 1):
            if kind == "type":
                shown = "***" if (len(val) > 2 and val[2]) else val[1]
                label = f"{val[0]} -> {shown}"
            else:
                label = val if isinstance(val, str) else " -> ".join(val)
            print(f"[{i}/{len(steps)}] {kind}: {label}")
            if kind == "type":
                step_type(page, val[0], val[1])
            elif kind == "click":
                note = step_click(page, val)
                print(f"          {note}")
            elif kind == "press":
                step_press(page, val)
            elif kind == "wait":
                time.sleep(float(val))
            elif kind == "waitfor":
                step_waitfor(page, val, timeout)
            elif kind == "goto":
                step_goto(page, val, timeout)
            time.sleep(delay)
        print("flow complete")
        if spawned_ctx and temp.get("dispose"):
            dispose_context(cdp, spawned_ctx)
            print(f"disposed temp context {spawned_ctx[:8]}")
        elif spawned_ctx:
            print(f"temp context {spawned_ctx[:8]} left open (24h TTL sweep applies; "
                  f"dispose now: bash scripts/temp-profile/run.sh dispose {spawned_ctx[:8]})")
    except (RuntimeError, TimeoutError) as e:
        print(f"FAILED: {e}", file=sys.stderr)
        if spawned_ctx:
            print(f"temp context {spawned_ctx[:8]} kept for inspection "
                  f"(dispose: bash scripts/temp-profile/run.sh dispose {spawned_ctx[:8]})",
                  file=sys.stderr)
        sys.exit(1)
    finally:
        cdp.close()


if __name__ == "__main__":
    main()