Update gallery_app.py

gallery_app.py

@@ -7,10 +7,10 @@ from fastapi import Response
 from engine import SorterEngine
 
 # ==========================================
-# STATE MANAGEMENT
+# STATE MANAGEMENT (OPTIMIZED)
 # ==========================================
 class AppState:
-    """Centralized application state with lazy loading."""
+    """Centralized application state with caching and dirty flags."""
     
     def __init__(self):
         # Profile Data
@@ -37,9 +37,23 @@ class AppState:
         
         # Data Caches
         self.all_images: List[str] = []
-        self.staged_data: Dict = {}
+        self._path_to_idx: Dict[str, int] = {}  # Reverse index for O(1) lookups
         self.green_dots: Set[int] = set()
         self.index_map: Dict[int, str] = {}
+        self.committed_indexes: Set[int] = set()  # NEW: Track which indexes are from persistent tags
+        
+        # Staged data cache with dirty flag
+        self._staged_data: Dict = {}
+        self._staged_dirty = True
+        
+        # Disk index cache
+        self._disk_index_cache: Dict[str, Dict[int, str]] = {}
+        self._disk_cache_valid = False
+        
+        # UI dirty flags for granular updates
+        self._dirty_sidebar = True
+        self._dirty_pagination = True
+        self._dirty_gallery = True
         
         # UI Containers (populated later)
         self.sidebar_container = None
@@ -80,30 +94,122 @@ class AppState:
         start = self.page * self.page_size
         return self.all_images[start : start + self.page_size]
 
