#!/usr/bin/env python3
"""
Multi-account muse.ai chat API.

Usage:
  muse-chat-api.py --account <name> send "message"
  muse-chat-api.py --account <name> messages [n]
  muse-chat-api.py --account <name> wait [timeout]

Accounts (configured on bl):
  muse-agent : CDP 9410 (smoke node) - ltd.pixels.ltd@gmail.com
  pip        : CDP 9420 (phone-test node) - piparada (phone OTP)
  646        : (not currently logged in - needs OTP)

The account maps to a CDP endpoint. The browser must be running
with the account logged in.
"""
import json, urllib.request, websocket, time, sys, argparse

ACCOUNTS = {
    "muse-agent": "http://127.0.0.1:9410/json/list",
    "pip": "http://127.0.0.1:9420/json/list",
    # "646": "http://127.0.0.1:9430/json/list",  # TODO: set up 646 node
}

def get_page(cdp_url):
    with urllib.request.urlopen(cdp_url, timeout=5) as r:
        ts = json.load(r)
    pages = [t for t in ts if t.get('type') == 'page']
    if not pages:
        print("ERROR: No page found", file=sys.stderr)
        sys.exit(1)
    return pages[0]

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

def cmd_send(ws, message):
    # Escape for JS
    msg_esc = message.replace('\\', '\\\\').replace('`', '\\`').replace('$', '\\$')
    result = ev(ws, f"""(async()=>{{
      const input = document.querySelector('[contenteditable="true"]') ||
                    document.querySelector('textarea[placeholder*="Message"]') ||
                    [...document.querySelectorAll('div[role="textbox"]')][0];
      if (!input) return 'NOINPUT';
      input.focus();
      document.execCommand('insertText', false, `{msg_esc}`);
      await new Promise(r=>setTimeout(r,500));
      const send = [...document.querySelectorAll('button')].find(b=>
        b.getAttribute('aria-label')&&b.getAttribute('aria-label').toLowerCase().includes('send')
      );
      if (send) {{ send.click(); return 'sent'; }}
      const ke = new KeyboardEvent('keydown', {{key:'Enter', code:'Enter', bubbles:true}});
      input.dispatchEvent(ke);
      return 'enter-sent';
    }})()""", True)
    print(result)

def cmd_messages(ws, n=5):
    result = ev(ws, f"""(() => {{
      const ps = [...document.querySelectorAll('p')].slice(-{n*2}).map(p=>p.innerText.slice(0,200));
      return ps.join('\\n---\\n');
    }})()""")
    print(result)

def cmd_wait(ws, timeout=30):
    print(f"Waiting {timeout}s for response...")
    # Simple: wait and then get messages
    time.sleep(timeout)
    cmd_messages(ws, 2)

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--account', required=True, choices=list(ACCOUNTS.keys()),
                   help='Account to use')
    p.add_argument('command', choices=['send', 'messages', 'wait'])
    p.add_argument('arg', nargs='?', default=None)
    args = p.parse_args()

    cdp_url = ACCOUNTS[args.account]
    page = get_page(cdp_url)
    ws = websocket.create_connection(page['webSocketDebuggerUrl'], timeout=15)

    try:
        if args.command == 'send':
            if not args.arg:
                print("ERROR: send requires a message", file=sys.stderr)
                sys.exit(1)
            cmd_send(ws, args.arg)
        elif args.command == 'messages':
            n = int(args.arg) if args.arg else 5
            cmd_messages(ws, n)
        elif args.command == 'wait':
            t = int(args.arg) if args.arg else 30
            cmd_wait(ws, t)
    finally:
        ws.close()

if __name__ == '__main__':
    main()