docsifter.py

#!/usr/bin/env python3
"""
Docsifter - Document organization and tagging application
"""

import tkinter as tk
from tkinter import ttk, messagebox, scrolledtext
import os
import pathlib
from typing import List, Dict
from dotenv import load_dotenv
import json
import threading
from datetime import datetime

try:
    import ollama
    OLLAMA_AVAILABLE = True
except ImportError:
    OLLAMA_AVAILABLE = False


class DocSifterApp:
    """Main application for document sifting and tagging"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("DocSifter - Document Tagger")
        self.root.geometry("800x700")
        
        # Tags database file
        self.tags_db_file = os.path.join(os.path.dirname(__file__), "tags_database.json")
        self.tags_db = self.load_tags_database()
        
        # Load environment variables
        load_dotenv()
        
        # Initialize local LLM
        self.client = None
        self.ollama_model = os.getenv("OLLAMA_MODEL", "llama3.2")
        
        if OLLAMA_AVAILABLE:
            try:
                ollama.list()
                self.client = ollama
            except Exception as e:
                print(f"Ollama not available: {e}")
        
        # Get default folders
        home = pathlib.Path.home()
        self.default_folders = {
            "Documents": str(home / "Documents"),
            "Desktop": str(home / "Desktop"),
            "Downloads": str(home / "Downloads")
        }
        
        self.setup_ui()
    
    def load_tags_database(self) -> Dict:
        """Load tags database from JSON file"""
        if os.path.exists(self.tags_db_file):
            try:
                with open(self.tags_db_file, 'r') as f:
                    return json.load(f)
            except Exception as e:
                print(f"Error loading tags database: {e}")
                return {}
        return {}
    
    def save_tags_database(self):
        """Save tags database to JSON file"""
        try:
            with open(self.tags_db_file, 'w') as f:
                json.dump(self.tags_db, f, indent=2)
        except Exception as e:
            print(f"Error saving tags database: {e}")
        
    def setup_ui(self):
        """Set up the user interface"""
        # Title
        title_label = tk.Label(
            self.root, 
            text="DocSifter - AI Document Tagger",
            font=("Helvetica", 16, "bold")
        )
        title_label.pack(pady=10)
        
        # Folder selection frame
        folder_frame = tk.Frame(self.root)
        folder_frame.pack(pady=10, padx=20, fill=tk.X)
        
        tk.Label(folder_frame, text="Select Folder:", font=("Helvetica", 12)).pack(side=tk.LEFT, padx=5)
        
        self.folder_var = tk.StringVar()
        self.folder_dropdown = ttk.Combobox(
            folder_frame,
            textvariable=self.folder_var,
            values=list(self.default_folders.keys()),
            state="readonly",
            width=30
        )
        self.folder_dropdown.pack(side=tk.LEFT, padx=5)
        self.folder_dropdown.current(0)  # Select Documents by default
        
        # Scan button
        scan_button = tk.Button(
            folder_frame,
            text="Scan & Tag Files",
            command=self.start_scan,
            bg="#4CAF50",
            fg="white",
            font=("Helvetica", 10, "bold"),
            padx=15,
            pady=5
        )
        scan_button.pack(side=tk.LEFT, padx=10)
        
        # Search frame
        search_frame = tk.Frame(self.root)
        search_frame.pack(pady=10, padx=20, fill=tk.X)
        
        tk.Label(search_frame, text="Search:", font=("Helvetica", 12)).pack(side=tk.LEFT, padx=5)
        
        self.search_var = tk.StringVar()
        self.search_var.trace('w', lambda *args: self.perform_search())
        search_entry = tk.Entry(
            search_frame,
            textvariable=self.search_var,
            font=("Helvetica", 11),
            width=40
        )
        search_entry.pack(side=tk.LEFT, padx=5)
        
        # Clear search button
        clear_button = tk.Button(
            search_frame,
            text="Clear",
            command=self.clear_search,
            font=("Helvetica", 10),
            padx=10
        )
        clear_button.pack(side=tk.LEFT, padx=5)
        
        # Status label
        self.status_label = tk.Label(
            self.root,
            text="Ready to scan",
            font=("Helvetica", 10),
            fg="gray"
        )
        self.status_label.pack(pady=5)
        
        # Progress bar
        self.progress = ttk.Progressbar(
            self.root,
            mode='indeterminate',
            length=760
        )
        self.progress.pack(pady=5, padx=20)
        
        # Results area
        results_label = tk.Label(
            self.root,
            text="Results:",
            font=("Helvetica", 12, "bold")
        )
        results_label.pack(pady=5, anchor=tk.W, padx=20)
        
        self.results_text = scrolledtext.ScrolledText(
            self.root,
            width=90,
            height=22,
            wrap=tk.WORD,
            font=("Courier", 9)
        )
        self.results_text.pack(pady=5, padx=20, fill=tk.BOTH, expand=True)
        
        # Info label
        if not self.client:
            info_text = "⚠️  Ollama not running. Using heuristic mode."
        else:
            info_text = f"✓ Local LLM ready ({self.ollama_model})"
        
        info_label = tk.Label(
            self.root,
            text=info_text,
            font=("Helvetica", 9),
            fg="orange" if not self.client else "green"
        )
        info_label.pack(pady=5)
        
    def clear_search(self):
        """Clear the search field and show all files"""
        self.search_var.set("")
        self.perform_search()
    
    def perform_search(self):
        """Search for files by filename or tags"""
        search_query = self.search_var.get().lower().strip()
        
        if not search_query:
            # If search is empty, show message
            if not self.tags_db:
                self.results_text.delete(1.0, tk.END)
                self.results_text.insert(tk.END, "No files tagged yet. Scan a folder to start tagging.\n")
            return
        
        # Clear results
        self.results_text.delete(1.0, tk.END)
        
        # Search through database
        results = []
        for file_path, data in self.tags_db.items():
            file_name = os.path.basename(file_path)
            tags = data.get('tags', [])
            
            # Check if query matches filename or any tag
            if (search_query in file_name.lower() or 
                any(search_query in tag.lower() for tag in tags)):
                results.append((file_path, data))
        
        # Display results
        if results:
            self.log_result(f"Search results for '{search_query}' ({len(results)} files found):\n")
            self.log_result("=" * 70 + "\n\n")
            
            for i, (file_path, data) in enumerate(results, 1):
                file_name = os.path.basename(file_path)
                tags = data.get('tags', [])
                tagged_date = data.get('tagged_date', 'Unknown')
                
                self.log_result(f"[{i}] {file_name}\n")
                self.log_result(f"    Path: {file_path}\n")
                self.log_result(f"    Tags: {', '.join(tags)}\n")
                self.log_result(f"    Tagged: {tagged_date}\n")
                self.log_result("\n")
        else:
            self.log_result(f"No files found matching '{search_query}'\n")
            self.log_result("\nTry searching by:\n")
            self.log_result("  - Filename (e.g., 'report', 'invoice')\n")
            self.log_result("  - Tag (e.g., 'document', 'image', 'pdf')\n")
    
    def start_scan(self):
        """Start the scanning process in a background thread"""
        selected_folder_name = self.folder_var.get()
        if not selected_folder_name:
            messagebox.showwarning("No Selection", "Please select a folder first.")
            return
        
        folder_path = self.default_folders[selected_folder_name]
        
        if not os.path.exists(folder_path):
            messagebox.showerror(
                "Folder Not Found",
                f"The folder '{folder_path}' does not exist."
            )
            return
        
        # Clear previous results
        self.results_text.delete(1.0, tk.END)
        self.status_label.config(text=f"Scanning {selected_folder_name}...", fg="blue")
        self.progress.start(10)
        
        # Run scanning in background thread
        thread = threading.Thread(
            target=self.scan_and_tag_files,
            args=(folder_path, selected_folder_name),
            daemon=True
        )
        thread.start()
        
    def scan_and_tag_files(self, folder_path: str, folder_name: str):
        """Scan files in the folder and tag them"""
        try:
            files = []
            for entry in os.scandir(folder_path):
                if entry.is_file() and not entry.name.startswith('.'):
                    files.append(entry.path)
            
            if not files:
                self.update_status("No files found in the selected folder.", "orange")
                self.progress.stop()
                return
            
            self.log_result(f"Found {len(files)} files in {folder_name}\n")
            self.log_result("=" * 70 + "\n\n")
            
            for i, file_path in enumerate(files[:20], 1):  # Limit to first 20 files
                file_name = os.path.basename(file_path)
                self.log_result(f"[{i}] Processing: {file_name}\n")
                
                # Generate tags
                tags = self.generate_tags(file_path)
                
                # Display tags
                tags_str = ", ".join(tags)
                self.log_result(f"    Tags: {tags_str}\n")
                
                # Save tags to database
                self.save_file_tags(file_path, tags)
                self.log_result(f"    ✓ Tags saved to database\n")
                
                self.log_result("\n")
            
            if len(files) > 20:
                self.log_result(f"\n(Showing first 20 of {len(files)} files)\n")
            
            # Save database to file
            self.save_tags_database()
            
            self.update_status(f"✓ Completed scanning {len(files[:20])} files", "green")
            
        except Exception as e:
            self.update_status(f"Error: {str(e)}", "red")
            self.log_result(f"\n\n❌ Error: {str(e)}\n")
        finally:
            self.progress.stop()
    
    def is_text_based_file(self, file_ext: str) -> bool:
        """Check if file is text-based (not binary like images, videos, audio)"""
        text_extensions = {
            '.pdf', '.doc', '.docx', '.txt', '.rtf', '.odt',
            '.xlsx', '.xls', '.csv', '.ods',
            '.pptx', '.ppt', '.odp',
            '.html', '.htm', '.xml', '.json', '.yaml', '.yml',
            '.md', '.markdown', '.rst',
            '.py', '.js', '.java', '.c', '.cpp', '.h', '.cs',
            '.go', '.rs', '.rb', '.php', '.swift', '.kt',
            '.css', '.scss', '.sass', '.less',
            '.sql', '.sh', '.bash', '.ps1',
            '.log', '.cfg', '.conf', '.ini', '.env'
        }
        return file_ext in text_extensions
    
    def generate_tags(self, file_path: str) -> List[str]:
        """Generate tags for a file using AI or fallback heuristics"""
        file_name = os.path.basename(file_path)
        file_ext = pathlib.Path(file_path).suffix.lower()
        
        # For binary files (images, videos, audio), only use basic heuristic tagging
        # Don't use AI as it can't read the content and would just guess
        if not self.is_text_based_file(file_ext):
            return self.generate_tags_heuristic(file_name, file_ext)[:10]
        
        # Try AI-based tagging for text-based documents
        if self.client:
            try:
                tags = self.generate_tags_ai(file_path, file_name, file_ext)
                if tags:
                    return tags[:10]  # Max 10 tags
            except Exception as e:
                print(f"AI tagging failed: {e}, falling back to heuristic")
        
        # Fallback to heuristic-based tagging
        return self.generate_tags_heuristic(file_name, file_ext)[:10]
    
    def generate_tags_ai(self, file_path: str, file_name: str, file_ext: str) -> List[str]:
        """Generate tags using local Ollama LLM"""
        stats = os.stat(file_path)
        file_size = stats.st_size
        
        prompt = f"""Given a file with these properties:
