import json
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor

notebook_id = "a0f8646b-90b1-4a65-bddc-07280aa06a99"
nlm_path = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"

print("Fetching source list...")
result = subprocess.run([nlm_path, "source", "list", notebook_id, "--json"], capture_output=True, text=True)
sources = json.loads(result.stdout)
print(f"Total sources: {len(sources)}")

results = []
counter = 0

def describe_source(src):
    global counter
    src_id = src["id"]
    src_title = src["title"]
    start = time.time()
    desc_res = subprocess.run([nlm_path, "source", "describe", src_id, "--json"], capture_output=True, text=True)
    duration = time.time() - start
    try:
        data = json.loads(desc_res.stdout)
        summary = data.get("summary", "")
        keywords = data.get("keywords", [])
    except Exception:
        summary = desc_res.stdout.strip()
        keywords = []
    
    print(f"[{len(results)+1}/{len(sources)}] Described '{src_title}' in {duration:.2f}s")
    return {
        "id": src_id,
        "title": src_title,
        "summary": summary,
        "keywords": keywords
    }

# Using ThreadPoolExecutor to run describes in parallel.
# Google APIs might have rate limits, so we choose a moderate worker count.
with ThreadPoolExecutor(max_workers=8) as executor:
    results = list(executor.map(describe_source, sources))

with open("source_summaries.json", "w", encoding="utf-8") as f:
    json.dump(results, f, indent=2, ensure_ascii=False)

print("Done! Saved to source_summaries.json")