Grok Bot

"""Chat verbs: grok_chat_status / grok_chat_read via UIA on the Grok Bot hwnd.

Never steal focus. Never return webhook keys. UIA first, not OCR. Never CDP 9222.
Uses grok_hwnd / _cb / _callback_ok from handlers.windows after the hwnd unwrap fix.
"""

from __future__ import annotations

import re
import time

from grok_detect import HOST_APP, detect_host
from handlers.windows import (
    RESTART_HINT,
    _cb,
    _callback_ok,
    _title_of,
    flatten_callback,
    grok_hwnd,
    host_launch_flags_ok,
    main_window_state,
    read_target_hwnd,
)

# Proven UIA on Grok Bot.exe hwnd (AdomLapper): Document named "Grok Bot",
# Edit "Prompt", heading / sand-conversation-heading ("New Bot"),
# Text/button "<agent> is working" while thinking.
DOCUMENT_NAME = "Grok Bot"
PROMPT_NAME = "Prompt"
HEADING_AID = "sand-conversation-heading"
VIEW_DETAILS = "View conversation details"

SIDEBAR_NOISE = {
    "bot list",
    "new bot or channel",
    "search",
    "prompt",
    "view conversation details",
    "conversation",
    "grok bot",
}

_WORKING_SUFFIXES = (" is working", " is thinking")

# Title-bar caption buttons. With no thread list on the page (blank window after a
# background launch, or the reconnect page) a Heading query can land on these, and
# focusedThread once read "Minimize" (wiki #5). Never report one as a thread.
CAPTION_BUTTONS = {
    "minimize", "maximize", "restore", "restore down", "close", "system",
    "system menu", "application",
}

# The page Grok Bot shows when the agent box (Grok Bot's computer) is down:
# "Reconnecting to your computer... Couldn't Reach Grok Bot's Computer ...
# Recover Grok Bot's Computer / Retry". Matched case-insensitively, both apostrophes.
_UNREACHABLE_MARKERS = (
    "couldn't reach", "couldn\u2019t reach",
    "recover grok bot's computer", "recover grok bot\u2019s computer",
)
_RECONNECT_MARKERS = ("reconnecting to your computer",)

# Per-callback UIA budget. A blank window used to run grok_chat_read into ab's 60 s
# verb timeout; every chat callback now gives up after this, and the Document read
# has an overall budget too.
UI_TIMEOUT = 8
DOC_READ_BUDGET = 15.0


def _ui(command: str, args: dict) -> dict:
    return _cb(command, args, timeout=UI_TIMEOUT)

# Never return webhook keys / bearer tokens in transcript.
_SECRET_RE = re.compile(
    r"(?:crsr_[A-Za-z0-9_-]+)"
    r"|(?:Bearer\s+[A-Za-z0-9._\-+=/]+)"
    r"|(?:Webhook\s+key\s*[:=]\s*\S+)",
    re.IGNORECASE,
)


def _as_controls(blob) -> list[dict]:
    merged = flatten_callback(blob)
    if isinstance(merged, list):
        return [c for c in merged if isinstance(c, dict)]
    if not isinstance(merged, dict):
        return []
    out: list[dict] = []
    for key in ("controls", "matches", "items", "result", "nodes", "windows"):
        val = merged.get(key)
        if isinstance(val, list):
            out.extend(c for c in val if isinstance(c, dict))
    data = merged.get("data")
    if isinstance(data, list):
        out.extend(c for c in data if isinstance(c, dict))
    best = merged.get("best")
    if isinstance(best, dict):
        out.append(best)
    if not out and (merged.get("name") or merged.get("hwnd") or merged.get("role")):
        out.append(merged)
    return out


def _ctrl_name(c: dict) -> str:
    return str(c.get("name") or c.get("Name") or c.get("title") or "")


def _ctrl_value(c: dict) -> str:
    for k in ("value", "Value", "currentValue", "text", "content"):
        v = c.get(k)
        if v is not None and str(v).strip():
            return str(v).strip()
    return ""


