#!/usr/bin/env python3
"""agent-cognitive-probe.py — Real-time Cognitive Sensing and Agent Menu Navigation.

Directly probes Cloud Muse / Hatch browser runtime via CDP:
1. Passive Cognitive Sensing (zero-click):
   - Token streaming / generation state (stop button presence)
   - Typing / thinking indicators
   - Status text displayed under/beside avatar (Connected, Thinking, Working)
   - Active context (Main chat vs Side chats with thread titles & snippets)
   - Input wait / parked approval detection
2. Active Profile Menu Navigation:
   - Status panel sliding surface navigation
   - Tabs: Activity (tasks & processes), Upcoming (timers & recurring cron loops),
     Approvals, and Identity.
3. Cognitive Lock Gate:
   - Protects single-threaded thought process from interruptions.
"""

import argparse
import json
import os
import subprocess
import sys
import time
import urllib.request

# Node to pinned CDP port mapping
NODE_CDP_PORTS = {
    "muse": 9410,
    "pip": 9420,
    "646": 9430,
    "opm": 9440,
    "def": 9450,
    "dev": 9455,
    "muse-main": 9410,
}

REMOTE_HOST = "100.123.153.75"  # bl control node


def is_running_on_bl():
    """Detect if we are running locally on bl or on tp/remote."""
    try:
        import socket
        hn = socket.gethostname().lower()
        if "bl" in hn:
            return True
    except Exception:
        pass
    # Check if network namespaces exist locally
    return os.path.exists("/var/run/netns/warp-muse") or os.path.exists("/run/netns/warp-muse")


def run_cdp_eval_inside_netns(node, js_code, timeout=8):
    """Executes a JS snippet against the node's browser via CDP inside its netns."""
    port = NODE_CDP_PORTS.get(node)
    if not port:
        return {"error": f"Unknown node '{node}'"}

    # Python runner script to execute inside the target environment
    runner_code = f'''
import json, sys, urllib.request
try:
    import websocket
except ImportError:
    print(json.dumps({{"error": "websocket package missing"}}))
    sys.exit(1)

try:
    with urllib.request.urlopen("http://127.0.0.1:{port}/json/list", timeout=3) as r:
        targets = json.load(r)
    pages = [t for t in targets if t.get("type") == "page"]
    if not pages:
        print(json.dumps({{"error": "No active page target"}}))
        sys.exit(0)
    
    ws_url = pages[0]["webSocketDebuggerUrl"]
    ws = websocket.create_connection(ws_url, timeout={timeout})
    
    ws.send(json.dumps({{
        "id": 1,
        "method": "Runtime.evaluate",
        "params": {{
            "expression": {json.dumps(js_code)},
            "returnByValue": True,
            "awaitPromise": True
        }}
    }}))
    
    res = None
    for _ in range(30):
        msg = json.loads(ws.recv())
        if msg.get("id") == 1:
            res = msg.get("result", {{}}).get("result", {{}}).get("value")
            break
            
    print(json.dumps({{"ok": True, "value": res}}))
except Exception as e:
    print(json.dumps({{"error": str(e)}}))
'''

    if is_running_on_bl():
        cmd = ["sudo", "ip", "netns", "exec", f"warp-{node}", "python3", "-c", runner_code]
    else:
        # Wrap via ssh to bl
        # Use python3 on bl directly executing inside netns
        remote_cmd = f"sudo ip netns exec warp-{node} python3 -c {subprocess.list2cmdline([runner_code])}"
        cmd = ["ssh", "-q", f"super@{REMOTE_HOST}", remote_cmd]

    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout + 5)
        if proc.returncode != 0 and not proc.stdout:
            return {"error": proc.stderr.strip() or f"Process exited with {proc.returncode}"}
        
        # Parse output line that contains valid json
        for line in proc.stdout.strip().splitlines():
            line = line.strip()
            if line.startswith("{") and line.endswith("}"):
                try:
                    data = json.loads(line)
                    if "ok" in data:
                        return data["value"]
                    if "error" in data:
                        return {"error": data["error"]}
                except Exception:
                    continue
        return {"error": proc.stdout.strip() or proc.stderr.strip()}
    except subprocess.TimeoutExpired:
        return {"error": "CDP probe timed out"}
    except Exception as e:
        return {"error": str(e)}


