import json
import os
import subprocess
import time
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build

NOTEBOOK_ID = "99090327-f66b-4d8d-8b9e-a3c00acb599f"
NLM_CLI = "/opt/ai-os/products/ceo/integrations/notebooklm-mcp-cli/.venv/bin/nlm"

with open("/root/.hermes/google_token.json") as f:
    token_data = json.load(f)
service = build("drive", "v3", credentials=Credentials.from_authorized_user_info(token_data))

# 1. Foundational files
foundational_files = [
    {"id": "1nuClqkMVQ6Lij2Y_eOQdzl_lmrTxuvna", "name": "[G]_01. MPP28PA_Macroeconomics Syllabus 2026_Final_V.pdf"},
    {"id": "1-hGIH751gBsATY_ZYIAeTGx6uSNBZ1MC", "name": "[G]_1. So Tay Hoc Vien MPP28 - digital.pdf"},
    {"id": "1powzueiHeF2-bH2l_V0HwF4bTcMvbclH", "name": "[G]_2.1. Library_Rules_Copying_Guidelines_(ENG).pdf"},
    {"id": "1zg3STNcUo7XGocw4g02ikbqmcsrdfuP8", "name": "[G]_2.2. Instructions_for_Searching_and_Using_Online_Databases_(ENG).pdf"},
    {"id": "1B8JynWqyOG2oPl36CaipfCohcbfzyI3h", "name": "[G]_2.3. List_of_Online_Databases_(ENG).pdf"},
    {"id": "1071axfmvd0tTWIssvU2LB4hlQMCPU-Pn", "name": "[G]_2.4. HuongDanTrichDanTLTK_MPP2028.pdf"},
]

# 2. Scan Macro readings folder
macro_readings_folder = "1MgBFnJfvGwcQoECWJPKDkWkq-6iuomLv"

def scan_folder(fid, path=""):
    res = service.files().list(
        q=f"'{fid}' in parents and trashed = false",
        fields="files(id, name, mimeType)"
    ).execute()
    items = []
    for f in res.get("files", []):
        cur_p = f"{path}/{f['name']}" if path else f["name"]
        if f["mimeType"] == "application/vnd.google-apps.folder":
            items.extend(scan_folder(f["id"], cur_p))
        else:
            items.append({"id": f["id"], "name": f["name"], "mimeType": f["mimeType"], "path": cur_p})
    return items

drive_files = scan_folder(macro_readings_folder)

lecture_prefix_map = {
    "Bài giảng 01": "[0511]",
    "Bài giảng 02": "[1211]",
    "Bài giảng 03": "[1411]",
    "Bài giảng 04": "[1411]",
    "Bài giảng 05": "[1511]",
    "Bài giảng 06": "[1511]",
    "Bài giảng 07": "[1611]",
    "Bài giảng 08": "[1611]",
    "Bài giảng 09": "[1711]",
    "Bài giảng 10": "[1711]",
    "Bài giảng 11": "[1911]",
    "Bài giảng 12": "[1911]",
    "Bài giảng 13": "[2011]",
    "Bài giảng 14": "[2011]",
    "Bài giảng 16": "[2811]",
    "Bài giảng 17": "[2811]",
    "Bài giảng 18": "[2911]"
}

items_to_add = []
for ff in foundational_files:
    items_to_add.append((ff["id"], ff["name"], "pdf"))

for f in drive_files:
    # Skip huge Krugman file
    if "Krugman, Obstfeld, and Melitz, 2015, International Economics_ Theory and Policy.pdf" in f["name"]:
        print(f"Skipping 480MB file: {f['name']}")
        continue

    # Determine title
    path_parts = f["path"].split("/")
    folder_name = path_parts[0] if path_parts else ""
    
    if folder_name == "Sách":
        title = f"[G]_{f['name']}"
    elif folder_name in lecture_prefix_map:
        prefix = lecture_prefix_map[folder_name]
        title = f"{prefix} {f['name']}"
    else:
        title = f"[G]_{f['name']}"
    
    doc_type = "pdf"
    if "presentation" in f["mimeType"]:
        doc_type = "slides"
    elif "spreadsheet" in f["mimeType"] or "excel" in f["mimeType"]:
        doc_type = "sheets"
    elif "word" in f["mimeType"] or "document" in f["mimeType"]:
        doc_type = "doc"
        
    items_to_add.append((f["id"], title, doc_type))

print(f"Total items to add to Macro notebook: {len(items_to_add)}")

# Get existing sources in notebook to avoid duplication
existing_res = subprocess.run([NLM_CLI, "source", "list", NOTEBOOK_ID, "--json"], capture_output=True, text=True)
try:
    existing_sources = json.loads(existing_res.stdout)
    existing_titles = set(s.get("title", "") for s in existing_sources)
except Exception:
    existing_titles = set()

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

success_count = 0
fail_count = 0
for idx, (fid, title, doc_type) in enumerate(items_to_add):
    if title in existing_titles:
        print(f"[{idx+1}/{len(items_to_add)}] Already exists: {title}")
        continue
        
    print(f"[{idx+1}/{len(items_to_add)}] Adding: {title} (ID: {fid})...")
    cmd = [
        NLM_CLI, "source", "add", NOTEBOOK_ID,
        "--drive", fid,
        "--title", title,
        "--type", doc_type
    ]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode == 0:
        print(f"  ✓ Added successfully")
        success_count += 1
    else:
        print(f"  ✗ Failed: {res.stderr.strip() or res.stdout.strip()}")
        fail_count += 1
    time.sleep(1)

print(f"\nCompleted sync! Added: {success_count}, Failed: {fail_count}, Total in queue: {len(items_to_add)}")