Update engine.py

engine.py

@@ -48,7 +48,7 @@ class SorterEngine:
     
     @staticmethod
     def init_db():
-        """Initializes tables, including the HISTORY log."""
+        """Initializes tables, including the HISTORY log and persistent tags."""
         with SorterEngine.get_connection() as conn:
             cursor = conn.cursor()
             
@@ -62,6 +62,21 @@ class SorterEngine:
             cursor.execute('''CREATE TABLE IF NOT EXISTS processed_log 
                 (source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''')
             
+            # NEW: Persistent tags table - survives after commit
+            # Maps: output_folder + filename -> original source info for tag restoration
+            cursor.execute('''CREATE TABLE IF NOT EXISTS persistent_tags 
+                (output_folder TEXT,
+                 filename TEXT,
+                 category TEXT,
+                 tag_index INTEGER,
+                 original_source_folder TEXT,
+                 committed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                 PRIMARY KEY (output_folder, filename))''')
+            
+            # Index for fast lookup by output folder
+            cursor.execute('''CREATE INDEX IF NOT EXISTS idx_persistent_tags_folder 
+                ON persistent_tags(output_folder)''')
+            
             # Seed categories if empty
             cursor.execute("SELECT COUNT(*) FROM categories")
             if cursor.fetchone()[0] == 0:
@@ -388,6 +403,88 @@ class SorterEngine:
             rows = cursor.fetchall()
             return {r[0]: {"cat": r[1], "name": r[2], "marked": r[3]} for r in rows}
 
