#!/usr/bin/env python3
"""Detach all users/* blocks from a Letta agent.

Usage:
    python detach_user_blocks.py [agent_id]

If agent_id is not provided, uses the VOID_AGENT_ID env var or defaults to void-3.
"""

import os
import sys
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

API_URL = "https://api.letta.com"
API_KEY = os.environ.get("CAMERON_LETTA_API_KEY") or os.environ.get("LETTA_API_KEY")

def get_agent_id():
    if len(sys.argv) > 1:
        return sys.argv[1]
    return os.environ.get("VOID_AGENT_ID", "agent-c52e371d-cb89-4f99-8634-b606bdc69cad")

def get_user_blocks(agent_id: str) -> list[dict]:
    """Get all user blocks, handling pagination."""
    headers = {"Authorization": f"Bearer {API_KEY}"}
    all_blocks = []
    after = None
    
    while True:
        url = f"{API_URL}/v1/agents/{agent_id}/core-memory/blocks"
        params = {"limit": 1000}
        if after:
            params["after"] = after
            
        resp = requests.get(url, headers=headers, params=params)
        resp.raise_for_status()
        blocks = resp.json()
        
        if not blocks:
            break
            
        # Filter for user blocks
        user_blocks = [b for b in blocks if b.get("label", "").startswith("users/")]
        all_blocks.extend(user_blocks)
        
        # Get the last block ID for pagination
        if len(blocks) < 1000:
            break
        after = blocks[-1]["id"]
    
    return all_blocks

def detach_block(agent_id: str, block_id: str) -> bool:
    """Detach a single block from an agent."""
    headers = {"Authorization": f"Bearer {API_KEY}"}
    url = f"{API_URL}/v1/agents/{agent_id}/core-memory/blocks/detach/{block_id}"
    
    resp = requests.patch(url, headers=headers)
    return resp.status_code == 200

def main():
    agent_id = get_agent_id()
    print(f"Agent: {agent_id}")
    
    print("Fetching user blocks...")
    blocks = get_user_blocks(agent_id)
    print(f"Found {len(blocks)} user blocks")
    
    if not blocks:
        print("No user blocks to detach!")
        return
    
    print("Detaching blocks...")
    detached = 0
    
    with ThreadPoolExecutor(max_workers=50) as executor:
        futures = {
            executor.submit(detach_block, agent_id, b["id"]): b 
            for b in blocks
        }
        
        for future in as_completed(futures):
            block = futures[future]
            try:
                if future.result():
                    detached += 1
                    if detached % 100 == 0:
                        print(f"  Progress: {detached}/{len(blocks)}")
            except Exception as e:
                print(f"  Error detaching {block['id']}: {e}")
    
    print(f"Detached {detached}/{len(blocks)} blocks")
    
    # Verify
    print("Verifying...")
    remaining = get_user_blocks(agent_id)
    print(f"Remaining user blocks: {len(remaining)}")

if __name__ == "__main__":
    main()