#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os, time, json, ssl, base64
import websocket
import logging

URL = "ws://127.0.0.1:18789/ws"
TOKEN = "123456"
SESSION_KEY = "agent:main:main"
DEVICE_STORAGE = os.path.expanduser("~/.openclaw/workspace/.xiaomeng_correct_device.json")

logging.basicConfig(format='%(asctime)s.%(msecs)03d %(message)s', datefmt='%H:%M:%S', level=logging.INFO)
log = logging.info

class DeviceIdentity:
    def __init__(self, storage_path):
        self.storage_path = storage_path
        self.load()

    def load(self):
        if not os.path.exists(self.storage_path):
            log(f"❌ Error: device file not found {self.storage_path}")
            os._exit(1)
        with open(self.storage_path, 'r') as f: data = json.load(f)
        from cryptography.hazmat.primitives.asymmetric import ed25519
        pem_lines = data['privateKeyPem'].strip().split('\n')
        b64 = ''.join([l for l in pem_lines if not l.startswith('---')])
        self.priv = ed25519.Ed25519PrivateKey.from_private_bytes(base64.b64decode(b64)[16:])
        self.device_id = data['deviceId']
        from cryptography.hazmat.primitives import serialization
        pub_bytes = self.priv.public_key().public_bytes(
            encoding=serialization.Encoding.Raw, format=serialization.PublicFormat.Raw
        )
        self.pub_b64 = base64.b64encode(pub_bytes).decode('ascii').rstrip('=').replace('+','-').replace('/','_')

    def sign(self, p):
        sig = self.priv.sign(p.encode('utf-8'))
        return base64.b64encode(sig).decode('ascii').rstrip('=').replace('+','-').replace('/','_')

class LatencyTester:
    def __init__(self):
        self.device = DeviceIdentity(DEVICE_STORAGE)
        self.t0_send = 0
        self.t1_start = 0
        self.t2_first = 0
        self.has_sent = False

    def start(self):
        log("🚀 Connecting to OpenClaw...")
        self.ws = websocket.WebSocketApp(URL, header={'Authorization': f'Bearer {TOKEN}'},
            on_open=self.on_open, on_message=self.on_message, on_error=self.on_error)
        self.ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

    def on_open(self, ws): log("🔗 Connected")
    def on_error(self, ws, e): log(f"❌ Error: {e}")

    def send_test_msg(self, ws):
        if self.has_sent: return
        self.has_sent = True
        msg_id = f"m-{int(time.time()*1000)}"
        idem_key = f"s-{int(time.time()*1000)}"
        req = {
            "type": "req",
            "id": msg_id,
            "method": "chat.send",
            "params": {
                "sessionKey": SESSION_KEY,
                "message": "（延迟测试）",
                "deliver": True,
                "idempotencyKey": idem_key
            }
        }
        self.t0_send = time.time()
        log(f"📤 [T0] Sending message (ID: {msg_id})...")
        ws.send(json.dumps(req))

    def on_message(self, ws, message):
        try:
            d = json.loads(message)
            ev = d.get('event')
            pl = d.get('payload', {})

            if ev == 'connect.challenge':
                ts = pl['ts']
                nonce = pl['nonce']
                auth_str = f"v2|{self.device.device_id}|webchat|webchat|operator|operator.admin,operator.write,operator.approvals,operator.pairing,operator.read|{ts}|{TOKEN}|{nonce}"
                sig = self.device.sign(auth_str)
                ws.send(json.dumps({
                    "type":"req","id":"auth","method":"connect",
                    "params":{
                        "minProtocol":3,"maxProtocol":3,"role":"operator",
                        "scopes":["operator.admin","operator.write","operator.approvals","operator.pairing","operator.read"],
                        "auth":{"token":TOKEN},
                        "device":{"id":self.device.device_id,"publicKey":self.device.pub_b64,"signature":sig,"signedAt":ts,"nonce":nonce},
                        "client":{"id":"webchat","version":"1.0","platform":"linux","mode":"webchat"}
                    }
                }))

            elif d.get('id') == 'auth':
                if d.get('ok'):
                    log("✅ Authentication successful, sending test message...")
                    self.send_test_msg(ws)
                else:
                    log(f"❌ Authentication failed: {d.get('error')}")

            elif ev == 'agent':
                stream = pl.get('stream')
                phase = pl.get('data', {}).get('phase')

                if stream == 'lifecycle' and phase == 'start':
                    self.t1_start = time.time()
                    diff = int((self.t1_start - self.t0_send) * 1000)
                    log(f"⏱️ [T1] Agent started (Phase: Start): +{diff} ms")

                elif stream == 'assistant':
                    if 'text' in pl.get('data', {}) or 'delta' in pl.get('data', {}):
                        if self.t2_first == 0:
                            self.t2_first = time.time()
                            self.print_report()
                            ws.close()

        except Exception as e:
            log(f"Error processing message: {e}")

    def print_report(self):
        t0 = self.t0_send
        t1 = self.t1_start
        t2 = self.t2_first
        if t1 == 0: t1 = t2
        p1 = int((t1 - t0) * 1000)
        p2 = int((t2 - t1) * 1000)
        total = int((t2 - t0) * 1000)
        print("\n" + "="*40)
        print("📊 Latency Report")
        print("="*40)
        print(f"1️⃣ Phase 1 (Scheduling/RAG/History): {p1} ms")
        print(f"2️⃣ Phase 2 (LLM Generation):         {p2} ms")
        print("-" * 40)
        print(f"🏁 Total TTFT:                       {total} ms")
        print("="*40 + "\n")

if __name__ == "__main__":
    LatencyTester().start()