+    # ==========================================
+    # PERSISTENT TAGS (NEW)
+    # ==========================================
+    
+    @staticmethod
+    def save_persistent_tag(output_folder: str, filename: str, category: str, 
+                            tag_index: int, source_folder: str):
+        """Save a tag permanently after commit."""
+        with SorterEngine.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                INSERT OR REPLACE INTO persistent_tags 
+                (output_folder, filename, category, tag_index, original_source_folder)
+                VALUES (?, ?, ?, ?, ?)
+            """, (output_folder, filename, category, tag_index, source_folder))
+            conn.commit()
+    
+    @staticmethod
+    def save_persistent_tags_batch(tags: List[Tuple[str, str, str, int, str]]):
+        """Batch save multiple persistent tags."""
+        with SorterEngine.transaction() as conn:
+            cursor = conn.cursor()
+            cursor.executemany("""
+                INSERT OR REPLACE INTO persistent_tags 
+                (output_folder, filename, category, tag_index, original_source_folder)
+                VALUES (?, ?, ?, ?, ?)
+            """, tags)
+    
+    @staticmethod
+    def get_persistent_tags(output_folder: str) -> Dict[str, Dict]:
+        """
+        Get all persistent tags for an output folder.
+        Returns: {filename: {"cat": category, "index": tag_index, "source": source_folder}}
+        """
+        with SorterEngine.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT filename, category, tag_index, original_source_folder
+                FROM persistent_tags
+                WHERE output_folder = ?
+            """, (output_folder,))
+            rows = cursor.fetchall()
+            return {
+                r[0]: {"cat": r[1], "index": r[2], "source": r[3]}
+                for r in rows
+            }
+    
+    @staticmethod
+    def get_persistent_tags_by_category(output_folder: str, category: str) -> Dict[int, str]:
+        """
+        Get persistent tags for a specific category in an output folder.
+        Returns: {tag_index: filename}
+        """
+        with SorterEngine.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                SELECT tag_index, filename
+                FROM persistent_tags
+                WHERE output_folder = ? AND category = ?
+            """, (output_folder, category))
+            rows = cursor.fetchall()
+            return {r[0]: r[1] for r in rows}
+    
+    @staticmethod
+    def delete_persistent_tag(output_folder: str, filename: str):
+        """Remove a persistent tag (e.g., if file is deleted)."""
+        with SorterEngine.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("""
+                DELETE FROM persistent_tags
+                WHERE output_folder = ? AND filename = ?
+            """, (output_folder, filename))
+            conn.commit()
+    
+    @staticmethod
+    def clear_persistent_tags(output_folder: str):
+        """Clear all persistent tags for an output folder."""
+        with SorterEngine.get_connection() as conn:
+            cursor = conn.cursor()
+            cursor.execute("DELETE FROM persistent_tags WHERE output_folder = ?", (output_folder,))
+            conn.commit()
+
     # ==========================================
     # BATCH COMMIT OPERATIONS (OPTIMIZED)
     # ==========================================
@@ -416,7 +513,7 @@ class SorterEngine:
 
     @staticmethod
     def commit_batch(file_list: List[str], output_root: str, cleanup_mode: str, operation: str = "Copy"):
-        """Commits files with batched DB operations."""
+        """Commits files with batched DB operations and saves persistent tags."""
         data = SorterEngine.get_staged_data()
         
         if not os.path.exists(output_root):
@@ -425,6 +522,7 @@ class SorterEngine:
         # Prepare batch operations
         to_delete_from_staging = []
         to_insert_to_log = []
+        persistent_tags_to_save = []
         
         for file_path in file_list:
             if not os.path.exists(file_path):
@@ -434,6 +532,14 @@ class SorterEngine:
             if file_path in data and data[file_path]['marked']:
                 info = data[file_path]
                 final_dst = SorterEngine._compute_final_destination(output_root, info['name'])
+                final_filename = os.path.basename(final_dst)
+                
+                # Extract tag index from filename
+                tag_index = None
+                try:
+                    tag_index = int(info['name'].rsplit('_', 1)[1].split('.')[0])
+                except (ValueError, IndexError):
+                    pass
                 
                 # Perform file operation
                 if operation == "Copy":
@@ -445,6 +551,16 @@ class SorterEngine:
                 
                 to_delete_from_staging.append((file_path,))
                 to_insert_to_log.append((file_path, info['cat'], operation))
+                
+                # Save persistent tag
+                if tag_index is not None:
+                    persistent_tags_to_save.append((
+                        output_root,           # output_folder
+                        final_filename,        # filename
+                        info['cat'],           # category
+                        tag_index,             # tag_index
+                        os.path.dirname(file_path)  # original source folder
+                    ))
             
             # Untagged files - cleanup
             elif cleanup_mode != "Keep":
@@ -464,16 +580,23 @@ class SorterEngine:
                 cursor.executemany("DELETE FROM staging_area WHERE original_path = ?", to_delete_from_staging)
             if to_insert_to_log:
                 cursor.executemany("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", to_insert_to_log)
+            if persistent_tags_to_save:
+                cursor.executemany("""
+                    INSERT OR REPLACE INTO persistent_tags 
+                    (output_folder, filename, category, tag_index, original_source_folder)
+                    VALUES (?, ?, ?, ?, ?)
+                """, persistent_tags_to_save)
 
     @staticmethod
     def commit_global(output_root: str, cleanup_mode: str, operation: str = "Copy", source_root: str = None):
-        """Commits ALL staged files with batched operations."""
+        """Commits ALL staged files with batched operations and saves persistent tags."""
         data = SorterEngine.get_staged_data()
         
         if not os.path.exists(output_root):
             os.makedirs(output_root, exist_ok=True)
         
         to_insert_to_log = []
+        persistent_tags_to_save = []
         
         # Process all staged items
         for old_p, info in data.items():
@@ -481,6 +604,14 @@ class SorterEngine:
                 continue
                 
             final_dst = SorterEngine._compute_final_destination(output_root, info['name'])
+            final_filename = os.path.basename(final_dst)
+            
+            # Extract tag index from filename
+            tag_index = None
+            try:
+                tag_index = int(info['name'].rsplit('_', 1)[1].split('.')[0])
+            except (ValueError, IndexError):
+                pass
             
             if operation == "Copy":
                 shutil.copy2(old_p, final_dst)
@@ -489,6 +620,16 @@ class SorterEngine:
             
             SorterEngine.fix_permissions(final_dst)
             to_insert_to_log.append((old_p, info['cat'], operation))
+            
+            # Save persistent tag
+            if tag_index is not None:
+                persistent_tags_to_save.append((
+                    output_root,           # output_folder
+                    final_filename,        # filename  
+                    info['cat'],           # category
+                    tag_index,             # tag_index
+                    os.path.dirname(old_p) # original source folder
+                ))
         
         # Global cleanup
         if cleanup_mode != "Keep" and source_root:
@@ -510,6 +651,12 @@ class SorterEngine:
             cursor.execute("DELETE FROM staging_area")
             if to_insert_to_log:
                 cursor.executemany("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", to_insert_to_log)
+            if persistent_tags_to_save:
+                cursor.executemany("""
+                    INSERT OR REPLACE INTO persistent_tags 
+                    (output_folder, filename, category, tag_index, original_source_folder)
+                    VALUES (?, ?, ?, ?, ?)
+                """, persistent_tags_to_save)
 
     # ==========================================
     # UTILITY OPERATIONS