#!/usr/bin/env python3
# GOLDEN PATH: container -> VM (34.139.37.135) -> bl (100.123.153.75) -> netns -> browser -> agent
# This API is the bridge. Every call traverses 4 hops. Respect the path.

"""
Multi-account muse.ai chat API with approval handling.

Approvals: The browser may show permission dialogs (e.g., "Allow pip to share
information with 34.139.37.135?"). The API detects these and handles them:
- Known-safe (our infrastructure IPs): auto-approve
- Unknown: raise APPROVAL_NEEDED, operator decides via chat

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]
  muse-chat-api.py --account <agent> approvals  # check pending approvals
  muse-chat-api.py --account <agent> upload <file> [--message txt] [--dry-run]
"""
import json, urllib.request, websocket, time, sys, argparse, importlib.util


def _load_accounts():
    """Build the accounts table from the NODES.md registry — new nodes
    propagate automatically instead of being hardcoded here."""
    path = "/home/super/Projects/NetVM/bin/netvm-registry.py"
    spec = importlib.util.spec_from_file_location("netvm_registry", path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    accounts = {}
    for node, rec in mod.load().items():
        accounts[node] = (node, "http://127.0.0.1:%d/json/list" %
                          rec["cdp_port"])
    return accounts


ACCOUNTS = _load_accounts()

# IPs we trust for auto-approval (our infrastructure)
TRUSTED_IPS = {
    "34.139.37.135",  # VM (gateway)
    "100.123.153.75", # bl (main compute)
    "100.81.31.9",    # VM tailnet
}

def get_page(node, 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 check_approvals(ws):
    """
    Check for browser permission dialogs.
    Returns list of (dialog_text, is_trusted, action_taken).
    """
    result = ev(ws, """(() => {
      const dialogs = [];
      // Look for permission prompts (common patterns)
      const body = document.body.innerText;
      // Check for "Allow ... to share" pattern
      if (body.includes('Allow') && body.includes('to share')) {
        // Find the dialog
        const els = [...document.querySelectorAll('*')].filter(el => {
          const t = el.innerText || '';
          return t.includes('Allow') && t.includes('to share') && t.length < 500;
        });
        for (const el of els.slice(0,3)) {
          dialogs.push(el.innerText.slice(0,200));
        }
      }
      // Check for other permission patterns
      const perm_btns = [...document.querySelectorAll('button')].filter(b => {
        const t = (b.innerText||'').toLowerCase();
        return t.includes('allow') || t.includes('deny') || t.includes('block');
      });
      if (perm_btns.length >= 2 && dialogs.length === 0) {
        // Might be a permission dialog
        const parent = perm_btns[0].closest('div');
        if (parent) dialogs.push(parent.innerText.slice(0,200));
      }
      return JSON.stringify(dialogs);
    })()""")
    try:
        dialogs = json.loads(result) if result else []
    except:
        dialogs = []
    
    actions = []
    for d in dialogs:
        # Extract IP if present
        import re
        ips = re.findall(r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b', d)
        is_trusted = any(ip in TRUSTED_IPS for ip in ips)
        if is_trusted:
            # Auto-approve: click "Allow once" or "Allow"
            clicked = ev(ws, """(async()=>{
              const b = [...document.querySelectorAll('button')].find(x=>{
                const t = (x.innerText||'').toLowerCase();
                return t.includes('allow once') || t === 'allow';
              });
              if (b) { b.click(); return 'clicked:'+b.innerText.slice(0,20); }
              return 'NOTFOUND';
            })()""", True)
            actions.append((d[:80], True, clicked))
        else:
            actions.append((d[:80], False, "APPROVAL_NEEDED"))
    
    return actions

def cmd_approvals(ws):
    """Check and handle pending approvals."""
    actions = check_approvals(ws)
    if not actions:
        print("No pending approvals")
        return
    for dialog, trusted, action in actions:
        print(f"Dialog: {dialog}")
        print(f"  Trusted: {trusted}, Action: {action}")
        if not trusted:
            print("  APPROVAL_NEEDED: Manual review required")
            sys.exit(2)

def cmd_send(ws, message):
    # Check approvals first
    actions = check_approvals(ws)
    for dialog, trusted, action in actions:
        if not trusted:
            print(f"APPROVAL_NEEDED: {dialog[:80]}", file=sys.stderr)
            sys.exit(2)
    
    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, width=200):
    # Check approvals first (non-blocking)
    check_approvals(ws)
    result = ev(ws, f"""(() => {{
      const ps = [...document.querySelectorAll('p')].slice(-{n*2}).map(p=>p.innerText.slice(0,{width}));
      return ps.join('\\n---\\n');
    }})()""")
    print(result)

def cmd_wait(ws, timeout=30):
    print(f"Waiting {timeout}s for response...")
    # Check approvals periodically during wait
    for i in range(timeout // 5):
        actions = check_approvals(ws)
        for dialog, trusted, action in actions:
            if not trusted:
                print(f"APPROVAL_NEEDED: {dialog[:80]}", file=sys.stderr)
                sys.exit(2)
        time.sleep(5)
    cmd_messages(ws, 2)

def cmd_sidechat_use(ws, chat_id):
    """Open sidebar, click the box with the chat name."""
    import base64
    b64 = base64.b64encode(chat_id.encode()).decode()
    js = (
        "(async()=>{"
        # Open sidebar first
        "const ob=[...document.querySelectorAll('button')].find("
        "b=>b.textContent.includes('Open chat and side chats'));"
        "if(ob)ob.click();"
        "await new Promise(r=>setTimeout(r,2000));"
        # Find and click the chat
        "const name=atob('" + b64 + "');"
        "const els=[...document.querySelectorAll('*')];"
        "const cands=els.filter(e=>e.textContent&&e.textContent.includes(name));"
        "cands.sort((a,b)=>a.textContent.length-b.textContent.length);"
        "const el=cands[0];"
        "if(!el)return 'NOTFOUND';"
        "el.click();"
        "await new Promise(r=>setTimeout(r,4000));"
        "return window.location.href;"
        "})()"
    )
    result = ev(ws, js, True)
    print(f"Navigated to: {result}")

def cmd_sidechat_main(ws):
    """Navigate back to main chat by clicking the Main chat button."""
    # Do NOT navigate to https://muse.ai/ (landing page) - it restores the
    # last-viewed chat which may be a side chat. Click Main chat instead.
    js = (
        "(async()=>{"
        "const ob=[...document.querySelectorAll('button')].find("
        "b=>b.textContent.includes('Open chat and side chats'));"
        "if(ob)ob.click();"
        "await new Promise(r=>setTimeout(r,2000));"
        "const els=[...document.querySelectorAll('*')];"
        "const cands=els.filter(e=>e.textContent&&e.textContent.trim()==='Main chat');"
        "const el=cands[0];"
        "if(!el)return 'NOTFOUND';"
        "el.click();"
        "await new Promise(r=>setTimeout(r,4000));"
        "return window.location.href;"
        "})()"
    )
    result = ev(ws, js, True)
    print(f"Back to main chat: {result}")

def cmd_sidechat_list(ws):
    """List side chats via DOM."""
    ev(ws, """(() => {
      const btn = [...document.querySelectorAll("button")].find(b => 
        b.textContent.includes("Open chat and side chats")
      );
      if (btn) btn.click();
    })()""")
    time.sleep(2)
    chats = ev(ws, """(() => {
      const text = document.body.innerText;
      const idx = text.indexOf("Side chats");
      if (idx === -1) return "NOSIDEBAR";
      return text.slice(idx, idx + 1000);
    })()""")
    print(chats[:500])

def cmd_sidechat_create(ws):
    """Create a new side chat via the 'New side chat' button. Returns the new thread URL."""
    # Ensure sidebar is open (same selector as cmd_sidechat_list: textContent)
    ev(ws, """(() => {
      const btn = [...document.querySelectorAll("button")].find(b =>
        b.textContent.includes("Open chat and side chats")
      );
      if (btn) btn.click();
    })()""")
    import time as _time
    _time.sleep(2)
    result = ev(ws, """(() => {
      const btn = [...document.querySelectorAll('button')].find(b =>
        (b.getAttribute('aria-label')||'') === 'New side chat');
      if (!btn) return 'NOT_FOUND';
      btn.click();
      return 'CLICKED';
    })()""")
    if result == 'NOT_FOUND':
        print("ERROR: 'New side chat' button not found (sidebar may be closed)", file=sys.stderr)
        sys.exit(1)
    import time as _time
    _time.sleep(5)
    url = ev(ws, "window.location.href")
    print(f"Created: {url}")
    return url




def cdp_call(ws, method, params=None):
    """Raw CDP method call for non-Runtime domains (DOM, Page, ...).
    The ev() helper only speaks Runtime.evaluate; uploads need DOM.
    Skips CDP event chatter (messages without our id) while waiting."""
    cdp_call.counter += 1
    cid = cdp_call.counter
    ws.send(json.dumps({"id": cid, "method": method,
                        "params": params or {}}))
    for _ in range(20):
        resp = json.loads(ws.recv())
        if resp.get("id") != cid:
            continue  # event, not our response
        if "error" in resp:
            raise RuntimeError("CDP %s failed: %s" % (method, resp["error"]))
        return resp.get("result", {})
    raise RuntimeError("CDP %s: no response after 20 reads" % method)


cdp_call.counter = 100


def ev1(ws, expr, await_p=False):
    """Runtime.evaluate that skips CDP event chatter while awaiting its
    response. ev() reads a single message and can catch an event
    instead (the known None-result quirk); uploads do several DOM
    calls first, so chatter is likely."""
    ws.send(json.dumps({
        "id": 1, "method": "Runtime.evaluate",
        "params": {"expression": expr, "returnByValue": True,
                   "awaitPromise": await_p}
    }))
    for _ in range(30):
        resp = json.loads(ws.recv())
        if resp.get("id") != 1:
            continue
        return resp.get("result", {}).get("result", {}).get("value")
    return None


def cmd_upload(ws, filepath, message=None, dry_run=False):
    """Attach a file to the chat composer via CDP DOM.setFileInputFiles.

    Dry-run stages the attachment and verifies the preview chip without
    sending. Otherwise sends (with optional caption) and verifies via
    read-back — never trust the CDP result alone.
    """
    import os
    if not os.path.isfile(filepath):
        print("ERROR: file not found: %s" % filepath, file=sys.stderr)
        sys.exit(1)
    actions = check_approvals(ws)
    for dialog, trusted, action in actions:
        if not trusted:
            print("APPROVAL_NEEDED: %s" % dialog[:80], file=sys.stderr)
            sys.exit(2)
    # The file input may only render after the attach button is clicked.
    node_id = 0
    for _ in range(2):
        doc = cdp_call(ws, "DOM.getDocument")
        root = doc["root"]["nodeId"]
        q = cdp_call(ws, "DOM.querySelector",
                     {"nodeId": root, "selector": "input[type=file]"})
        node_id = q.get("nodeId", 0)
        if node_id:
            break
        ev(ws, """(() => {
          const b = [...document.querySelectorAll('button')].find(x =>
            (x.getAttribute('aria-label')||'').toLowerCase().includes('attach'));
          if (b) { b.click(); return 'clicked'; }
          return 'notfound';
        })()""")
        time.sleep(2)
    if not node_id:
        print("ERROR: no file input in composer", file=sys.stderr)
        sys.exit(1)
    cdp_call(ws, "DOM.setFileInputFiles",
             {"nodeId": node_id, "files": [os.path.abspath(filepath)]})
    time.sleep(2)
    # Read-back: the attachment must be staged. React removes the
    # file input after change, so input.files is unreliable — the
    # composer's own signal is the "Remove attachment" button plus the
    # chip text beside it.
    base = os.path.basename(filepath)
    preview = ev1(ws, """(() => {
      const btns = [...document.querySelectorAll('button')].filter(b =>
        (b.getAttribute('aria-label') || '') === 'Remove attachment');
      if (!btns.length) return JSON.stringify([]);
      const chip = btns[0].closest('div');
      const text = chip ? chip.innerText.slice(0, 120) : '';
      return JSON.stringify([{button: true, chip: text}]);
    })()""")
    print("preview: %s" % preview)
    try:
        hits = json.loads(preview) if preview else []
    except Exception:
        hits = []
    if not hits:
        print("ERROR: attachment not staged (no Remove-attachment button)",
              file=sys.stderr)
        sys.exit(1)
    if dry_run:
        print("DRY-RUN OK: '%s' staged in composer, not sent." % base)
        return
    if message:
        msg_esc = message.replace('\\', '\\\\').replace('`', '\\`').replace('$', '\\$')
        ev(ws, """(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, `%s`);
          return 'caption-set';
        })()""" % msg_esc, True)
        time.sleep(1)
    sent = ev(ws, """(async()=>{
      const send = [...document.querySelectorAll('button')].find(b =>
        b.getAttribute('aria-label') &&
        b.getAttribute('aria-label').toLowerCase().includes('send'));
      if (send) { send.click(); return 'sent'; }
      return 'NOSEND';
    })()""", True)
    print("send: %s" % sent)
    time.sleep(5)
    recent = ev1(ws, """(() => {
      const ps = [...document.querySelectorAll('p')].slice(-10)
        .map(p => p.innerText.slice(0,200));
      return ps.join('\\n---\\n');
    })()""")
    if base in (recent or ''):
        print("VERIFIED: attachment '%s' present in chat." % base)
    else:
        print("WARNING: attachment not found in read-back; "
              "check the composer manually.")

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', 'approvals', 'sidechat', 'upload'])
    p.add_argument('arg', nargs='*', default=[])
    p.add_argument('--dry-run', action='store_true',
                 help='upload: stage attachment without sending')
    p.add_argument('--message', default=None,
                 help='upload: caption text sent with the file')
    args = p.parse_args()

    node, cdp_url = ACCOUNTS[args.account]
    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, " ".join(args.arg))
        elif args.command == 'messages':
            n = int(args.arg[0]) if args.arg else 5
            w = int(args.arg[1]) if len(args.arg) > 1 else 200
            cmd_messages(ws, n, w)
        elif args.command == 'wait':
            t = int(args.arg[0]) if args.arg else 30
            cmd_wait(ws, t)
        elif args.command == 'approvals':
            cmd_approvals(ws)
        elif args.command == 'sidechat':
            if not args.arg:
                print("ERROR: sidechat requires subcommand (use|main|list|create)", file=sys.stderr)
                sys.exit(1)
            sub = args.arg[0]
            if sub == "create":
                cmd_sidechat_create(ws)
            elif sub == "use":
                if len(args.arg) < 2:
                    print("ERROR: sidechat use requires chat_id", file=sys.stderr)
                    sys.exit(1)
                cmd_sidechat_use(ws, args.arg[1])
            elif sub == "main":
                cmd_sidechat_main(ws)
            elif sub == "list":
                cmd_sidechat_list(ws)
            else:
                print(f"ERROR: unknown sidechat subcommand: {sub}", file=sys.stderr)
                sys.exit(1)
        elif args.command == 'upload':
            if not args.arg:
                print("ERROR: upload requires a filepath", file=sys.stderr)
                sys.exit(1)
            cmd_upload(ws, args.arg[0], message=args.message,
                       dry_run=args.dry_run)
    finally:
        ws.close()

if __name__ == '__main__':
    main()