JS_PASSIVE_COGNITIVE = """(() => {
    // 1. Generation & Thinking signals
    const stopBtn = document.querySelector('[data-testid="hatch-composer-stop-button"]');
    const isGenerating = !!stopBtn;

    const typingEl = document.querySelector('[data-testid="hatch-chat-typing-indicator"]');
    const isTyping = !!typingEl && typingEl.textContent.trim().length > 0;

    // 2. Avatar / Status text
    const statusTextEl = document.querySelector('.group\\\\/status-avatar span, [class*="status-avatar"] span, span[class*="text-body-status"]');
    const avatarStatus = statusTextEl ? (statusTextEl.innerText || '').trim() : '';

    // 3. Thread context & active URL
    const url = window.location.href;
    const isMainChat = url === 'https://muse.ai/' || url.endsWith('/thread/new');
    const pageTitle = document.title;

    // 4. Side chats overview
    const sideRows = Array.from(document.querySelectorAll('[data-testid="hatch-thread-row"]')).map(r => {
        const text = (r.innerText || '').trim().replace(/\\\\n+/g, ' — ');
        return text;
    }).slice(0, 5);

    // 5. Input wait / Parked prompt cards
    // Detect buttons asking for approval / input in the message flow
    const actionBtns = Array.from(document.querySelectorAll('div[data-message-id] button')).map(b => (b.innerText || '').trim()).filter(t => /approve|confirm|proceed|resume|start|allow/i.test(t));
    const isInputWait = actionBtns.length > 0 || /asking for input|pending approval/i.test(document.body.innerText.slice(-600));

    // 6. Status panel state
    const panel = document.querySelector('[data-testid="hatch-status-panel-sliding-surface"]');
    const panelOpen = !!panel && panel.getBoundingClientRect().width > 0;

    return {
        url: url,
        title: pageTitle,
        is_main_chat: isMainChat,
        is_generating: isGenerating,
        is_typing: isTyping,
        avatar_status: avatarStatus,
        is_input_wait: isInputWait,
        pending_actions: actionBtns,
        panel_open: panelOpen,
        side_chats: sideRows
    };
})()"""


def get_passive_cognitive_state(node):
    """Gathers passive cognitive signals without altering UI state."""
    if not is_running_on_bl():
        try:
            cmd = ["ssh", "-q", "-o", "ConnectTimeout=5", f"super@{REMOTE_HOST}",
                   f"python3 /home/super/Projects/NetVM/bin/agent-cognitive-probe.py status {node} --json"]
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if proc.returncode == 0 and proc.stdout.strip():
                data = json.loads(proc.stdout.strip())
                if isinstance(data, list) and len(data) > 0:
                    return data[0]
                elif isinstance(data, dict):
                    return data
        except Exception as e:
            return {
                "node": node,
                "status": "DARK",
                "error": f"Remote delegation failed: {e}",
                "cognitive_lock": False,
                "lock_reason": None,
            }

    raw = run_cdp_eval_inside_netns(node, JS_PASSIVE_COGNITIVE)
    if not isinstance(raw, dict) or "error" in raw:
        return {
            "node": node,
            "status": "DARK",
            "error": raw.get("error", "Unknown error") if isinstance(raw, dict) else str(raw),
            "cognitive_lock": False,
            "lock_reason": None,
        }

    is_generating = raw.get("is_generating", False)
    is_typing = raw.get("is_typing", False)
    is_input_wait = raw.get("is_input_wait", False)
    avatar_status = raw.get("avatar_status", "")
    is_main = raw.get("is_main_chat", True)

    # Determine synthesized cognitive state
    if is_generating or is_typing or "thinking" in avatar_status.lower():
        state = "THINKING"
        locked = True
        reason = "Agent is actively generating tokens / thinking (stop button active)"
    elif is_input_wait:
        state = "INPUT_WAIT"
        locked = True
        reason = "Agent is waiting for operator or system input on a parked prompt"
    elif "working" in avatar_status.lower() or "making" in avatar_status.lower():
        state = "WORKING"
        locked = True
        reason = f"Avatar status indicates work in progress: '{avatar_status}'"
    elif not is_main:
        state = "SIDECHAT_IDLE"
        locked = False
        reason = None
    else:
        state = "IDLE"
        locked = False
        reason = None

    return {
        "node": node,
        "status": state,
        "cognitive_lock": locked,
        "lock_reason": reason,
        "details": raw,
    }


