#!/usr/bin/env python3
"""Muse.ai chat API via CDP (headless smoke profile).
Usage:
  muse-chat-api.py send "hello"          # send message to current chat
  muse-chat-api.py messages             # print recent messages
  muse-chat-api.py wait [timeout]       # wait for new response\n  muse-chat-api.py navigate <url>       # go to thread URL
Requires: websocket-client (pip install --break-system-packages websocket-client)
CDP relay must be up: http://10.201.87.2:9410/json/list
"""
import json, sys, time, urllib.request, subprocess, os
import websocket

CDP_URL = "http://10.201.87.2:9410/json/list"

def get_page():
    with urllib.request.urlopen(CDP_URL, timeout=5) as r:
        targets = json.load(r)
    pages = [t for t in targets if t.get('type')=='page' and 'muse.ai' in t.get('url','')]
    if not pages:
        raise RuntimeError("no muse.ai page found")
    return pages[0]

def revive_browser():
    script = os.path.expanduser("~/Projects/NetVM/bin/netvm-chrome.sh")
    print("CDP disconnect detected. Reviving Chromium smoke profile...", file=sys.stderr)
    subprocess.Popen([script, "--headless", "smoke", "https://muse.ai"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(4)

def connect(retries=2):
    for attempt in range(retries + 1):
        try:
            page = get_page()
            ws = websocket.create_connection(page['webSocketDebuggerUrl'], timeout=10)
            return ws
        except Exception as e:
            if attempt < retries:
                revive_browser()
            else:
                raise RuntimeError(f"Failed to connect to CDP after {retries} retries: {e}")


def ev(ws, expr, await_promise=False):
    ws.send(json.dumps({"id":1,"method":"Runtime.evaluate",
        "params":{"expression":expr,"returnByValue":True,"awaitPromise":await_promise}}))
    r = json.loads(ws.recv())
    return r.get('result',{}).get('result',{}).get('value')

def send_message(text):
    ws = connect()
    try:
        result = ev(ws, f"""
(async () => {{
  const ta = document.querySelector('textarea[aria-label="Message"]');
  if (!ta) return 'ERROR: no composer';
  ta.focus();
  document.execCommand('insertText', false, {json.dumps(text)});
  await new Promise(r => setTimeout(r, 200));
  ta.dispatchEvent(new KeyboardEvent('keydown', {{key:'Enter', code:'Enter', keyCode:13, bubbles:true}}));
  return 'sent';
}})()
""", await_promise=True)
        return result
    finally:
        ws.close()

def get_messages(limit=10):
    ws = connect()
    try:
        return ev(ws, f"""
Array.from(document.querySelectorAll('p')).slice(-{limit}).map(e=>e.innerText).join('\\n---\\n')
""")
    finally:
        ws.close()

def navigate(url):
    """Navigate to a URL (e.g., a thread: https://muse.ai/thread/<id>)."""
    ws = connect()
    try:
        ws.send(json.dumps({"id":1,"method":"Page.navigate","params":{"url":url}}))
        r = json.loads(ws.recv())
        return r.get('result',{}).get('frameId','navigated')
    finally:
        ws.close()

def wait_for_response(timeout=60, poll=2):
    """Wait for a new assistant message after sending."""
    ws = connect()
    try:
        before = ev(ws, "document.body.innerText.length")
        start = time.time()
        while time.time() - start < timeout:
            time.sleep(poll)
            # Check if there's a "stop" button (indicates generating)
            generating = ev(ws, """
!!document.querySelector('button[aria-label*="Stop"], button[aria-label*="stop"]')
""")
            if not generating:
                # Give it a moment to finish rendering
                time.sleep(2)
                return get_messages(5)
        return "TIMEOUT waiting for response"
    finally:
        ws.close()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == "send" and len(sys.argv) > 2:
        print(send_message(sys.argv[2]))
    elif cmd == "messages":
        print(get_messages(int(sys.argv[2]) if len(sys.argv)>2 else 10))
    elif cmd == "navigate" and len(sys.argv) > 2:
        print(navigate(sys.argv[2]))
    elif cmd == "wait":
        print(wait_for_response(int(sys.argv[2]) if len(sys.argv)>2 else 60))
    else:
        print("unknown command")