+    # ==========================================
+    # CACHED PROPERTY: STAGED DATA
+    # ==========================================
+    @property
+    def staged_data(self) -> Dict:
+        """Cached staged data - only fetches from DB when dirty."""
+        if self._staged_dirty:
+            self._staged_data = SorterEngine.get_staged_data()
+            self._staged_dirty = False
+        return self._staged_data
+    
+    def mark_staged_dirty(self):
+        """Mark staged data for refresh on next access."""
+        self._staged_dirty = True
+    
+    def invalidate_disk_cache(self):
+        """Invalidate disk index cache (call after commits)."""
+        self._disk_cache_valid = False
+        self._disk_index_cache.clear()
+
+    # ==========================================
+    # OPTIMIZED INDEX BUILDING
+    # ==========================================
+    def rebuild_path_index(self):
+        """Build reverse path->index mapping for O(1) lookups."""
+        self._path_to_idx = {path: idx for idx, path in enumerate(self.all_images)}
+    
+    def compute_green_dots(self) -> Set[int]:
+        """O(staged) instead of O(all_images)."""
+        dots = set()
+        for path in self.staged_data.keys():
+            if path in self._path_to_idx:
+                dots.add(self._path_to_idx[path] // self.page_size)
+        return dots
+    
+    def get_disk_index_map(self, category: str) -> Dict[int, str]:
+        """Cached disk scan for category indexes."""
+        if not self._disk_cache_valid or category not in self._disk_index_cache:
+            self._rebuild_disk_cache_for_category(category)
+        return self._disk_index_cache.get(category, {})
+    
+    def _rebuild_disk_cache_for_category(self, category: str):
+        """Scan disk for existing files in category folder."""
+        cat_path = os.path.join(self.output_dir, category)
+        index_map = {}
+        if os.path.exists(cat_path):
+            try:
+                with os.scandir(cat_path) as entries:
+                    for entry in entries:
+                        if entry.is_file() and entry.name.startswith(category):
+                            idx = _extract_index(entry.name)
+                            if idx is not None:
+                                index_map[idx] = entry.path
+            except PermissionError:
+                pass
+        self._disk_index_cache[category] = index_map
+        self._disk_cache_valid = True
+
+    # ==========================================
+    # DIRTY FLAG HELPERS
+    # ==========================================
+    def mark_all_dirty(self):
+        self._dirty_sidebar = True
+        self._dirty_pagination = True
+        self._dirty_gallery = True
+    
+    def mark_gallery_dirty(self):
+        self._dirty_gallery = True
+    
+    def mark_sidebar_dirty(self):
+        self._dirty_sidebar = True
+
+
 state = AppState()
 
 # ==========================================
-# IMAGE SERVING API
+# IMAGE SERVING API (OPTIMIZED WITH CACHING)
 # ==========================================
 
 @app.get('/thumbnail')
 async def get_thumbnail(path: str, size: int = 400, q: int = 50):
-    """Serve WebP thumbnail with dynamic quality."""
+    """Serve WebP thumbnail with caching headers."""
     if not os.path.exists(path):
         return Response(status_code=404)
+    
     img_bytes = await run.cpu_bound(SorterEngine.compress_for_web, path, q, size)
-    return Response(content=img_bytes, media_type="image/webp") if img_bytes else Response(status_code=500)
+    
+    if img_bytes:
+        return Response(
+            content=img_bytes, 
+            media_type="image/webp",
+            headers={
+                "Cache-Control": "public, max-age=3600",  # 1 hour browser cache
+                "ETag": f'"{hash(path + str(os.path.getmtime(path)))}"'
+            }
+        )
+    return Response(status_code=500)
 
 @app.get('/full_res')
 async def get_full_res(path: str):
     """Serve full resolution image."""
     if not os.path.exists(path):
         return Response(status_code=404)
+    
     img_bytes = await run.cpu_bound(SorterEngine.compress_for_web, path, 90, None)
-    return Response(content=img_bytes, media_type="image/webp") if img_bytes else Response(status_code=500)
+    
+    if img_bytes:
+        return Response(
+            content=img_bytes, 
+            media_type="image/webp",
+            headers={"Cache-Control": "public, max-age=7200"}  # 2 hour cache for full res
+        )
+    return Response(status_code=500)
 
 # ==========================================
-# CORE LOGIC
+# CORE LOGIC (OPTIMIZED)
 # ==========================================
 
 def load_images():
@@ -113,43 +219,52 @@ def load_images():
         return
     
     state.all_images = SorterEngine.get_images(state.source_dir, recursive=True)
+    state.rebuild_path_index()  # Build reverse index
     
     # Reset page if out of bounds
     if state.page >= state.total_pages:
         state.page = 0
     
+    state.mark_staged_dirty()  # Force refresh of staged data
     refresh_staged_info()
+    state.mark_all_dirty()
     refresh_ui()
 
 def refresh_staged_info():
-    """Update staged data and index maps."""
+    """Update staged data and index maps (optimized) - includes persistent tags."""
-    state.staged_data = SorterEngine.get_staged_data()
+    # Green dots using optimized O(staged) lookup
-    
+    state.green_dots = state.compute_green_dots()
-    # Update green dots (pages with staged images)
-    state.green_dots.clear()
-    staged_keys = set(state.staged_data.keys())
-    for idx, img_path in enumerate(state.all_images):
-        if img_path in staged_keys:
-            state.green_dots.add(idx // state.page_size)
     
     # Build index map for active category
     state.index_map.clear()
+    state.committed_indexes.clear()  # Track which are committed vs staged
     
-    # Add staged images
+    # 1. Add staged images for current category (pending commits) - these are "yellow" 
+    staged_indexes = set()
     for orig_path, info in state.staged_data.items():
         if info['cat'] == state.active_cat:
             idx = _extract_index(info['name'])
             if idx is not None:
                 state.index_map[idx] = orig_path
+                staged_indexes.add(idx)
     
-    # Add committed images from disk
+    # 2. Add committed images from disk (cached scan)
-    cat_path = os.path.join(state.output_dir, state.active_cat)
+    disk_map = state.get_disk_index_map(state.active_cat)
-    if os.path.exists(cat_path):
+    for idx, path in disk_map.items():
-        for filename in os.listdir(cat_path):
+        if idx not in state.index_map:
-            if filename.startswith(state.active_cat):
+            state.index_map[idx] = path
-                idx = _extract_index(filename)
+            state.committed_indexes.add(idx)
-                if idx is not None and idx not in state.index_map:
+    
-                    state.index_map[idx] = os.path.join(cat_path, filename)
+    # 3. Load persistent tags for this output folder (NEW)
+    # This shows which indexes are "taken" even if files moved elsewhere
+    persistent = SorterEngine.get_persistent_tags_by_category(state.output_dir, state.active_cat)
+    for idx, filename in persistent.items():
+        if idx not in state.index_map:
+            # Check if file still exists in output
+            full_path = os.path.join(state.output_dir, state.active_cat, filename)
+            if os.path.exists(full_path):
+                state.index_map[idx] = full_path
+                state.committed_indexes.add(idx)
 
 def _extract_index(filename: str) -> Optional[int]:
     """Extract numeric index from filename (e.g., 'Cat_042.jpg' -> 42)."""
@@ -159,7 +274,7 @@ def _extract_index(filename: str) -> Optional[int]:
         return None
 
 # ==========================================
-# ACTIONS
+# ACTIONS (OPTIMIZED)
 # ==========================================
 
 def action_tag(img_path: str, manual_idx: Optional[int] = None):
@@ -178,17 +293,23 @@ def action_tag(img_path: str, manual_idx: Optional[int] = None):
     
     SorterEngine.stage_image(img_path, state.active_cat, name)
     
-    # Only auto-increment if we used the default next_index (not manual)
+    # Only auto-increment if we used the default next_index
     if manual_idx is None:
         state.next_index = idx + 1
     
+    state.mark_staged_dirty()
     refresh_staged_info()
+    state.mark_sidebar_dirty()
+    state.mark_gallery_dirty()
     refresh_ui()
 
 def action_untag(img_path: str):
     """Remove staging from an image."""
     SorterEngine.clear_staged_item(img_path)
+    state.mark_staged_dirty()
     refresh_staged_info()
+    state.mark_sidebar_dirty()
+    state.mark_gallery_dirty()
     refresh_ui()
 
 def action_delete(img_path: str):
@@ -205,6 +326,9 @@ def action_apply_page():
     
     SorterEngine.commit_batch(batch, state.output_dir, state.cleanup_mode, state.batch_mode)
     ui.notify(f"Page processed ({state.batch_mode})", type='positive')
+    
+    state.invalidate_disk_cache()  # Disk content changed
+    state.mark_staged_dirty()
     load_images()
 
 async def action_apply_global():
@@ -217,11 +341,13 @@ async def action_apply_global():
         state.batch_mode,
         state.source_dir
     )
+    state.invalidate_disk_cache()  # Disk content changed
+    state.mark_staged_dirty()
     load_images()
     ui.notify("Global apply complete!", type='positive')
 
 # ==========================================
-# UI COMPONENTS
+# UI COMPONENTS (OPTIMIZED)
 # ==========================================
 
 def open_zoom_dialog(path: str, title: Optional[str] = None, show_untag: bool = False, show_jump: bool = False):
@@ -232,9 +358,9 @@ def open_zoom_dialog(path: str, title: Optional[str] = None, show_untag: bool =
             
             with ui.row().classes('gap-2'):
                 # Jump to page button
-                if show_jump and path in state.all_images:
+                if show_jump and path in state._path_to_idx:  # O(1) lookup
                     def jump_to_image():
-                        img_idx = state.all_images.index(path)
+                        img_idx = state._path_to_idx[path]  # O(1) instead of list.index()
                         target_page = img_idx // state.page_size
                         dialog.close()
                         set_page(target_page)
@@ -262,21 +388,37 @@ def open_zoom_dialog(path: str, title: Optional[str] = None, show_untag: bool =
 
 def render_sidebar():
     """Render category management sidebar."""
+    if not state._dirty_sidebar:
+        return
+    
     state.sidebar_container.clear()
     
     with state.sidebar_container:
         ui.label("🏷️ Category Manager").classes('text-xl font-bold mb-2 text-white')
         
-        # Number grid (1-25)
+        # Legend for number grid colors
+        with ui.row().classes('gap-4 mb-2 text-xs'):
+            ui.label("🟢 Committed").classes('text-green-400')
+            ui.label("🟡 Staged").classes('text-yellow-400')
+            ui.label("⚫ Free").classes('text-gray-500')
+        
+        # Number grid (1-25) with color coding
         with ui.grid(columns=5).classes('gap-1 mb-4 w-full'):
             for i in range(1, 26):
                 is_used = i in state.index_map
-                color = 'green' if is_used else 'grey-9'
+                is_committed = i in state.committed_indexes
+                
+                # Color logic: green=committed, yellow=staged, grey=free
+                if is_committed:
+                    color = 'green'
+                elif is_used:
+                    color = 'yellow'
+                else:
+                    color = 'grey-9'
                 
                 def make_click_handler(num: int):
                     def handler():
                         if num in state.index_map:
-                            # Number is used - open preview
                             img_path = state.index_map[num]
                             is_staged = img_path in state.staged_data
                             open_zoom_dialog(
@@ -286,8 +428,8 @@ def render_sidebar():
                                 show_jump=True
                             )
                         else:
-                            # Number is free - set as next index
                             state.next_index = num
+                            state._dirty_sidebar = True
                             render_sidebar()
                     return handler
                 
@@ -300,7 +442,9 @@ def render_sidebar():
         
         def on_category_change(e):
             state.active_cat = e.value
+            state.mark_staged_dirty()
             refresh_staged_info()
+            state._dirty_sidebar = True
             render_sidebar()
         
         ui.select(
@@ -319,7 +463,9 @@ def render_sidebar():
                 if new_cat_input.value:
                     SorterEngine.add_category(new_cat_input.value)
                     state.active_cat = new_cat_input.value
+                    state.mark_staged_dirty()
                     refresh_staged_info()
+                    state._dirty_sidebar = True
                     render_sidebar()
             
             ui.button(icon='add', on_click=add_category).props('flat color=green')
@@ -328,7 +474,9 @@ def render_sidebar():
         with ui.expansion('Danger Zone', icon='warning').classes('w-full text-red-400 mt-2'):
             def delete_category():
                 SorterEngine.delete_category(state.active_cat)
+                state.mark_staged_dirty()
                 refresh_staged_info()
+                state._dirty_sidebar = True
                 render_sidebar()
             
             ui.button('DELETE CATEGORY', color='red', on_click=delete_category).classes('w-full')
@@ -343,23 +491,33 @@ def render_sidebar():
             
             def reset_index():
                 state.next_index = (max(state.index_map.keys()) + 1) if state.index_map else 1
+                state._dirty_sidebar = True
                 render_sidebar()
             
             ui.button('🔄', on_click=reset_index).props('flat color=white')
+    
+    state._dirty_sidebar = False
 
 def render_gallery():
     """Render image gallery grid."""
+    if not state._dirty_gallery:
+        return
+    
     state.grid_container.clear()
     batch = state.get_current_batch()
     
+    # Pre-fetch staged data keys for O(1) lookup in loop
+    staged_keys = set(state.staged_data.keys())
+    
     with state.grid_container:
         with ui.grid(columns=state.grid_cols).classes('w-full gap-3'):
             for img_path in batch:
-                render_image_card(img_path)
+                render_image_card(img_path, img_path in staged_keys)
+    
+    state._dirty_gallery = False
 
-def render_image_card(img_path: str):
+def render_image_card(img_path: str, is_staged: bool):
-    """Render individual image card."""
+    """Render individual image card with lazy loading."""
-    is_staged = img_path in state.staged_data
     thumb_size = 800
     
     with ui.card().classes('p-2 bg-gray-900 border border-gray-700 no-shadow'):
@@ -376,10 +534,10 @@ def render_image_card(img_path: str):
                     on_click=lambda p=img_path: action_delete(p)
                 ).props('flat size=sm dense color=red')
         
-        # Thumbnail
+        # Thumbnail with lazy loading
         ui.image(f"/thumbnail?path={img_path}&size={thumb_size}&q={state.preview_quality}") \
             .classes('w-full h-64 bg-black rounded') \
-            .props('fit=contain no-spinner')
+            .props('fit=contain loading=lazy')  # Native lazy loading
         
         # Tagging UI
         if is_staged:
@@ -402,9 +560,13 @@ def render_image_card(img_path: str):
 
 def render_pagination():
     """Render pagination controls."""
+    if not state._dirty_pagination:
+        return
+    
     state.pagination_container.clear()
     
     if state.total_pages <= 1:
+        state._dirty_pagination = False
         return
     
     with state.pagination_container:
@@ -437,17 +599,42 @@ def render_pagination():
             # Next button
             if state.page < state.total_pages - 1:
                 ui.button('▶', on_click=lambda: set_page(state.page + 1)).props('flat color=white')
+    
+    state._dirty_pagination = False
 
 def set_page(p: int):
     """Navigate to specific page."""
     state.page = max(0, min(p, state.total_pages - 1))
+    state._dirty_pagination = True
+    state._dirty_gallery = True
     refresh_ui()
+    
+    # Preload next page in background
+    asyncio.create_task(preload_adjacent_pages())
+
+async def preload_adjacent_pages():
+    """Preload thumbnails for adjacent pages in background."""
+    pages_to_preload = []
+    
+    if state.page < state.total_pages - 1:
+        next_start = (state.page + 1) * state.page_size
+        pages_to_preload.extend(state.all_images[next_start:next_start + state.page_size])
+    
+    if pages_to_preload:
+        await run.cpu_bound(
+            SorterEngine.load_batch_parallel,
+            pages_to_preload,
+            state.preview_quality
+        )
 
 def refresh_ui():
-    """Refresh all UI components."""
+    """Refresh dirty UI components only."""
-    render_sidebar()
+    if state._dirty_sidebar:
-    render_pagination()
+        render_sidebar()
-    render_gallery()
+    if state._dirty_pagination:
+        render_pagination()
+    if state._dirty_gallery:
+        render_gallery()
 
 def handle_keyboard(e):
     """Handle keyboard navigation."""
@@ -501,14 +688,14 @@ def build_header():
                     ui.slider(
                         min=2, max=8, step=1,
                         value=state.grid_cols,
-                        on_change=lambda e: (setattr(state, 'grid_cols', e.value), refresh_ui())
+                        on_change=lambda e: (setattr(state, 'grid_cols', e.value), state.mark_gallery_dirty(), refresh_ui())
                     ).props('color=green')
                     
                     ui.label('Preview Quality:')
                     ui.slider(
                         min=10, max=100, step=10,
                         value=state.preview_quality,
-                        on_change=lambda e: (setattr(state, 'preview_quality', e.value), refresh_ui())
+                        on_change=lambda e: (setattr(state, 'preview_quality', e.value), state.mark_gallery_dirty(), refresh_ui())
                     ).props('color=green label-always')
             
             ui.switch('Dark', value=True, on_change=lambda e: ui.dark_mode().set_value(e.value)) \
@@ -563,6 +750,9 @@ build_main_content()
 
 ui.keyboard(on_key=handle_keyboard)
 ui.dark_mode().enable()
+
+# Initial load with all dirty flags set
+state.mark_all_dirty()
 load_images()
 
 ui.run(title="NiceSorter", host="0.0.0.0", port=8080, reload=False)