1#!/usr/bin/env python3
2"""
3Cleanup Manager for NotebookLM Skill
4Manages cleanup of skill data and browser state
5"""
6 
7import shutil
8import argparse
9from pathlib import Path
10from typing import Dict, List, Any
11 
12 
13class CleanupManager:
14    """
15    Manages cleanup of NotebookLM skill data
16 
17    Features:
18    - Preview what will be deleted
19    - Selective cleanup options
20    - Library preservation
21    - Safe deletion with confirmation
22    """
23 
24    def __init__(self):
25        """Initialize the cleanup manager"""
26        # Skill directory paths
27        self.skill_dir = Path(__file__).parent.parent
28        self.data_dir = self.skill_dir / "data"
29 
30    def get_cleanup_paths(self, preserve_library: bool = False) -> Dict[str, Any]:
31        """
32        Get paths that would be cleaned up
33 
34        Args:
35            preserve_library: Keep library.json if True
36 
37        Returns:
38            Dict with paths and sizes
39 
40        Note: .venv is NEVER deleted - it's part of the skill infrastructure
41        """
42        paths = {
43            'browser_state': [],
44            'sessions': [],
45            'library': [],
46            'auth': [],
47            'other': []
48        }
49 
50        total_size = 0
51 
52        if self.data_dir.exists():
53            # Browser state
54            browser_state_dir = self.data_dir / "browser_state"
55            if browser_state_dir.exists():
56                for item in browser_state_dir.iterdir():
57                    size = self._get_size(item)
58                    paths['browser_state'].append({
59                        'path': str(item),
60                        'size': size,
61                        'type': 'dir' if item.is_dir() else 'file'
62                    })
63                    total_size += size
64 
65            # Sessions
66            sessions_file = self.data_dir / "sessions.json"
67            if sessions_file.exists():
68                size = sessions_file.stat().st_size
69                paths['sessions'].append({
70                    'path': str(sessions_file),
71                    'size': size,
72                    'type': 'file'
73                })
74                total_size += size
75 
76            # Library (unless preserved)
77            if not preserve_library:
78                library_file = self.data_dir / "library.json"
79                if library_file.exists():
80                    size = library_file.stat().st_size
81                    paths['library'].append({
82                        'path': str(library_file),
83                        'size': size,
84                        'type': 'file'
85                    })
86                    total_size += size
87 
88            # Auth info
89            auth_info = self.data_dir / "auth_info.json"
90            if auth_info.exists():
91                size = auth_info.stat().st_size
92                paths['auth'].append({
93                    'path': str(auth_info),
94                    'size': size,
95                    'type': 'file'
96                })
97                total_size += size
98 
99            # Other files in data dir (but NEVER .venv!)
100            for item in self.data_dir.iterdir():
101                if item.name not in ['browser_state', 'sessions.json', 'library.json', 'auth_info.json']:
102                    size = self._get_size(item)
103                    paths['other'].append({
104                        'path': str(item),
105                        'size': size,
106                        'type': 'dir' if item.is_dir() else 'file'
107                    })
108                    total_size += size
109 
110        return {
111            'categories': paths,
112            'total_size': total_size,
113            'total_items': sum(len(items) for items in paths.values())
114        }
115 
116    def _get_size(self, path: Path) -> int:
117        """Get size of file or directory in bytes"""
118        if path.is_file():
119            return path.stat().st_size
120        elif path.is_dir():
121            total = 0
122            try:
123                for item in path.rglob('*'):
124                    if item.is_file():
125                        total += item.stat().st_size
126            except Exception:
127                pass
128            return total
129        return 0
130 
131    def _format_size(self, size: int) -> str:
132        """Format size in human-readable form"""
133        for unit in ['B', 'KB', 'MB', 'GB']:
134            if size < 1024:
135                return f"{size:.1f} {unit}"
136            size /= 1024
137        return f"{size:.1f} TB"
138 
139    def perform_cleanup(
140        self,
141        preserve_library: bool = False,
142        dry_run: bool = False
143    ) -> Dict[str, Any]:
144        """
145        Perform the actual cleanup
146 
147        Args:
148            preserve_library: Keep library.json if True
149            dry_run: Preview only, don't delete
150 
151        Returns:
152            Dict with cleanup results
153        """
154        cleanup_data = self.get_cleanup_paths(preserve_library)
155        deleted_items = []
156        failed_items = []
157        deleted_size = 0
158 
159        if dry_run:
160            return {
161                'dry_run': True,
162                'would_delete': cleanup_data['total_items'],
163                'would_free': cleanup_data['total_size']
164            }
165 
166        # Perform deletion
167        for category, items in cleanup_data['categories'].items():
168            for item_info in items:
169                path = Path(item_info['path'])
170                try:
171                    if path.exists():
172                        if path.is_dir():
173                            shutil.rmtree(path)
174                        else:
175                            path.unlink()
176                        deleted_items.append(str(path))
177                        deleted_size += item_info['size']
178                        print(f"  ✅ Deleted: {path.name}")
179                except Exception as e:
180                    failed_items.append({
181                        'path': str(path),
182                        'error': str(e)
183                    })
184                    print(f"  ❌ Failed: {path.name} ({e})")
185 
186        # Recreate browser_state dir if everything was deleted
187        if not preserve_library and not failed_items:
188            browser_state_dir = self.data_dir / "browser_state"
189            browser_state_dir.mkdir(parents=True, exist_ok=True)
190 
191        return {
192            'deleted_items': deleted_items,
193            'failed_items': failed_items,
194            'deleted_size': deleted_size,
195            'deleted_count': len(deleted_items),
196            'failed_count': len(failed_items)
197        }
198 
199    def print_cleanup_preview(self, preserve_library: bool = False):
200        """Print a preview of what will be cleaned"""
201        data = self.get_cleanup_paths(preserve_library)
202 
203        print("\n🔍 Cleanup Preview")
204        print("=" * 60)
205 
206        for category, items in data['categories'].items():
207            if items:
208                print(f"\n📁 {category.replace('_', ' ').title()}:")
209                for item in items:
210                    path = Path(item['path'])
211                    size_str = self._format_size(item['size'])
212                    type_icon = "📂" if item['type'] == 'dir' else "📄"
213                    print(f"  {type_icon} {path.name:<30} {size_str:>10}")
214 
215        print("\n" + "=" * 60)
216        print(f"Total items: {data['total_items']}")
217        print(f"Total size: {self._format_size(data['total_size'])}")
218 
219        if preserve_library:
220            print("\n📚 Library will be preserved")
221 
222        print("\nThis preview shows what would be deleted.")
223        print("Use --confirm to actually perform the cleanup.")
224 
225 
226def main():
227    """Command-line interface for cleanup management"""
228    parser = argparse.ArgumentParser(
229        description='Clean up NotebookLM skill data',
230        formatter_class=argparse.RawDescriptionHelpFormatter,
231        epilog="""
232Examples:
233  # Preview what will be deleted
234  python cleanup_manager.py
235 
236  # Perform cleanup (delete everything)
237  python cleanup_manager.py --confirm
238 
239  # Cleanup but keep library
240  python cleanup_manager.py --confirm --preserve-library
241 
242  # Force cleanup without preview
243  python cleanup_manager.py --confirm --force
244        """
245    )
246 
247    parser.add_argument(
248        '--confirm',
249        action='store_true',
250        help='Actually perform the cleanup (without this, only preview)'
251    )
252 
253    parser.add_argument(
254        '--preserve-library',
255        action='store_true',
256        help='Keep the notebook library (library.json)'
257    )
258 
259    parser.add_argument(
260        '--force',
261        action='store_true',
262        help='Skip confirmation prompt'
263    )
264 
265    args = parser.parse_args()
266 
267    # Initialize manager
268    manager = CleanupManager()
269 
270    if args.confirm:
271        # Show preview first unless forced
272        if not args.force:
273            manager.print_cleanup_preview(args.preserve_library)
274 
275            print("\n⚠️  WARNING: This will delete the files shown above!")
276            print("   Note: .venv is preserved (part of skill infrastructure)")
277            response = input("Are you sure? (yes/no): ")
278 
279            if response.lower() != 'yes':
280                print("Cleanup cancelled.")
281                return
282 
283        # Perform cleanup
284        print("\n🗑️ Performing cleanup...")
285        result = manager.perform_cleanup(args.preserve_library, dry_run=False)
286 
287        print(f"\n✅ Cleanup complete!")
288        print(f"  Deleted: {result['deleted_count']} items")
289        print(f"  Freed: {manager._format_size(result['deleted_size'])}")
290 
291        if result['failed_count'] > 0:
292            print(f"  ⚠️ Failed: {result['failed_count']} items")
293 
294    else:
295        # Just show preview
296        manager.print_cleanup_preview(args.preserve_library)
297        print("\n💡 Note: Virtual environment (.venv) is never deleted")
298        print("   It's part of the skill infrastructure, not user data")
299 
300 
301if __name__ == "__main__":
302    main()