def _extract_text(blob) -> str:
    merged = flatten_callback(blob)
    if isinstance(merged, str):
        return merged
    if isinstance(merged, list):
        parts = []
        for item in merged:
            if isinstance(item, str) and item.strip():
                parts.append(item)
            elif isinstance(item, dict):
                t = _extract_text(item)
                if t:
                    parts.append(t)
        return "\n".join(parts)
    if not isinstance(merged, dict):
        return ""
    for k in ("text", "value", "content", "all", "visible"):
        v = merged.get(k)
        if isinstance(v, str) and v.strip():
            return v
    data = merged.get("data")
    if isinstance(data, str) and data.strip():
        return data
    if isinstance(data, dict):
        for k in ("text", "value", "content"):
            v = data.get(k)
            if isinstance(v, str) and v.strip():
                return v
    best = merged.get("best")
    if isinstance(best, dict):
        t = _extract_text(best)
        if t:
            return t
    return ""


def _looks_working(name: str) -> bool:
    nl = (name or "").strip().lower()
    if not nl:
        return False
    if nl.endswith(" is working") or nl.endswith(" is thinking"):
        return True
    if nl in ("working", "thinking"):
        return True
    return False


def _strip_working(name: str) -> str:
    n = (name or "").strip()
    nl = n.lower()
    for suf in _WORKING_SUFFIXES:
        if nl.endswith(suf):
            return n[: len(n) - len(suf)].strip()
    return ""


def _redact_secrets(text: str) -> str:
    if not text:
        return text
    return _SECRET_RE.sub("[redacted]", text)


def _window_missing(verb: str) -> dict:
    info = detect_host()
    thinking = False
    ms = main_window_state()
    if ms["mainWindowHidden"]:
        out = {"success": False, "errorCode": "main_window_hidden", "thinking": False, "idle": False,
               "agentName": None, "hwnd": None, "title": None, "hostApp": HOST_APP,
               "windowEtiquette": "background",
               "_hint": "Grok Bot's main window is hidden (only small secondary windows are visible). "
                        "Its hwnd is not remembered on this PC, so it cannot be read in place. "
                        "grok_open {\"restart\":true,\"reason\":\"...\"} relaunches it in the background.",
               "_next": "grok_open {\"restart\":true,\"reason\":\"...\"}"}
        out.update(ms)
        return out
    return {
        "success": False,
        "errorCode": "window_not_found",
        "error": "No Grok Bot window.",
        "thinking": thinking,
        "idle": False,
        "agentName": None,
        "hwnd": None,
        "title": None,
        "hostApp": HOST_APP,
        "hostAppFallback": info.get("hostAppFallback"),
        "windowEtiquette": "background",
        "_hint": (
            "No Grok Bot hwnd (desktop_find_window titleContains exact \"Grok Bot\"). "
            f"Call grok_open {{}} then grok_{verb} {{}}. Never steal focus. Never kill Cursor. "
            "Do not treat Google Chrome / ai-grok-bot as the host."
        ),
        "_next": f"grok_open {{}} then grok_{verb} {{}}",
    }


def _find_working(hwnd: int) -> tuple[bool, str]:
    """Return (thinking, working_control_name). Read-only; never focuses."""
    queries = (
        {"hwnd": hwnd, "contains": "is working"},
        {"hwnd": hwnd, "nameContains": "is working"},
        {"hwnd": hwnd, "contains": "working"},
    )
    for args in queries:
        blob = _ui("desktop_find_controls", args)
        ctrls = _as_controls(blob)
        if not ctrls:
            ctrls = _as_controls(_ui("desktop_find_control", args))
        for c in ctrls:
            n = _ctrl_name(c) or _ctrl_value(c)
            if _looks_working(n):
                return True, n
    return False, ""


def _is_thread_name(n: str) -> bool:
    nl = (n or "").strip().lower()
    return bool(nl) and nl not in SIDEBAR_NOISE and nl not in CAPTION_BUTTONS