JS_NAVIGATE_MENU_TEMPLATE = """(async () => {
    // 1. Ensure panel is open
    let panel = document.querySelector('[data-testid="hatch-status-panel-sliding-surface"]');
    if (!panel) {
        // Try clicking avatar or trigger
        const avatarBtn = document.querySelector('[role="img"][aria-label*="avatar"], button[aria-label*="avatar"], .group\\\\/status-avatar');
        if (avatarBtn) avatarBtn.click();
        await new Promise(r => setTimeout(r, 350));
    }
    
    // 2. Click requested tab if specified
    const targetTab = "%(tab)s";
    if (targetTab && targetTab !== "all") {
        const btn = document.querySelector('button[aria-label="' + targetTab + '"]');
        if (btn) {
            btn.click();
            await new Promise(r => setTimeout(r, 400));
        }
    }

    panel = document.querySelector('[data-testid="hatch-status-panel-sliding-surface"]');
    if (!panel) return {error: "Status panel not rendered"};

    // Read full text and structural items
    const rawText = panel.innerText || '';
    const lines = rawText.split('\\n').map(s => s.trim()).filter(Boolean);

    // Extract activity / timer items
    const items = [];
    const buttons = Array.from(panel.querySelectorAll('button')).filter(b => !b.getAttribute('aria-label') && (b.innerText || '').length > 0);
    buttons.forEach(b => {
        const text = (b.innerText || '').trim();
        const parts = text.split('\\n').map(s => s.trim()).filter(Boolean);
        if (parts.length >= 2) {
            items.push({
                title: parts[0],
                detail: parts[1],
                time: parts.length > 2 ? parts[2] : null
            });
        }
    });

    return {
        raw_text: rawText,
        lines: lines,
        items: items
    };
})()"""


def get_agent_menu(node, tab="all"):
    """Navigates and extracts data from the agent profile / status panel."""
    if not is_running_on_bl():
        try:
            cmd = ["ssh", "-q", "-o", "ConnectTimeout=5", f"super@{REMOTE_HOST}",
                   f"python3 /home/super/Projects/NetVM/bin/agent-cognitive-probe.py menu {node} {tab} --json"]
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
            if proc.returncode == 0 and proc.stdout.strip():
                return json.loads(proc.stdout.strip())
        except Exception as e:
            return {
                "node": node,
                "tab": tab,
                "data": {"error": f"Remote delegation failed: {e}"}
            }

    tab_map = {
        "activity": "Activity",
        "upcoming": "Upcoming",
        "approvals": "Approvals",
        "identity": "Identity",
        "all": "all",
    }
    target_tab = tab_map.get(tab.lower(), "Activity")

    # If all is requested, gather activity, upcoming, and approvals
    if target_tab == "all":
        result = {}
        for sub_tab in ["Activity", "Upcoming", "Approvals"]:
            js = JS_NAVIGATE_MENU_TEMPLATE % {"tab": sub_tab}
            tab_res = run_cdp_eval_inside_netns(node, js)
            result[sub_tab.lower()] = tab_res
        return {
            "node": node,
            "menu": result
        }

    js = JS_NAVIGATE_MENU_TEMPLATE % {"tab": target_tab}
    res = run_cdp_eval_inside_netns(node, js)
    return {
        "node": node,
        "tab": target_tab,
        "data": res
    }


