#!/usr/bin/env python3
"""
Simulate the browser's exact WebSocket behavior to debug Overview tab
"""
import websocket
import json
import time

TOKEN = "6a938dff0b0f06dcebabb84639d5b901d4f1aa523a9c4250"
URL = "ws://127.0.0.1:18789"

class BrowserSimulator:
    def __init__(self):
        self.ws = None
        self.connected = False
        self.connect_handled = False
        self.pending = {}  # Like pendingRequests Map
        self.request_counter = 0
        
    def update_stats(self, status):
        """Simulate updateStats function"""
        print(f"\n>>> updateStats called")
        if not status:
            print("    ERROR: status is empty!")
            return
            
        sessions = status.get('sessions', [])
        tokens = sum(s.get('tokens', 0) for s in sessions)
        cost = sum(s.get('cost', 0) for s in sessions)
        
        print(f"    Active Sessions: {len(sessions)}")
        print(f"    Agents: 3")
        print(f"    Total Tokens: {tokens}")
        print(f"    Est. Cost: ${cost:.2f}")
        
    def send_request(self, method, params):
        """Simulate sendRequest function"""
        if self.ws and self.ws.connected:
            self.request_counter += 1
            req_id = f"{int(time.time() * 1000)}_{self.request_counter}"
            self.pending[req_id] = method  # Track by ID
            
            msg = {
                "type": "req",
                "id": req_id,
                "method": method,
                "params": params
            }
            print(f"\n[SEND] {method} id={req_id}")
            self.ws.send(json.dumps(msg))
            return req_id
        return None
        
    def handle_message(self, msg):
        """Simulate handleMessage function"""
        msg_type = msg.get("type")
        
        # Handle connect challenge
        if msg_type == "event" and msg.get("event") == "connect.challenge":
            print("\n[RECV] connect.challenge")
            connect_params = {
                "minProtocol": 3,
                "maxProtocol": 3,
                "client": {
                    "id": "openclaw-control-ui",
                    "version": "1.0.0",
                    "platform": "web",
                    "mode": "ui"
                },
                "role": "operator",
                "scopes": ["operator.read", "operator.write", "operator.admin"],
                "caps": ["tool-events"],
                "userAgent": "Mozilla/5.0",
                "locale": "en-US",
                "auth": {"token": TOKEN}
            }
            self.send_request("connect", connect_params)
            return
            
        # Handle responses
        if msg_type == "res":
            # Check if connect response (first response)
            if not self.connect_handled:
                self.connect_handled = True
                if msg.get("ok"):
                    self.connected = True
                    print("\n[CONN] Connected! Requesting data...")
                    self.send_request("status", {})
                    time.sleep(0.1)  # Small delay like browser
                    self.send_request("sessions.list", {})
                else:
                    print(f"\n[ERR] Connect failed: {msg.get('error')}")
                return
            
            # Look up request method by response ID
            resp_id = msg.get("id")
            method = self.pending.get(resp_id)
            
            print(f"\n[RECV] Response id={resp_id}, method={method}")
            print(f"       pendingSize={len(self.pending)}")
            
            if resp_id in self.pending:
                del self.pending[resp_id]
            
            if method == "status":
                print(">>> Handling as status response")
                self.update_stats(msg.get("payload"))
                return
                
            if method == "sessions.list":
                print(">>> Handling as sessions.list response")
                sessions = msg.get("payload", {}).get("sessions", [])
                print(f"    Sessions: {len(sessions)}")
                return
                
            print(f">>> UNHANDLED: method={method}")
    
    def run(self):
        print("="*50)
        print("Browser WebSocket Simulator")
        print("="*50)
        
        self.ws = websocket.create_connection(URL)
        
        # Process messages
        try:
            while True:
                data = self.ws.recv()
                msg = json.loads(data)
                self.handle_message(msg)
        except websocket._exceptions.WebSocketConnectionClosedException:
            print("\n[DISC] Connection closed")
        
        self.ws.close()
        print("\nDone!")

if __name__ == "__main__":
    sim = BrowserSimulator()
    sim.run()