import subprocess
import json
import time

notebook_id = "e8a30630-0df2-4932-8740-9ffe7d8882bc"
nlm_bin = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"
gapi_bin = "/root/.hermes/profiles/khoi/skills/productivity/google-workspace/scripts/google_api.py"

# Get current sources in notebook to avoid re-adding
res = subprocess.run([nlm_bin, "source", "list", notebook_id, "--json"], capture_output=True, text=True)
try:
    current_sources = json.loads(res.stdout)
    existing_titles = set(s.get("title", "") for s in current_sources)
except Exception:
    existing_titles = set()

print(f"Already in notebook: {len(existing_titles)} sources")

# 1. Micro PA folder files
micro_pa_files = [
    ("1gZbwzsYyXpIEvduOLIEmZZrNfxnZhEoh", "[0311] SAB, Ch. 1 - Micro For Public Decisions Overview.pdf", "pdf"),
    ("1mVd8jxkiQN4KU_sAiNHZX1xa4xGpdxqW", "GM, Ch.1- 10 Ly Thuyet Kinh Te Hoc.pdf", "pdf"),
]

# 2. Reading subfolders
reading_folders = [
    "1ErHXsmL1qrVyyn7jZLUCr3BVejWWUJk7",
    "1uZ-j2KphGnGTax0Nwro0KlSZfBlaYyh7",
    "1jsv4G_NhNlABIDJVAVEQLqIhLMT5Xhz4",
    "1uNCQ2b8je8R8MlwBr5zqfoRXDI1cwwDT",
    "1k7CcgXLjT_bzenNsbNqdab8MEA7c3my1",
    "1XfRWW9qWQ_ViXc4ldp5lwZk1SE8fk2lc",
    "1jspwpzDnZJoJ4os4AmqiJciPSmXl5zDV",
    "1jjUmNcsr7EsACFrDK8LiCrce7_M7klG5",
    "1RZVs10AAc92-9JJ94cpmMf0dyvTvKJQO",
    "1KnrVlyemfznOvsBjQbo6j26ChBL_9Urm",
]

to_add = list(micro_pa_files)

for folder_id in reading_folders:
    query = f"'{folder_id}' in parents and trashed = false"
    r = subprocess.run(['python3', gapi_bin, 'drive', 'search', '--raw-query', query], capture_output=True, text=True)
    try:
        files = json.loads(r.stdout)
        for f in files:
            mime = f.get('mimeType', '')
            name = f.get('name', '')
            fid = f.get('id', '')
            if mime == 'application/pdf':
                to_add.append((fid, name, 'pdf'))
            elif mime in ['application/vnd.google-apps.document', 'application/msword', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document']:
                to_add.append((fid, name, 'doc'))
    except Exception as e:
        print("Error reading folder", folder_id, e)

print(f"Total candidate files to add: {len(to_add)}")

added_count = 0
for fid, title, ftype in to_add:
    # Check if title already in notebook
    clean_title = title.replace(".pdf", "").replace(".doc", "").replace(".docx", "")
    if any(clean_title in ext for ext in existing_titles):
        print(f"Skipping (already exists): {title}")
        continue
    
    print(f"Adding [{ftype}] {title} ({fid})...")
    cmd = [nlm_bin, 'source', 'add', notebook_id, '--drive', fid, '--title', title, '--type', ftype]
    sub_res = subprocess.run(cmd, capture_output=True, text=True)
    if sub_res.returncode == 0:
        print(f"  ✓ Added: {title}")
        added_count += 1
        existing_titles.add(title)
    else:
        print(f"  ✗ Failed: {sub_res.stderr.strip()} | {sub_res.stdout.strip()}")
    time.sleep(1)

print(f"\nDone! Successfully added {added_count} new sources.")