magicseek__nblm/scripts/source_manager.py

#!/usr/bin/env python3
"""
Unified source ingestion for NotebookLM.
"""

import argparse
import asyncio
import json
import re
import sys
import tempfile
from pathlib import Path
from typing import List, Optional, Union

from agent_browser_client import AgentBrowserClient
from auth_manager import AuthManager
from config import DEFAULT_SESSION_ID
from notebook_manager import NotebookLibrary
from notebooklm_wrapper import NotebookLMWrapper, NotebookLMError
from zlibrary.downloader import ZLibraryDownloader
from zlibrary import epub_converter


def _resolve_notebook_target(args, file_title: str) -> tuple[Optional[str], bool]:
    """Resolve which notebook to use based on CLI args.

    Returns:
        (notebook_id, create_new) tuple:
        - (id, False) = use existing notebook with given ID
        - (None, True) = create new notebook
        - Raises SystemExit if invalid combination
    """
    library = NotebookLibrary()

    # Explicit notebook ID takes priority
    if args.notebook_id:
        return (args.notebook_id, False)

    # --use-active flag
    if args.use_active:
        active = library.get_active_notebook()
        if not active:
            print("❌ No active notebook set.", file=sys.stderr)
            print("   Set one with: python scripts/run.py notebook_manager.py activate --id <id>", file=sys.stderr)
            print("   Or use --create-new to create a new notebook", file=sys.stderr)
            raise SystemExit(1)
        print(f"📓 Using active notebook: \"{active.get('name', 'Unnamed')}\"")
        return (active.get("id"), False)

    # --create-new flag
    if args.create_new:
        print(f"📓 Will create new notebook: \"{file_title}\"")
        return (None, True)

    # Neither flag provided - show error with options
    active = library.get_active_notebook()
    print("❌ No notebook specified. Please choose one of:", file=sys.stderr)
    print(file=sys.stderr)
    if active:
        print(f"  --use-active    Upload to active notebook: \"{active.get('name', 'Unnamed')}\"", file=sys.stderr)
    print(f"  --create-new    Create new notebook named after the file", file=sys.stderr)
    print(f"  --notebook-id   Specify notebook ID explicitly", file=sys.stderr)
    raise SystemExit(1)


