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

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

The <agent> name matches the NetVM node, chrome-box profile, and
the `agent` column in ACCOUNTS.md. No translation.

Accounts:
  muse : CDP 9410 - ltd.pixels.ltd@gmail.com (email_otp) - ACTIVE
  pip  : CDP 9420 - piparada (phone_otp) - PENDING (needs re-auth)
  646  : (not set up yet)

See ACCOUNTS.md for the full registry.
"""
import json, urllib.request, websocket, time, sys, argparse

# Maps agent name -> (node, CDP URL). Agent == node == profile.
ACCOUNTS = {
    "muse": ("muse", "http://127.0.0.1:9410/json/list"),
    "pip": ("pip", "http://127.0.0.1:9420/json/list"),
}

def get_page(node, cdp_url):
    # Run from within the node's netns via netvm-exec
    # For now, assume we're already in the netns (called via netvm-exec.sh)
    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):
    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...")
    time.sleep(timeout)
    cmd_messages(ws, 2)

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--account', required=True, choices=list(ACCOUNTS.keys()),
                   help='Agent name (matches node, profile, ACCOUNTS.md)')
    p.add_argument('command', choices=['send', 'messages', 'wait'])
    p.add_argument('arg', nargs='?', default=None)
    args = p.parse_args()

    node, cdp_url = ACCOUNTS[args.account]
    # Note: Must be run via: netvm-exec.sh <node> -- python3 muse-chat-api.py ...
    # The node name matches the account name.
    page = get_page(node, 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()