1#!/usr/bin/env python3
2"""
3Notebook Library Management for NotebookLM
4Manages a library of NotebookLM notebooks with metadata
5Based on the MCP server implementation
6"""
7 
8import json
9import argparse
10import uuid
11import os
12from pathlib import Path
13from typing import Dict, List, Optional, Any
14from datetime import datetime
15 
16 
17class NotebookLibrary:
18    """Manages a collection of NotebookLM notebooks with metadata"""
19 
20    def __init__(self):
21        """Initialize the notebook library"""
22        # Store data within the skill directory
23        skill_dir = Path(__file__).parent.parent
24        self.data_dir = skill_dir / "data"
25        self.data_dir.mkdir(parents=True, exist_ok=True)
26 
27        self.library_file = self.data_dir / "library.json"
28        self.notebooks: Dict[str, Dict[str, Any]] = {}
29        self.active_notebook_id: Optional[str] = None
30 
31        # Load existing library
32        self._load_library()
33 
34    def _load_library(self):
35        """Load library from disk"""
36        if self.library_file.exists():
37            try:
38                with open(self.library_file, 'r') as f:
39                    data = json.load(f)
40                    self.notebooks = data.get('notebooks', {})
41                    self.active_notebook_id = data.get('active_notebook_id')
42                    print(f"📚 Loaded library with {len(self.notebooks)} notebooks")
43            except Exception as e:
44                print(f"⚠️ Error loading library: {e}")
45                self.notebooks = {}
46                self.active_notebook_id = None
47        else:
48            self._save_library()
49 
50    def _save_library(self):
51        """Save library to disk"""
52        try:
53            data = {
54                'notebooks': self.notebooks,
55                'active_notebook_id': self.active_notebook_id,
56                'updated_at': datetime.now().isoformat()
57            }
58            with open(self.library_file, 'w') as f:
59                json.dump(data, f, indent=2)
60        except Exception as e:
61            print(f"❌ Error saving library: {e}")
62 
63    def add_notebook(
64        self,
65        url: str,
66        name: str,
67        description: str,
68        topics: List[str],
69        content_types: Optional[List[str]] = None,
70        use_cases: Optional[List[str]] = None,
71        tags: Optional[List[str]] = None
72    ) -> Dict[str, Any]:
73        """
74        Add a new notebook to the library
75 
76        Args:
77            url: NotebookLM notebook URL
78            name: Display name for the notebook
79            description: What's in this notebook
80            topics: Topics covered
81            content_types: Types of content (optional)
82            use_cases: When to use this notebook (optional)
83            tags: Additional tags for organization (optional)
84 
85        Returns:
86            The created notebook object
87        """
88        # Generate ID from name
89        notebook_id = name.lower().replace(' ', '-').replace('_', '-')
90 
91        # Check for duplicates
92        if notebook_id in self.notebooks:
93            raise ValueError(f"Notebook with ID '{notebook_id}' already exists")
94 
95        # Create notebook object
96        notebook = {
97            'id': notebook_id,
98            'url': url,
99            'name': name,
100            'description': description,
101            'topics': topics,
102            'content_types': content_types or [],
103            'use_cases': use_cases or [],
104            'tags': tags or [],
105            'created_at': datetime.now().isoformat(),
106            'updated_at': datetime.now().isoformat(),
107            'use_count': 0,
108            'last_used': None
109        }
110 
111        # Add to library
112        self.notebooks[notebook_id] = notebook
113 
114        # Set as active if it's the first notebook
115        if len(self.notebooks) == 1:
116            self.active_notebook_id = notebook_id
117 
118        self._save_library()
119 
120        print(f"✅ Added notebook: {name} ({notebook_id})")
121        return notebook
122 
123    def remove_notebook(self, notebook_id: str) -> bool:
124        """
125        Remove a notebook from the library
126 
127        Args:
128            notebook_id: ID of notebook to remove
129 
130        Returns:
131            True if removed, False if not found
132        """
133        if notebook_id in self.notebooks:
134            del self.notebooks[notebook_id]
135 
136            # Clear active if it was removed
137            if self.active_notebook_id == notebook_id:
138                self.active_notebook_id = None
139                # Set new active if there are other notebooks
140                if self.notebooks:
141                    self.active_notebook_id = list(self.notebooks.keys())[0]
142 
143            self._save_library()
144            print(f"✅ Removed notebook: {notebook_id}")
145            return True
146 
147        print(f"⚠️ Notebook not found: {notebook_id}")
148        return False
149 
150    def update_notebook(
151        self,
152        notebook_id: str,
153        name: Optional[str] = None,
154        description: Optional[str] = None,
155        topics: Optional[List[str]] = None,
156        content_types: Optional[List[str]] = None,
157        use_cases: Optional[List[str]] = None,
158        tags: Optional[List[str]] = None,
159        url: Optional[str] = None
160    ) -> Dict[str, Any]:
161        """
162        Update notebook metadata
163 
164        Args:
165            notebook_id: ID of notebook to update
166            Other args: Fields to update (None = keep existing)
167 
168        Returns:
169            Updated notebook object
170        """
171        if notebook_id not in self.notebooks:
172            raise ValueError(f"Notebook not found: {notebook_id}")
173 
174        notebook = self.notebooks[notebook_id]
175 
176        # Update fields if provided
177        if name is not None:
178            notebook['name'] = name
179        if description is not None:
180            notebook['description'] = description
181        if topics is not None:
182            notebook['topics'] = topics
183        if content_types is not None:
184            notebook['content_types'] = content_types
185        if use_cases is not None:
186            notebook['use_cases'] = use_cases
187        if tags is not None:
188            notebook['tags'] = tags
189        if url is not None:
190            notebook['url'] = url
191 
192        notebook['updated_at'] = datetime.now().isoformat()
193 
194        self._save_library()
195        print(f"✅ Updated notebook: {notebook['name']}")
196        return notebook
197 
198    def get_notebook(self, notebook_id: str) -> Optional[Dict[str, Any]]:
199        """Get a specific notebook by ID"""
200        return self.notebooks.get(notebook_id)
201 
202    def list_notebooks(self) -> List[Dict[str, Any]]:
203        """List all notebooks in the library"""
204        return list(self.notebooks.values())
205 
206    def search_notebooks(self, query: str) -> List[Dict[str, Any]]:
207        """
208        Search notebooks by query
209 
210        Args:
211            query: Search query (searches name, description, topics, tags)
212 
213        Returns:
214            List of matching notebooks
215        """
216        query_lower = query.lower()
217        results = []
218 
219        for notebook in self.notebooks.values():
220            # Search in various fields
221            searchable = [
222                notebook['name'].lower(),
223                notebook['description'].lower(),
224                ' '.join(notebook['topics']).lower(),
225                ' '.join(notebook['tags']).lower(),
226                ' '.join(notebook.get('use_cases', [])).lower()
227            ]
228 
229            if any(query_lower in field for field in searchable):
230                results.append(notebook)
231 
232        return results
233 
234    def select_notebook(self, notebook_id: str) -> Dict[str, Any]:
235        """
236        Set a notebook as active
237 
238        Args:
239            notebook_id: ID of notebook to activate
240 
241        Returns:
242            The activated notebook
243        """
244        if notebook_id not in self.notebooks:
245            raise ValueError(f"Notebook not found: {notebook_id}")
246 
247        self.active_notebook_id = notebook_id
248        self._save_library()
249 
250        notebook = self.notebooks[notebook_id]
251        print(f"✅ Activated notebook: {notebook['name']}")
252        return notebook
253 
254    def get_active_notebook(self) -> Optional[Dict[str, Any]]:
255        """Get the currently active notebook"""
256        if self.active_notebook_id:
257            return self.notebooks.get(self.active_notebook_id)
258        return None
259 
260    def increment_use_count(self, notebook_id: str) -> Dict[str, Any]:
261        """
262        Increment usage counter for a notebook
263 
264        Args:
265            notebook_id: ID of notebook that was used
266 
267        Returns:
268            Updated notebook
269        """
270        if notebook_id not in self.notebooks:
271            raise ValueError(f"Notebook not found: {notebook_id}")
272 
273        notebook = self.notebooks[notebook_id]
274        notebook['use_count'] += 1
275        notebook['last_used'] = datetime.now().isoformat()
276 
277        self._save_library()
278        return notebook
279 
280    def get_stats(self) -> Dict[str, Any]:
281        """Get library statistics"""
282        total_notebooks = len(self.notebooks)
283        total_topics = set()
284        total_use_count = 0
285 
286        for notebook in self.notebooks.values():
287            total_topics.update(notebook['topics'])
288            total_use_count += notebook['use_count']
289 
290        # Find most used
291        most_used = None
292        if self.notebooks:
293            most_used = max(
294                self.notebooks.values(),
295                key=lambda n: n['use_count']
296            )
297 
298        return {
299            'total_notebooks': total_notebooks,
300            'total_topics': len(total_topics),
301            'total_use_count': total_use_count,
302            'active_notebook': self.get_active_notebook(),
303            'most_used_notebook': most_used,
304            'library_path': str(self.library_file)
305        }
306 
307 
308def main():
309    """Command-line interface for notebook management"""
310    parser = argparse.ArgumentParser(description='Manage NotebookLM library')
311 
312    subparsers = parser.add_subparsers(dest='command', help='Commands')
313 
314    # Add command
315    add_parser = subparsers.add_parser('add', help='Add a notebook')
316    add_parser.add_argument('--url', required=True, help='NotebookLM URL')
317    add_parser.add_argument('--name', required=True, help='Display name')
318    add_parser.add_argument('--description', required=True, help='Description')
319    add_parser.add_argument('--topics', required=True, help='Comma-separated topics')
320    add_parser.add_argument('--use-cases', help='Comma-separated use cases')
321    add_parser.add_argument('--tags', help='Comma-separated tags')
322 
323    # List command
324    subparsers.add_parser('list', help='List all notebooks')
325 
326    # Search command
327    search_parser = subparsers.add_parser('search', help='Search notebooks')
328    search_parser.add_argument('--query', required=True, help='Search query')
329 
330    # Activate command
331    activate_parser = subparsers.add_parser('activate', help='Set active notebook')
332    activate_parser.add_argument('--id', required=True, help='Notebook ID')
333 
334    # Remove command
335    remove_parser = subparsers.add_parser('remove', help='Remove a notebook')
336    remove_parser.add_argument('--id', required=True, help='Notebook ID')
337 
338    # Stats command
339    subparsers.add_parser('stats', help='Show library statistics')
340 
341    args = parser.parse_args()
342 
343    # Initialize library
344    library = NotebookLibrary()
345 
346    # Execute command
347    if args.command == 'add':
348        topics = [t.strip() for t in args.topics.split(',')]
349        use_cases = [u.strip() for u in args.use_cases.split(',')] if args.use_cases else None
350        tags = [t.strip() for t in args.tags.split(',')] if args.tags else None
351 
352        notebook = library.add_notebook(
353            url=args.url,
354            name=args.name,
355            description=args.description,
356            topics=topics,
357            use_cases=use_cases,
358            tags=tags
359        )
360        print(json.dumps(notebook, indent=2))
361 
362    elif args.command == 'list':
363        notebooks = library.list_notebooks()
364        if notebooks:
365            print("\n📚 Notebook Library:")
366            for notebook in notebooks:
367                active = " [ACTIVE]" if notebook['id'] == library.active_notebook_id else ""
368                print(f"\n  📓 {notebook['name']}{active}")
369                print(f"     ID: {notebook['id']}")
370                print(f"     Topics: {', '.join(notebook['topics'])}")
371                print(f"     Uses: {notebook['use_count']}")
372        else:
373            print("📚 Library is empty. Add notebooks with: notebook_manager.py add")
374 
375    elif args.command == 'search':
376        results = library.search_notebooks(args.query)
377        if results:
378            print(f"\n🔍 Found {len(results)} notebooks:")
379            for notebook in results:
380                print(f"\n  📓 {notebook['name']} ({notebook['id']})")
381                print(f"     {notebook['description']}")
382        else:
383            print(f"🔍 No notebooks found for: {args.query}")
384 
385    elif args.command == 'activate':
386        notebook = library.select_notebook(args.id)
387        print(f"Now using: {notebook['name']}")
388 
389    elif args.command == 'remove':
390        if library.remove_notebook(args.id):
391            print("Notebook removed from library")
392 
393    elif args.command == 'stats':
394        stats = library.get_stats()
395        print("\n📊 Library Statistics:")
396        print(f"  Total notebooks: {stats['total_notebooks']}")
397        print(f"  Total topics: {stats['total_topics']}")
398        print(f"  Total uses: {stats['total_use_count']}")
399        if stats['active_notebook']:
400            print(f"  Active: {stats['active_notebook']['name']}")
401        if stats['most_used_notebook']:
402            print(f"  Most used: {stats['most_used_notebook']['name']} ({stats['most_used_notebook']['use_count']} uses)")
403        print(f"  Library path: {stats['library_path']}")
404 
405    else:
406        parser.print_help()
407 
408 
409if __name__ == "__main__":
410    main()