def _read_agent_name(hwnd: int, working_name: str = "") -> str | None:
    """Focused thread name, or None when the page shows no thread (never a caption button)."""
    stripped = _strip_working(working_name)
    if stripped:
        return stripped
    queries = (
        {"hwnd": hwnd, "automationId": HEADING_AID},
        {"hwnd": hwnd, "name": HEADING_AID},
        {"hwnd": hwnd, "role": "Heading"},
        {"hwnd": hwnd, "name": VIEW_DETAILS},
    )
    for args in queries:
        blob = _ui("desktop_find_control", args)
        for c in _as_controls(blob):
            n = (_ctrl_name(c) or _ctrl_value(c)).strip()
            if _is_thread_name(n):
                return n
        blob2 = _ui("desktop_find_controls", args)
        for c in _as_controls(blob2):
            n = (_ctrl_name(c) or _ctrl_value(c)).strip()
            if _is_thread_name(n):
                return n
    return None


def _prompt_value(hwnd: int) -> str:
    blob = _ui("desktop_ui_get", {"hwnd": hwnd, "name": PROMPT_NAME, "role": "Edit"})
    text = _extract_text(blob)
    if text:
        return text.strip()
    blob = _ui("desktop_find_control", {"hwnd": hwnd, "name": PROMPT_NAME, "role": "Edit"})
    for c in _as_controls(blob):
        v = _ctrl_value(c)
        if v:
            return v
    return ""


def _strip_prompt_garbage(text: str, prompt_val: str) -> str:
    """Drop the Prompt compose box if it is tacked onto the Document transcript."""
    out = text or ""
    pv = (prompt_val or "").strip()
    if pv and out.rstrip().endswith(pv):
        out = out.rstrip()[: -len(pv)].rstrip()
    # Trailing isolated "Prompt" label (compose chrome, not a message).
    lines = out.splitlines()
    while lines and lines[-1].strip().lower() in ("prompt", ""):
        if lines[-1].strip() == "":
            lines.pop()
            continue
        if lines[-1].strip().lower() == "prompt":
            lines.pop()
            continue
        break
    return "\n".join(lines).rstrip()


def _read_document_text(hwnd: int) -> tuple[str, dict | None]:
    """desktop_ui_text {hwnd, role:Document, range:all} — proven Document 'Grok Bot'."""
    last = None
    deadline = time.monotonic() + DOC_READ_BUDGET
    for args in (
        {"hwnd": hwnd, "role": "Document", "range": "all"},
        {"hwnd": hwnd, "name": DOCUMENT_NAME, "role": "Document", "range": "all"},
        {"hwnd": hwnd, "name": DOCUMENT_NAME, "range": "all"},
    ):
        if time.monotonic() > deadline:
            break
        last = _ui("desktop_ui_text", args)
        if not _callback_ok(last):
            continue
        text = _extract_text(last)
        if text.strip():
            return text, last
    return "", last


def computer_state_from_text(text: str) -> tuple[str, str | None]:
    """(state, bannerLine) from the page text. unreachable | reconnecting | no_banner | unknown.

    no_banner is NOT proof the box is up: the banner only shows on some pages.
    """
    if not (text or "").strip():
        return "unknown", None
    for line in text.splitlines():
        ll = line.strip().lower()
        if any(m in ll for m in _UNREACHABLE_MARKERS):
            return "unreachable", line.strip()[:200]
    for line in text.splitlines():
        ll = line.strip().lower()
        if any(m in ll for m in _RECONNECT_MARKERS):
            return "reconnecting", line.strip()[:200]
    return "no_banner", None


def computer_fields(state: str, banner: str | None) -> dict:
    return {
        "computerState": state,
        "computerReachable": False if state in ("unreachable", "reconnecting") else None,
        "computerBanner": banner,
    }


def page_probe(hwnd: int) -> dict:
    """One bounded Document read: is there a page yet, and is Grok Bot's computer down?"""
    raw, last = _read_document_text(hwnd)
    state, banner = computer_state_from_text(raw)
    return {"hasDocument": bool((raw or "").strip()), "text": raw or "", "last": last,
            "state": state, "banner": banner}


