import asyncio
from pathlib import Path
from notebooklm import NotebookLMClient

RELATIVE_TARGETS = [
    "../../AGENTS.md",
    "../../README.md",
    "../../docs/README.md",
    "../../docs/guides",
    "../../docs/architecture",
    "../../docs/specs",
    "../../docs/plans",
    "../../docs/known_issues.md",
]

NOTEBOOK_TITLE = "Stickman Studio"

def gather_markdown_files(base_path: Path):
    """Recursively collect all .md files from specified relative paths."""
    files_to_upload = []
    
    for rel_path in RELATIVE_TARGETS:
        target = (base_path / rel_path).resolve()
        
        if target.is_file() and target.suffix.lower() in [".md", ".markdown"]:
            files_to_upload.append(target)
        elif target.is_dir():
            for md_file in target.rglob("*.md"):
                files_to_upload.append(md_file)
        else:
            print(f"⚠️ Target path not found: {target}")

    return list(dict.fromkeys(files_to_upload))


async def main():
    script_dir = Path(__file__).parent.resolve()

    async with await NotebookLMClient.from_storage() as client:

        # 1. Locate 'Stickman Studio'
        print(f"🔍 Searching for notebook: '{NOTEBOOK_TITLE}'...")
        notebooks = await client.notebooks.list()
        
        target_notebook = next(
            (nb for nb in notebooks if getattr(nb, 'title', None) == NOTEBOOK_TITLE or getattr(nb, 'display_name', None) == NOTEBOOK_TITLE), 
            None
        )

        if not target_notebook:
            print(f"❌ Could not find notebook named '{NOTEBOOK_TITLE}'.")
            return

        notebook_id = getattr(target_notebook, 'id', None) or getattr(target_notebook, 'name', '')
        print(f"✅ Found notebook ID: {notebook_id}")

        # 2. Safely filter and delete ONLY non-chat file sources
        existing_sources = await client.sources.list(notebook_id=notebook_id)
        
        sources_to_delete = []
        for source in existing_sources:
            # Check source properties to identify saved chats
            mime_type = getattr(source, 'original_mime_type', None) or getattr(source, 'originalMimeType', '')
            source_type = getattr(source, 'source_type', None) or getattr(source, 'type', '')
            
            # SAFEGUARD: Keep any source identified as a chat note
            is_chat = 'chat' in str(mime_type).lower() or 'chat' in str(source_type).lower()
            
            if is_chat:
                source_title = getattr(source, 'title', None) or getattr(source, 'display_name', 'Chat')
                print(f"🛡️ Preserving past chat: '{source_title}'")
            else:
                sources_to_delete.append(source)

        if sources_to_delete:
            print(f"\n🗑️ Removing {len(sources_to_delete)} file sources (keeping chats intact)...")
            for source in sources_to_delete:
                source_id = getattr(source, 'id', None) or getattr(source, 'name', '')
                source_title = getattr(source, 'title', None) or getattr(source, 'display_name', 'File Source')
                await client.sources.delete(notebook_id=notebook_id, source_id=source_id)
                print(f"   - Removed: {source_title}")
        else:
            print("ℹ️ No old file sources to remove.")

        # 3. Collect target Markdown files
        md_files = gather_markdown_files(script_dir)
        print(f"\n📁 Found {len(md_files)} local markdown files to attach.")

        # 4. Upload Markdown files
        for md_path in md_files:
            rel_display = md_path.relative_to(script_dir.parent.parent)
            print(f"   ⬆️ Uploading: {rel_display}")
            await client.sources.add_file(
                notebook_id=notebook_id,
                file_path=str(md_path)
            )

        print(f"\n🎉 Sync complete! All local docs updated without touching your chat notes.")


if __name__ == "__main__":
    asyncio.run(main())