import sys, json, time, websocket, os

if len(sys.argv) < 2:
    print("Usage: python3 nostr_export.py <hex_pubkey>")
    sys.exit(1)

pubkey = sys.argv[1]
if len(pubkey) != 64:
    print("Error: pubkey must be 64-character hex.")
    sys.exit(1)

out_dir = f"nostr_backup_{pubkey[:8]}"
os.makedirs(out_dir, exist_ok=True)

relays = [
    "wss://relay.damus.io",
    "wss://relay.primal.net",
    "wss://nos.lol",
    "wss://relay.nostr.band"
]

events = {}

print(f"Starting Nostr backup for {pubkey}")
print("Connecting to relays to fetch Kind 0 (profile), 1 (posts), 3 (follows), 10002 (relays)...")

for r in relays:
    try:
        ws = websocket.create_connection(r, timeout=5)
        req = ["REQ", "backup", {"authors": [pubkey], "kinds": [0, 1, 3, 10002]}]
        ws.send(json.dumps(req))
        
        ws.settimeout(2)
        deadline = time.time() + 10  # 10 seconds per relay max
        while time.time() < deadline:
            try:
                msg = json.loads(ws.recv())
                if msg[0] == "EOSE":
                    break
                if msg[0] == "EVENT":
                    ev = msg[2]
                    events[ev['id']] = ev
            except websocket.WebSocketTimeoutException:
                break
        ws.close()
        print(f"  [OK] {r} -> {len(events)} total unique events so far")
    except Exception as e:
        print(f"  [FAIL] {r} -> {e}")

if not events:
    print("No events found or all relays failed.")
    sys.exit(0)

# Categorize
profile = None
follows = None
relay_list = None
posts = []

for ev in events.values():
    if ev['kind'] == 0:
        if not profile or ev['created_at'] > profile['created_at']:
            profile = ev
    elif ev['kind'] == 3:
        if not follows or ev['created_at'] > follows['created_at']:
            follows = ev
    elif ev['kind'] == 10002:
        if not relay_list or ev['created_at'] > relay_list['created_at']:
            relay_list = ev
    elif ev['kind'] == 1:
        posts.append(ev)

# Sort posts
for p in posts:
    if 'created_at' not in p:
        p['created_at'] = 0
posts.sort(key=lambda x: x['created_at'], reverse=True)

# Save JSON
if profile:
    with open(f"{out_dir}/profile.json", "w") as f:
        json.dump(profile, f, indent=2)

if follows:
    with open(f"{out_dir}/follows.json", "w") as f:
        json.dump(follows, f, indent=2)

if relay_list:
    with open(f"{out_dir}/relays.json", "w") as f:
        json.dump(relay_list, f, indent=2)

if posts:
    with open(f"{out_dir}/posts.json", "w") as f:
        json.dump(posts, f, indent=2)

# Export posts as readable markdown
with open(f"{out_dir}/posts.md", "w") as f:
    f.write(f"# Nostr Posts Archive - {pubkey}\n\n")
    for p in posts:
        dt = time.strftime('%Y-%m-%d %H:%M:%S', time.gmtime(p.get('created_at', 0)))
        f.write(f"### {dt} (ID: {p['id']})\n")
        f.write(f"{p.get('content', '')}\n\n---\n\n")

print(f"\nDone! Backup saved to directory: {out_dir}/")
print(f"  - Profile: {'Yes' if profile else 'No'}")
print(f"  - Follows: {len(follows.get('tags', [])) if follows else 0} contacts")
print(f"  - Posts: {len(posts)}")
print(f"Readable Markdown archive created: {out_dir}/posts.md")