_PROBE_CACHE: dict = {}
PROBE_TTL = 30.0


def probe_computer(max_age: float = PROBE_TTL) -> dict:
    """computer_* fields for readiness / webhook verbs. Read-only, bounded, never raises.

    Cached for 30 s: Hydrogen polls grok_readiness thousands of times a day (16k on crr), and
    each uncached probe is several UIA callbacks on the user's PC."""
    hit = _PROBE_CACHE.get("v")
    if hit and time.monotonic() - _PROBE_CACHE.get("t", 0) < max_age:
        return dict(hit, computerStateAgeSeconds=int(time.monotonic() - _PROBE_CACHE["t"]))
    out = _probe_computer_uncached()
    _PROBE_CACHE.update(v=out, t=time.monotonic())
    return dict(out)


def _probe_computer_uncached() -> dict:
    try:
        hwnd, _win = grok_hwnd(allow_cursor=False)
        if not hwnd:
            out = computer_fields("unknown", None)
            out["computerStateReason"] = "no Grok Bot window to read"
            return out
        pr = page_probe(hwnd)
        out = computer_fields(pr["state"], pr["banner"])
        if not pr["hasDocument"]:
            out["computerStateReason"] = "Grok Bot window has no page yet (blank or loading)"
        return out
    except Exception as e:  # never let a probe break the verb it decorates
        out = computer_fields("unknown", None)
        out["computerStateReason"] = f"probe failed: {e.__class__.__name__}"
        return out


UNREACHABLE_HINT = (
    "Grok Bot shows \"Couldn't Reach Grok Bot's Computer\": the agent box is down, so webhook "
    "events are accepted (HTTP 200) but no reply can come back. Ask the user to press "
    "Recover Grok Bot's Computer / Retry in Grok Bot, then re-check grok_chat_status."
)
NO_PAGE_REASON = "no page exposed to UIA (window blank or still loading)"


def _no_page(hwnd: int, win, probe: dict, verb: str) -> dict:
    """Fast answer when the window has no Document, or shows the unreachable page."""
    state = probe["state"]
    down = state in ("unreachable", "reconnecting")
    flags_ok = None if down else host_launch_flags_ok()
    out = {
        "success": False,
        "errorCode": "computer_unreachable" if down else "no_transcript",
        "thinking": False,
        "idle": True,
        "agentName": None,
        "agent": None,
        "focusedThread": None,
        "focusedThreadReason": ("Grok Bot is showing the reconnect page, not a thread" if down
                                else NO_PAGE_REASON),
        "text": "",
        "length": 0,
        "visibleText": _redact_secrets(probe["text"])[:1000],
        "hwnd": hwnd,
        "title": _title_of(win),
        "windowEtiquette": "background",
        "_hint": (UNREACHABLE_HINT if down else
                  "The Grok Bot window shows no page to UIA (only caption buttons). It may still be "
                  "loading: retry in 30 s, or grok_screenshot to see it. Answered fast instead of "
                  "waiting out the verb timeout."),
        "_next": ("grok_chat_status {} after the user presses Retry" if down
                  else f"wait 30 s, then grok_{verb} {{}}"),
    }
    out.update(computer_fields(state, probe["banner"]))
    out["launchedWithAccessibility"] = flags_ok
    if flags_ok is False:
        out["focusedThreadReason"] = "Grok Bot's page is not exposed to UIA (started without the accessibility switches)"
        out["_hint"] = RESTART_HINT
        out["_next"] = "grok_open {\"restart\":true,\"reason\":\"Grok Bot page invisible to UIA\"}"
    return out


def _select_sidebar_thread(hwnd: int, name: str) -> dict:
    """Best-effort click a sidebar row by UIA name. Does not steal OS focus beyond Invoke."""
    last = None
    for args in (
        {"hwnd": hwnd, "name": name},
        {"hwnd": hwnd, "name": name, "role": "ListItem"},
        {"hwnd": hwnd, "name": name, "role": "Button"},
        {"hwnd": hwnd, "nameContains": name},
    ):
        last = _ui("desktop_ui_click", args)
        if _callback_ok(last) and not (isinstance(last, dict) and last.get("error")):
            return last
    return last if isinstance(last, dict) else {"success": False}