JS_LIVE_SCREEN = """(() => {
    const ps = Array.from(document.querySelectorAll('p')).map(p => (p.innerText || '').trim()).filter(Boolean);
    const actionBtns = Array.from(document.querySelectorAll('div[data-message-id] button, [role="log"] button, div[role="status"] button')).map(b => (b.innerText || '').trim()).filter(Boolean);
    const stopBtn = !!document.querySelector('[data-testid="hatch-composer-stop-button"]');
    const typing = !!document.querySelector('[data-testid="hatch-chat-typing-indicator"]:not(:empty)');
    const statusTextEl = document.querySelector('.group\\\\/status-avatar span, [class*="status-avatar"] span, span[class*="text-body-status"]');
    const avatarStatus = statusTextEl ? (statusTextEl.innerText || '').trim() : '';

    return {
        title: document.title,
        url: window.location.href,
        is_main_chat: window.location.href === 'https://muse.ai/' || window.location.href.endsWith('/thread/new'),
        is_generating: stopBtn,
        is_typing: typing,
        avatar_status: avatarStatus,
        recent_paragraphs: ps.slice(-8),
        action_buttons: actionBtns.filter(t => /approve|confirm|proceed|resume|start|allow|review/i.test(t))
    };
})()"""


def get_agent_live_screen(node: str) -> dict:
    """Extracts live active chat text, thoughts, prompt blocks, and threads."""
    if not is_running_on_bl():
        try:
            cmd = ["ssh", "-q", "-o", "ConnectTimeout=5", f"super@{REMOTE_HOST}",
                   f"python3 /home/super/Projects/NetVM/bin/agent-cognitive-probe.py read {node} --json"]
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
            if proc.returncode == 0 and proc.stdout.strip():
                return json.loads(proc.stdout.strip())
        except Exception as e:
            return {"node": node, "error": f"Remote delegation failed: {e}"}

    screen = run_cdp_eval_inside_netns(node, JS_LIVE_SCREEN)
    if not isinstance(screen, dict) or "error" in screen:
        return {"node": node, "error": screen.get("error", "Failed to inspect screen") if isinstance(screen, dict) else str(screen)}

    sidechats = []
    try:
        cli_cmd = ["/home/super/Projects/NetVM/bin/muse-cli-node", node, "threads"]
        proc = subprocess.run(cli_cmd, capture_output=True, text=True, timeout=8)
        if proc.returncode == 0 and proc.stdout.strip():
            threads_data = json.loads(proc.stdout.strip())
            if isinstance(threads_data, list):
                sidechats = threads_data[:8]
    except Exception:
        pass

    return {
        "node": node,
        "screen": screen,
        "sidechats": sidechats
    }


JS_NAVIGATE_MAIN_CHAT = """(() => {
    try {
        const url = window.location.href;
        if (url === 'https://muse.ai/' || url === 'https://muse.ai/thread/new') {
            return { ok: true, already_main: true };
        }
        const candidates = Array.from(document.querySelectorAll('a, button, [role="button"], [data-testid]'));
        const mainBtn = candidates.find(b => {
            const t = (b.innerText || '').trim().toLowerCase();
            return t === 'main chat' || b.getAttribute('aria-label') === 'Main chat' || b.getAttribute('data-testid') === 'hatch-sidebar-main-chat';
        });
        if (mainBtn) {
            mainBtn.click();
            return { ok: true, method: 'click' };
        }
        window.location.href = 'https://muse.ai/';
        return { ok: true, method: 'navigate' };
    } catch (e) {
        return { error: String(e) };
    }
})()"""


def navigate_to_main_chat(node: str) -> dict:
    """Navigates the agent browser session back to Main Chat (https://muse.ai/)."""
    if not is_running_on_bl():
        try:
            cmd = ["ssh", "-q", "-o", "ConnectTimeout=5", f"super@{REMOTE_HOST}",
                   f"python3 /home/super/Projects/NetVM/bin/agent-cognitive-probe.py nav-main {node} --json"]
            proc = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if proc.returncode == 0 and proc.stdout.strip():
                return json.loads(proc.stdout.strip())
        except Exception as e:
            return {"error": str(e)}

    return run_cdp_eval_inside_netns(node, JS_NAVIGATE_MAIN_CHAT)


def cmd_nav_main(args):
    node = getattr(args, "agent", None) or getattr(args, "node", None)
    res = navigate_to_main_chat(node)
    if getattr(args, "json", False):
        print(json.dumps(res, indent=2))
    else:
        if isinstance(res, dict) and res.get("ok"):
            m = res.get("method") or ("already in main chat" if res.get("already_main") else "default")
            print(f"✓ Refocused agent '{node}' to Main Chat ({m})")
        else:
            err = res.get("error") if isinstance(res, dict) else str(res)
            print(f"❌ Failed to refocus '{node}' to Main Chat: {err}")