- Filename: {file_name}
- Extension: {file_ext}
- Size: {file_size} bytes

Generate tags based ONLY on the filename and extension. Do NOT guess or hallucinate content you cannot see.

Rules:
- If the filename clearly indicates document type (e.g., "invoice", "resume", "I797"), use that
- For financial/immigration docs with specific patterns, include appropriate tags
- If filename is generic (e.g., "document.pdf", "file.txt"), return basic tags only
- Do NOT invent or guess specific details not evident in the filename
- Maximum 10 tags, minimum 2 tags

For specific document types:
- Immigration: I797, I20, visa, passport, greencard, EAD (only if in filename)
- Financial: bank statement, tax return, W2, 1099, invoice, receipt (only if in filename)
- General: document type from extension, year if present in filename

Return ONLY a JSON array of strings.
Example for "invoice_jan2024.pdf": ["document", "pdf", "financial", "invoice", "2024"]
Example for "document.pdf": ["document", "pdf"]"""

        response = self.client.chat(
            model=self.ollama_model,
            messages=[
                {"role": "system", "content": "You are a conservative file tagger. Only tag based on clear evidence in filename and extension. Never guess or hallucinate. Return only JSON arrays."},
                {"role": "user", "content": prompt}
            ],
            options={"temperature": 0.3, "num_predict": 150}
        )
        
        tags_json = response['message']['content'].strip()
        if '```' in tags_json:
            import re
            json_match = re.search(r'\[.*\]', tags_json, re.DOTALL)
            if json_match:
                tags_json = json_match.group(0)
        
        try:
            tags = json.loads(tags_json)
            # Validate tags are reasonable
            if not isinstance(tags, list) or len(tags) == 0:
                return []
            # Filter out any non-string tags
            tags = [str(tag) for tag in tags if isinstance(tag, str) and tag.strip()]
            return tags
        except (json.JSONDecodeError, ValueError):
            print(f"Failed to parse AI response: {tags_json}")
            return []
    
    def generate_tags_heuristic(self, file_name: str, file_ext: str) -> List[str]:
        """Generate tags using heuristic rules (fallback)"""
        tags = []
        
        # Extension-based tags
        ext_map = {
            '.pdf': ['document', 'pdf'],
            '.doc': ['document', 'word'],
            '.docx': ['document', 'word'],
            '.txt': ['text', 'document'],
            '.jpg': ['image', 'photo'],
            '.jpeg': ['image', 'photo'],
            '.png': ['image', 'graphic'],
            '.gif': ['image', 'animation'],
            '.mp4': ['video', 'media'],
            '.mov': ['video', 'media'],
            '.mp3': ['audio', 'music'],
            '.wav': ['audio', 'sound'],
            '.xlsx': ['spreadsheet', 'excel'],
            '.xls': ['spreadsheet', 'excel'],
            '.csv': ['data', 'spreadsheet'],
            '.pptx': ['presentation', 'powerpoint'],
            '.ppt': ['presentation', 'powerpoint'],
            '.zip': ['archive', 'compressed'],
            '.rar': ['archive', 'compressed'],
            '.py': ['code', 'python'],
            '.js': ['code', 'javascript'],
            '.html': ['code', 'web'],
            '.css': ['code', 'web'],
        }
        
        if file_ext in ext_map:
            tags.extend(ext_map[file_ext])
        else:
            tags.append('file')
        
        # Name-based tags
        name_lower = file_name.lower()
        
        # Immigration documents
        if 'i797' in name_lower or 'i-797' in name_lower:
            tags.extend(['immigration', 'i797', 'uscis', 'approval'])
        if 'i20' in name_lower or 'i-20' in name_lower:
            tags.extend(['immigration', 'i20', 'student', 'visa'])
        if 'visa' in name_lower:
            tags.extend(['immigration', 'visa', 'travel'])
        if 'passport' in name_lower:
            tags.extend(['immigration', 'passport', 'travel', 'id'])
        if 'ead' in name_lower or 'work permit' in name_lower:
            tags.extend(['immigration', 'ead', 'work-permit'])
        if 'green' in name_lower and 'card' in name_lower:
            tags.extend(['immigration', 'greencard', 'permanent-resident'])
        
        # Financial documents
        if 'bank' in name_lower and 'statement' in name_lower:
            tags.extend(['financial', 'bank-statement', 'banking'])
        if 'tax' in name_lower:
            tags.extend(['financial', 'tax', 'irs'])
        if 'w2' in name_lower or 'w-2' in name_lower:
            tags.extend(['financial', 'w2', 'tax', 'income'])
        if '1099' in name_lower:
            tags.extend(['financial', '1099', 'tax', 'income'])
        if 'invoice' in name_lower:
            tags.extend(['financial', 'invoice', 'billing'])
        if 'receipt' in name_lower:
            tags.extend(['financial', 'receipt', 'expense'])
        if 'payslip' in name_lower or 'paystub' in name_lower:
            tags.extend(['financial', 'payslip', 'income', 'salary'])
        if 'credit' in name_lower and ('card' in name_lower or 'report' in name_lower):
            tags.extend(['financial', 'credit'])
        
        # General documents
        if 'resume' in name_lower or 'cv' in name_lower:
            tags.extend(['resume', 'career', 'professional'])
        if 'report' in name_lower:
            tags.append('report')
        if 'screenshot' in name_lower:
            tags.append('screenshot')
        if 'contract' in name_lower or 'agreement' in name_lower:
            tags.extend(['legal', 'contract'])
        if 'lease' in name_lower or 'rental' in name_lower:
            tags.extend(['legal', 'lease', 'housing'])
        
        # Year detection
        for year in ['2025', '2024', '2023', '2022', '2021']:
            if year in name_lower:
                tags.append(year)
                if year in ['2024', '2025']:
                    tags.append('recent')
                break
        
        # Remove duplicates while preserving order
        seen = set()
        unique_tags = []
        for tag in tags:
            if tag not in seen:
                seen.add(tag)
                unique_tags.append(tag)
        
        return unique_tags[:10]
    
    def save_file_tags(self, file_path: str, tags: List[str]):
        """Save tags for a file to the database"""
        self.tags_db[file_path] = {
            'tags': tags,
            'tagged_date': datetime.now().isoformat(),
            'filename': os.path.basename(file_path)
        }
    
    def log_result(self, message: str):
        """Log a message to the results area"""
        self.results_text.insert(tk.END, message)
        self.results_text.see(tk.END)
        self.root.update_idletasks()
    
    def update_status(self, message: str, color: str = "black"):
        """Update the status label"""
        self.status_label.config(text=message, fg=color)
        self.root.update_idletasks()


def main():
    """Main entry point"""
    root = tk.Tk()
    app = DocSifterApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()