#!/usr/bin/env python3
"""
NIP-65 Relay Optimizer
Checks your Nostr relay list (kind 10002 or kind 3) and tests latency.
Helps you remove dead/slow relays to improve client performance.
"""
import json, time, sys
try:
    import websocket
except ImportError:
    print("Please install websocket-client: pip install websocket-client")
    sys.exit(1)
from concurrent.futures import ThreadPoolExecutor

def fetch_relays(pubkey, bootstraps=["wss://relay.primal.net", "wss://nos.lol"]):
    for url in bootstraps:
        try:
            ws = websocket.create_connection(url, timeout=5)
            # Try kind 10002
            ws.send(json.dumps(["REQ", "sub1", {"kinds": [10002], "authors": [pubkey], "limit": 1}]))
            relays = set()
            while True:
                msg = json.loads(ws.recv())
                if msg[0] == 'EOSE': break
                if msg[0] == 'EVENT':
                    tags = msg[2].get('tags', [])
                    for t in tags:
                        if t[0] == 'r': relays.add(t[1])
            if relays:
                ws.close()
                return list(relays)
                
            # Fallback to kind 3
            ws.send(json.dumps(["REQ", "sub2", {"kinds": [3], "authors": [pubkey], "limit": 1}]))
            while True:
                msg = json.loads(ws.recv())
                if msg[0] == 'EOSE': break
                if msg[0] == 'EVENT':
                    content = msg[2].get('content', '{}')
                    try:
                        r_dict = json.loads(content)
                        for r in r_dict.keys():
                            relays.add(r)
                    except Exception: pass
            ws.close()
            if relays:
                return list(relays)
        except Exception:
            continue
    return []

def check_relay(url):
    try:
        start = time.time()
        ws = websocket.create_connection(url, timeout=3)
        lat = (time.time() - start) * 1000
        ws.close()
        return (url, True, lat)
    except Exception:
        return (url, False, 0)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python3 nip65_optimizer.py <pubkey_hex>")
        sys.exit(1)
        
    pubkey = sys.argv[1]
    if pubkey.startswith("npub"):
        print("Please provide the hex pubkey, not npub. You can convert it using tools like nostr.band")
        sys.exit(1)
        
    print(f"Fetching NIP-65 relays for {pubkey[:8]}...")
    relays = fetch_relays(pubkey)
    if not relays:
        print("No relays found. Make sure you have a NIP-65 (kind 10002) or kind 3 event.")
        sys.exit(1)
        
    print(f"Found {len(relays)} relays. Testing latencies (3s timeout)...")
    with ThreadPoolExecutor(max_workers=20) as executor:
        results = list(executor.map(check_relay, relays))
    
    alive = sorted([(u, lat) for u, ok, lat in results if ok], key=lambda x: x[1])
    dead = [u for u, ok, lat in results if not ok]
    
    print(f"\n🟢 FASTEST (KEEP THESE):")
    for url, lat in alive[:5]:
        print(f" - {url} ({lat:.0f}ms)")
        
    if dead:
        print(f"\n🔴 UNREACHABLE/SLOW (REMOVE THESE):")
        for url in dead:
            print(f" - {url}")