def cmd_read(args):
    node = getattr(args, "agent", None) or getattr(args, "node", None)
    data = get_agent_live_screen(node)
    if getattr(args, "json", False):
        print(json.dumps(data, indent=2))
        return

    if "error" in data:
        print(f"\n❌ Error inspecting screen for {node}: {data['error']}\n")
        return

    screen = data.get("screen", {})
    sidechats = data.get("sidechats", [])

    url = screen.get("url", "")
    mode = "Main Chat" if screen.get("is_main_chat") else "Side Chat"
    title = screen.get("title", "")
    avatar = screen.get("avatar_status") or "Connected"
    gen = "🧠 GENERATING" if screen.get("is_generating") else ("💭 TYPING" if screen.get("is_typing") else "🟢 SETTLED")

    print(f"\n=== LIVE THOUGHT STREAM & ACTIVE CHAT: {node.upper()} ===")
    print(f"Context: {mode} ({url})")
    print(f"Title:   {title}")
    print(f"Status:  {avatar} | State: {gen}")

    paragraphs = screen.get("recent_paragraphs", [])
    if paragraphs:
        print("\n--- ACTIVE CONVERSATION & THOUGHT PARAGRAPHS ---")
        for p in paragraphs:
            print(f"  • {p}\n")
    else:
        print("\n  (No text paragraphs visible in current viewport)")

    actions = screen.get("action_buttons", [])
    if actions:
        print("--- PENDING ACTION CARDS / APPROVAL BUTTONS ---")
        for a in actions:
            print(f"  ⚠️ [PROMPT ACTION] {a}")
        print()

    if sidechats:
        print("--- RECENT SIDE CHATS & TOPICS ---")
        for sc in sidechats:
            sid = sc.get("session_id", "")[:8]
            stitle = sc.get("title") or "(Untitled sidechat)"
            upd = sc.get("updated", "")
            print(f"  • [{sid}] {stitle} ({upd})")
        print()


def format_cognitive_badge(status):
    badges = {
        "IDLE": "🟢 IDLE",
        "SIDECHAT_IDLE": "💬 SIDE_IDLE",
        "THINKING": "🧠 THINKING",
        "INPUT_WAIT": "⏸️ INPUT_WAIT",
        "BUSY_SIDECHAT": "💬 SIDECHAT",
        "WORKING": "⚙️ WORKING",
        "DARK": "⚫ DARK",
    }
    return badges.get(status, f"❓ {status}")


def cmd_status(args):
    nodes = getattr(args, "agents", None) or getattr(args, "nodes", None) or ["muse", "pip", "646", "opm", "dev", "def"]
    results = []
    for n in nodes:
        results.append(get_passive_cognitive_state(n))

    if getattr(args, "json", False):
        print(json.dumps(results, indent=2))
        return

    print("\n=== AGENT COGNITIVE SENSOR (LIVE DOM PROBE) ===")
    print(f"{'AGENT':<10} {'COGNITIVE STATE':<18} {'LOCK':<8} {'UNDER-AVATAR':<15} {'DETAILS / CURRENT THOUGHT':<40}")
    print("-" * 95)
    for r in results:
        node = r["node"]
        status = r["status"]
        badge = format_cognitive_badge(status)
        locked = "LOCKED" if r.get("cognitive_lock") else "OPEN"
        det = r.get("details", {})
        avatar_status = det.get("avatar_status", "-")
        reason = r.get("lock_reason") or det.get("title", "Settled")
        if status == "DARK":
            reason = r.get("error", "CDP unreachable")
            avatar_status = "DARK"
        print(f"{node:<10} {badge:<18} {locked:<8} {avatar_status:<15} {reason[:40]:<40}")
    print()