def handle_chat_status(_args: dict) -> dict:

    """thinking=true when UIA exposes '<agent> is working' / working chip. Never steals focus."""
    hwnd, win, hidden = read_target_hwnd(allow_cursor=False)
    if not hwnd:
        hwnd, win, hidden = read_target_hwnd(allow_cursor=True)
    if not hwnd:
        return _window_missing("chat_status")

    probe = page_probe(hwnd)
    if not probe["hasDocument"] or probe["state"] in ("unreachable", "reconnecting"):
        out = _no_page(hwnd, win, probe, "chat_status")
        # status is an observation, not a failed read: the window exists
        out["success"] = True
        out.pop("text", None); out.pop("length", None)
        return out

    thinking, working_name = _find_working(hwnd)
    agent = _read_agent_name(hwnd, working_name)
    title = _title_of(win) or HOST_APP
    out = {
        "success": True,
        "thinking": bool(thinking),
        "idle": (not thinking),
        "agentName": agent,
        "agent": agent,
        "focusedThread": agent,
        "hwnd": hwnd,
        "title": title,
        "windowEtiquette": "background",
        "_hint": (
            f"Grok Bot hwnd={hwnd} focusedThread={agent!r} "
            + ("thinking (UIA name like '<agent> is working')." if thinking else "idle (no working chip).")
            + " Read-only UIA — focus not stolen. grok_chat_read for the visible transcript."
        ),
    }
    if agent is None:
        out["focusedThreadReason"] = ("main window is hidden: the heading and the working chip are not "
                                      "readable until it is shown" if hidden else "no conversation heading on the page")
    out["mainWindowHidden"] = hidden
    if hidden:
        out["thinking"] = None; out["idle"] = None
    out.update(computer_fields(probe["state"], probe["banner"]))
    return out