class SourceManager:
    """Unified source ingestion for NotebookLM."""

    def __init__(
        self,
        auth_manager: Optional[AuthManager] = None,
        client: Optional[AgentBrowserClient] = None,
        downloader_cls=ZLibraryDownloader,
        converter=epub_converter,
    ):
        self.auth = auth_manager or AuthManager()
        self.client = client or AgentBrowserClient(session_id=DEFAULT_SESSION_ID)
        self.downloader_cls = downloader_cls
        self.converter = converter

    @staticmethod
    def _is_zlibrary_url(url: str) -> bool:
        domains = ["zlib.li", "z-lib.org", "zlibrary.org", "zh.zlib.li"]
        return any(domain in url for domain in domains)

    @staticmethod
    def _sanitize_title(file_path: Path) -> str:
        title = file_path.stem
        title = re.sub(r'_part\d+$', '', title)
        title = title.replace('_', ' ')
        title = re.sub(r'\[.*?\]', '', title)
        title = re.sub(r'\(.*?\)', '', title)
        title = re.sub(r'\s+', ' ', title).strip()
        if len(title) > 50:
            title = title[:50] + "..."
        return title

    @staticmethod
    def _extract_notebook_id_from_url(notebook_url: str) -> Optional[str]:
        if not notebook_url:
            return None
        match = re.search(r"/notebook/([^/?#]+)", notebook_url)
        if match:
            return match.group(1)
        return None

    async def _wait_for_sources_ready_async(
        self,
        wrapper: NotebookLMWrapper,
        notebook_id: str,
        source_ids: List[str],
    ) -> Optional[dict]:
        """Wait for sources to be ready using async wrapper."""
        if not source_ids:
            return None

        unique_ids = list(dict.fromkeys(source_ids))
        total = len(unique_ids)
        print(f"Waiting for NotebookLM to process {total} source(s)...", file=sys.stderr, flush=True)

        last_ready = None
        while True:
            sources = await wrapper.list_sources(notebook_id)
            status_by_id = {src["source_id"]: src for src in sources}

            ready_count = 0
            for source_id in unique_ids:
                source = status_by_id.get(source_id)
                if not source:
                    continue
                if source.get("is_ready"):
                    ready_count += 1

            if last_ready is None or ready_count != last_ready:
                print(f"Ready: {ready_count}/{total}", file=sys.stderr, flush=True)
                last_ready = ready_count

            if ready_count >= total:
                return None

            await asyncio.sleep(2)

    async def add_from_file(
        self,
        file_path: Union[Path, List[Path]],
        notebook_id: Optional[str] = None,
        source_label: str = "upload",
    ) -> dict:
        """Upload local file(s) to NotebookLM."""
        paths = file_path if isinstance(file_path, list) else [file_path]
        for path in paths:
            if not Path(path).exists():
                raise FileNotFoundError(f"File not found: {path}")

        title = self._sanitize_title(Path(paths[0]))
        created_notebook = False

        notebook_url = None
        resolved_notebook_id = notebook_id

        if not self.auth.is_authenticated("google"):
            return {
                "success": False,
                "error": "Google authentication required",
                "recovery": "Run: python scripts/run.py auth_manager.py setup",
            }

        async with NotebookLMWrapper() as wrapper:
            if not notebook_id:
                try:
                    result = await wrapper.create_notebook(title)
                    notebook_id = result["id"]
                    created_notebook = True
                    resolved_notebook_id = notebook_id
                except NotebookLMError as e:
                    return {
                        "success": False,
                        "error": e.message,
                        "recovery": e.recovery,
                    }
            else:
                library = NotebookLibrary()
                notebook = library.get_notebook(notebook_id)
                if notebook:
                    notebook_url = notebook.get("url")
                    resolved_notebook_id = self._extract_notebook_id_from_url(notebook_url) or notebook_id
                else:
                    is_uuid = bool(re.fullmatch(r"[a-f0-9-]{36}", notebook_id or "", re.IGNORECASE))
                    if not is_uuid:
                        return {
                            "success": False,
                            "error": f"Notebook '{notebook_id}' not found in library",
                            "recovery": "Run: python scripts/run.py notebook_manager.py list",
                        }

            if not notebook_url:
                notebook_url = f"https://notebooklm.google.com/notebook/{resolved_notebook_id}"

            # Upload files using the wrapper
            source_ids = []
            try:
                for path in paths:
                    result = await wrapper.add_file(resolved_notebook_id, Path(path))
                    if result.get("source_id"):
                        source_ids.append(result["source_id"])
            except NotebookLMError as e:
                return {
                    "success": False,
                    "error": e.message,
                    "recovery": e.recovery,
                }

            if created_notebook:
                try:
                    library = NotebookLibrary()
                    description = f"Imported from {source_label}: {title}"
                    library.add_notebook(
                        url=notebook_url,
                        name=title,
                        description=description,
                        topics=[source_label],
                    )
                except Exception:
                    pass

            if not source_ids:
                return {"success": False, "error": "Upload failed"}

            # Wait for sources to be ready
            try:
                wait_error = await self._wait_for_sources_ready_async(wrapper, resolved_notebook_id, source_ids)
            except NotebookLMError as e:
                return {
                    "success": False,
                    "error": e.message,
                    "recovery": e.recovery,
                }
            if wait_error:
                return wait_error

        if len(paths) > 1:
            return {
                "success": True,
                "notebook_id": notebook_id,
                "source_ids": source_ids,
                "title": title,
                "chunks": len(paths)
            }

        return {
            "success": True,
            "notebook_id": notebook_id,
            "source_id": source_ids[0],
            "title": title
        }

    async def add_from_zlibrary(self, url: str, notebook_id: Optional[str] = None) -> dict:
        """Download from Z-Library and upload to NotebookLM."""
        if not self.auth.is_authenticated("zlibrary"):
            raise RuntimeError(
                "Z-Library authentication required. "
                "Run: python scripts/run.py auth_manager.py setup --service zlibrary"
            )

        self.client.connect()
        try:
            # Let restore_auth handle navigation - don't navigate here
            self.auth.restore_auth("zlibrary", client=self.client)
            downloader = self.downloader_cls(self.client)
            file_path, file_format = downloader.download(url)
            self.auth.save_auth("zlibrary", client=self.client)
        finally:
            self.client.disconnect()

        if file_format == "epub" or Path(file_path).suffix.lower() == ".epub":
            output_path = Path(tempfile.gettempdir()) / f"{Path(file_path).stem}.md"
            converted = self.converter.convert_epub_to_markdown(file_path, output_path)
            return await self.add_from_file(converted, notebook_id, source_label="zlibrary")

        return await self.add_from_file(Path(file_path), notebook_id, source_label="zlibrary")

    async def add_from_url(self, url: str, notebook_id: Optional[str] = None) -> dict:
        """Smart routing based on URL pattern."""
        if self._is_zlibrary_url(url):
            return await self.add_from_zlibrary(url, notebook_id)
        raise ValueError(f"Unsupported URL: {url}")


async def async_main():
    parser = argparse.ArgumentParser(description="Add sources to NotebookLM")
    parser.add_argument("command", choices=["add"], help="Command to run")
    parser.add_argument("--url", help="Source URL")
    parser.add_argument("--file", help="Local file path")
    parser.add_argument("--notebook-id", help="Existing notebook ID")
    parser.add_argument("--use-active", action="store_true",
                        help="Upload to currently active notebook")
    parser.add_argument("--create-new", action="store_true",
                        help="Create a new notebook for the upload")

    args = parser.parse_args()

    # Validate mutually exclusive options
    if args.use_active and args.create_new:
        print("❌ Cannot use both --use-active and --create-new", file=sys.stderr)
        raise SystemExit(1)
    if args.notebook_id and (args.use_active or args.create_new):
        print("❌ Cannot use --notebook-id with --use-active or --create-new", file=sys.stderr)
        raise SystemExit(1)

    manager = SourceManager()

    if args.command == "add":
        if args.url:
            # For URLs, derive title from URL
            file_title = Path(args.url).stem or "Untitled"
            notebook_id, create_new = _resolve_notebook_target(args, file_title)
            result = await manager.add_from_url(args.url, notebook_id)
        elif args.file:
            file_title = Path(args.file).stem
            notebook_id, create_new = _resolve_notebook_target(args, file_title)
            result = await manager.add_from_file(Path(args.file), notebook_id)
        else:
            raise SystemExit("Provide --url or --file")

        print(json.dumps(result, indent=2))


def main():
    asyncio.run(async_main())


if __name__ == "__main__":
    main()