def cmd_menu(args):
    node = getattr(args, "agent", None) or getattr(args, "node", None)
    tab = getattr(args, "tab", "activity") or "activity"
    data = get_agent_menu(node, tab)

    if getattr(args, "json", False):
        print(json.dumps(data, indent=2))
        return

    print(f"\n=== AGENT MENU: {node.upper()} (TAB: {tab.upper()}) ===")
    if tab.lower() == "all":
        menu = data.get("menu", {})
        for tname, tdata in menu.items():
            print(f"\n--- {tname.upper()} ---")
            if isinstance(tdata, dict) and "error" in tdata:
                print(f"  Error: {tdata['error']}")
            elif isinstance(tdata, dict):
                items = tdata.get("items", [])
                if items:
                    for it in items:
                        t_str = f" [{it['time']}]" if it.get("time") else ""
                        print(f"  • {it['title']}: {it['detail']}{t_str}")
                else:
                    raw = tdata.get("raw_text", "")
                    for line in raw.split("\n"):
                        if line.strip():
                            print(f"  {line.strip()}")
        print()
        return

    tdata = data.get("data", {})
    if isinstance(tdata, dict) and "error" in tdata:
        print(f"  Error: {tdata['error']}")
    elif isinstance(tdata, dict):
        items = tdata.get("items", [])
        if items:
            for it in items:
                t_str = f" [{it['time']}]" if it.get("time") else ""
                print(f"  • {it['title']}: {it['detail']}{t_str}")
        else:
            raw = tdata.get("raw_text", "")
            for line in raw.split("\n"):
                if line.strip():
                    print(f"  {line.strip()}")
    print()


def cmd_lock(args):
    node = getattr(args, "agent", None) or getattr(args, "node", None)
    state = get_passive_cognitive_state(node)
    if args.json:
        print(json.dumps(state, indent=2))
    else:
        if state.get("cognitive_lock"):
            print(f"🔴 COGNITIVE LOCK ENGAGED on '{node}' ({state['status']}): {state.get('lock_reason')}")
        else:
            print(f"🟢 COGNITIVELY IDLE: Agent '{node}' is free to receive new work without interruption.")

    if state.get("cognitive_lock") and not getattr(args, "force", False):
        sys.exit(1)
    sys.exit(0)


def main():
    parser = argparse.ArgumentParser(description="Real-time Cognitive Sensing and Agent Menu Navigation")
    sub = parser.add_subparsers(dest="command")

    p_status = sub.add_parser("status", help="Show cognitive state for all or selected agents")
    p_status.add_argument("nodes", nargs="*", help="Optional agent names")
    p_status.add_argument("--json", action="store_true", help="Output JSON")

    p_menu = sub.add_parser("menu", help="Navigate agent profile menu (tasks, timers, approvals, identity)")
    p_menu.add_argument("node", help="Agent name (muse, pip, 646, opm, dev, def)")
    p_menu.add_argument("tab", nargs="?", default="activity", choices=["activity", "upcoming", "approvals", "identity", "all"], help="Menu tab to view")
    p_menu.add_argument("--json", action="store_true", help="Output JSON")

    p_lock = sub.add_parser("lock", help="Check cognitive lock before dispatching work")
    p_lock.add_argument("node", help="Agent name")
    p_lock.add_argument("--force", action="store_true", help="Bypass lock check")
    p_lock.add_argument("--json", action="store_true", help="Output JSON")

    p_read = sub.add_parser("read", help="Extract live active chat text, thought stream, and side chats")
    p_read.add_argument("node", help="Agent name")
    p_read.add_argument("--json", action="store_true", help="Output JSON")

    p_nav = sub.add_parser("nav-main", help="Navigate agent browser session back to Main Chat")
    p_nav.add_argument("node", help="Agent name")
    p_nav.add_argument("--json", action="store_true", help="Output JSON")

    args = parser.parse_args()
    if not args.command:
        # Default to status
        args.nodes = []
        args.json = False
        cmd_status(args)
        return

    if args.command == "status":
        cmd_status(args)
    elif args.command == "menu":
        cmd_menu(args)
    elif args.command == "read":
        cmd_read(args)
    elif args.command == "lock":
        cmd_lock(args)
    elif args.command == "nav-main":
        cmd_nav_main(args)


if __name__ == "__main__":
    main()