def handle_chat_read(args: dict) -> dict:
    """Visible transcript via desktop_ui_text Document range=all.

    limit is CHARACTERS from the END (not messages). Values under 500 are coerced
    to max(limit*250, 4000). Pass 4000+ to search for a nonce.

    Optional args.thread / args.agent: try to click that sidebar row then re-read.
    Always returns focusedThread (heading / agent name). After cold launch the focused
    thread may not be grok-bridge — callers should pass thread/agent or wait after grok_open.
    """
    hwnd, win, hidden = read_target_hwnd(allow_cursor=False)
    if not hwnd:
        hwnd, win, hidden = read_target_hwnd(allow_cursor=True)
    if not hwnd:
        miss = _window_missing("chat_read")
        miss["text"] = ""
        miss["length"] = 0
        miss["focusedThread"] = None
        miss["agent"] = None
        return miss

    probe = page_probe(hwnd)
    if not probe["hasDocument"] or probe["state"] in ("unreachable", "reconnecting"):
        return _no_page(hwnd, win, probe, "chat_read")

    wanted = str(args.get("thread") or args.get("agent") or "").strip()
    select_attempt = None
    if wanted and not hidden:
        select_attempt = _select_sidebar_thread(hwnd, wanted)

    thinking, working_name = _find_working(hwnd)
    agent = _read_agent_name(hwnd, working_name)
    if select_attempt is None:
        raw, last = probe["text"], probe["last"]    # nothing clicked: the probe read is current
    else:
        raw, last = _read_document_text(hwnd)

    # Cold start: focused thread is often not grok-bridge. If nothing requested and
    # transcript is empty/tiny, try the grok-bridge sidebar row once.
    tiny = len((raw or "").strip()) < 40
    agent_l = (agent or "").strip().lower()
    if not hidden and not wanted and tiny and "grok-bridge" not in agent_l and "grok bridge" not in agent_l:
        select_attempt = _select_sidebar_thread(hwnd, "grok-bridge")
        thinking, working_name = _find_working(hwnd)
        agent = _read_agent_name(hwnd, working_name)
        raw, last = _read_document_text(hwnd)
        tiny = len((raw or "").strip()) < 40
    prompt_val = _prompt_value(hwnd)
    text = _strip_prompt_garbage(raw, prompt_val)
    text = _redact_secrets(text)

    # limit is CHARACTERS from the END, not messages. Callers who pass 20/40
    # almost certainly meant "last N messages" — coerce so a nonce search
    # is not truncated to the last 20 chars ("replied ").
    LIMIT_HINT = (
        "`limit` is characters from the end, not messages. "
        "Values under 500 are coerced. Pass 4000+ to search for a nonce."
    )
    limit = args.get("limit")
    try:
        limit_n = int(limit) if limit is not None and str(limit).strip() != "" else None
    except (TypeError, ValueError):
        limit_n = None
    if limit_n is not None and limit_n < 0:
        limit_n = None
    limit_meta = {}
    if limit_n:
        limit_meta["limitRequested"] = limit_n
        if limit_n < 500:
            chars = max(limit_n * 250, 4000)
            limit_meta["limitCoerced"] = True
        else:
            chars = limit_n
        limit_meta["limitChars"] = chars
        text = text[-chars:]

    ok = bool(text) or _callback_ok(last)
    focused = agent
    if wanted and not ok:
        err = "thread_select_failed"
    elif not text.strip():
        err = "focused_thread_empty"
    else:
        err = None

    if not ok and not text:
        out = {
            "success": False,
            "errorCode": "transcript_unavailable",
            "agentName": agent,
            "agent": agent,
            "focusedThread": focused,
            "thinking": bool(thinking),
            "idle": (not thinking),
            "text": "",
            "length": 0,
            "hwnd": hwnd,
            "title": _title_of(win),
            "windowEtiquette": "background",
            "_hint": (
                "desktop_ui_text {hwnd, role:Document, range:all} returned no text. "
                "After cold launch the focused thread may not be grok-bridge. "
                "Pass args.thread / args.agent to click that sidebar row, or grok_open then wait. "
                "Never OCR-first. Never CDP 9222. "
                + LIMIT_HINT
                + " Default 4000+ when polling for a nonce."
            ),
            "_next": "grok_open {} then wait, or grok_chat_read {\"thread\":\"grok-bridge\"}",
        }
        if select_attempt is not None:
            out["selectAttempt"] = select_attempt
        out.update(limit_meta)
        return out

    out = {
        "success": True if text.strip() else False,
        "agentName": agent,
        "agent": agent,
        "focusedThread": focused,
        "thinking": bool(thinking),
        "idle": (not thinking),
        "text": text,
        "length": len(text),
        "hwnd": hwnd,
        "title": _title_of(win),
        "windowEtiquette": "background",
        "_hint": (
            f"Read UIA Document range=all on Grok Bot hwnd={hwnd} "
            f"({len(text)} chars, focusedThread={focused!r}, "
            + ("thinking" if thinking else "idle")
            + "). Latest messages are at the end. Webhook keys redacted. "
            "Optional thread/agent clicks a sidebar row then re-reads. "
            "If empty after cold start, grok_open then wait or pass thread:grok-bridge. "
            + LIMIT_HINT
            + (" Default 4000+ when polling for a nonce." if not limit_n else "")
        ),
    }
    out.update(limit_meta)
    out.update(computer_fields(*computer_state_from_text(raw)))
    out["mainWindowHidden"] = hidden
    if focused is None:
        out["focusedThreadReason"] = ("main window is hidden; read in place (thread switching needs it shown)"
                                      if hidden else "no conversation heading on the page")
    if err and not text.strip():
        out["success"] = False
        out["errorCode"] = err
        out["_next"] = "grok_open {} then wait, or grok_chat_read {\"thread\":\"grok-bridge\"}"
    if wanted:
        out["requestedThread"] = wanted
    if select_attempt is not None:
        out["selectAttempt"] = select_attempt
    return out