import os
import shutil
import datetime
import asyncio
import aiofiles
from collections import defaultdict
import tkinter as tk
from tkinter import filedialog, ttk, messagebox
import threading
import time
 
class AdvancedFileOrganizer:
    def __init__(self, master):
        self.master = master
        self.master.title("Advanced File Organizer")
        self.master.geometry("800x600")
 
        self.search_dir = tk.StringVar()
        self.output_dir = tk.StringVar()
        self.dry_run = tk.BooleanVar(value=True)
        self.log_text = tk.StringVar()
 
        self.create_widgets()
 
        self.extension_mapping = {
            # Existing mappings
            '.jpg': 'Images', '.jpeg': 'Images', '.png': 'Images', '.gif': 'Images',
            '.bmp': 'Images', '.tiff': 'Images', '.eps': 'Images', '.psd': 'Images',
            '.webp': 'Images', '.svg': 'Images', '.ico': 'Images', '.heic': 'Images', 
            '.raw': 'Images', '.psd': 'Images', '.xcf': 'Images', 
 
            '.lnk': 'Shortcuts',
 
            '.pdf': 'Documents', 
            '.txt': 'Documents', '.md': 'Documents', '.rtf': 'Documents', 
            '.doc': 'Documents', '.docx': 'Documents', '.odt': 'Documents',
 
            '.csv': 'Spreadsheets', '.xlsx': 'Spreadsheets', '.xls': 'Spreadsheets', '.ods': 'Spreadsheets',
 
            '.wav': 'Audio', '.mp3': 'Audio', '.aiff': 'Audio', '.aac': 'Audio',
            '.wma': 'Audio', '.flac': 'Audio', '.mid': 'Audio', '.midi': 'Audio',
            '.m4a': 'Audio', '.ogg': 'Audio', '.m4p': 'Audio', 
 
            '.mp4': 'Videos', '.avi': 'Videos', '.mov': 'Videos', '.wmv': 'Videos',  
            '.flv': 'Videos', '.mkv': 'Videos', '.webm': 'Videos', '.mpg': 'Videos', 
            '.mpeg': 'Videos', '.aep': 'Videos', '.prproj': 'Videos', 
 
            '.stl': '3D Prints', '.ply': '3D Prints', '.3mf': '3D Prints', 
            '.gcode': '3D Prints', '.obj': '3D Prints',
 
            '.ttc': 'Fonts', '.otf': 'Fonts', '.eot': 'Fonts', '.fon': 'Fonts', 
            '.ttf': 'Fonts', '.woff': 'Fonts', '.woff2': 'Fonts',
 
            '.zip': 'Compressed', '.rar': 'Compressed', '.7z': 'Compressed', 
            '.tar': 'Compressed', '.gz': 'Compressed',
 
            '.exe': 'Executables', '.msi': 'Executables', '.app': 'Executables',
 
            '.py': 'Code', '.java': 'Code', '.cpp': 'Code', '.h': 'Code', 
            '.js': 'Code', '.html': 'Code', '.css': 'Code', '.php': 'Code', 
            '.sh': 'Code', '.bat': 'Code',
 
            '.ppt': 'Presentations', '.pptx': 'Presentations', '.key': 'Presentations', '.odp': 'Presentations',
 
            '.epub': 'eBooks', '.mobi': 'eBooks', '.azw': 'eBooks',
 
            '.db': 'Databases', '.sqlite': 'Databases', '.sql': 'Databases',
 
            '.vdi': 'Virtual Machines', '.vmdk': 'Virtual Machines', '.ova': 'Virtual Machines',
 
            '.indd': 'Design', '.qxd': 'Design', '.ai': 'Design', 
            '.cdr': 'Design', '.dwg': 'Design', '.dxf': 'Design', 
 
            '.iso': 'Disk Images', '.dmg': 'Disk Images',
 
            '.json': 'Data Files', '.xml': 'Data Files', '.yaml': 'Data Files',
 
            # New mappings
            '.step': 'Design', '.stp': 'Design', '.iges': 'Design', '.igs': 'Design',
            '.eps': 'Design', '.svg': 'Design',
            '.bak': 'Backups', '.old': 'Backups', '.tmp': 'Backups',
            '.ini': 'ConfigFiles', '.cfg': 'ConfigFiles', '.conf': 'ConfigFiles', '.config': 'ConfigFiles',
            '.ps1': 'Code', '.vbs': 'Code', '.bash': 'Code',
            '.log': 'Logs', '.syslog': 'Logs'
        }
 
    def create_widgets(self):
        # Search directory selection
        ttk.Label(self.master, text="Search Directory:").grid(row=0, column=0, padx=5, pady=5, sticky="w")
        ttk.Entry(self.master, textvariable=self.search_dir, width=50).grid(row=0, column=1, padx=5, pady=5)
        ttk.Button(self.master, text="Browse", command=self.browse_search_dir).grid(row=0, column=2, padx=5, pady=5)
 
        # Output directory selection
        ttk.Label(self.master, text="Output Directory:").grid(row=1, column=0, padx=5, pady=5, sticky="w")
        ttk.Entry(self.master, textvariable=self.output_dir, width=50).grid(row=1, column=1, padx=5, pady=5)
        ttk.Button(self.master, text="Browse", command=self.browse_output_dir).grid(row=1, column=2, padx=5, pady=5)
 
        # Dry run checkbox
        ttk.Checkbutton(self.master, text="Dry Run", variable=self.dry_run).grid(row=2, column=0, padx=5, pady=5, sticky="w")
 
        # Run button
        ttk.Button(self.master, text="Run", command=self.run_organizer).grid(row=2, column=1, padx=5, pady=5)
 
        # Log display
        self.log_display = tk.Text(self.master, wrap=tk.WORD, width=80, height=20)
        self.log_display.grid(row=3, column=0, columnspan=3, padx=5, pady=5)
 
        # Scrollbar for log display
        scrollbar = ttk.Scrollbar(self.master, orient="vertical", command=self.log_display.yview)
        scrollbar.grid(row=3, column=3, sticky="ns")
        self.log_display.configure(yscrollcommand=scrollbar.set)
 
        # Commit changes button (initially disabled)
        self.commit_button = ttk.Button(self.master, text="Commit Changes", command=self.run_commit_changes, state="disabled")
        self.commit_button.grid(row=4, column=1, padx=5, pady=5)
 
        # Progress bar
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(self.master, variable=self.progress_var, maximum=100)
        self.progress_bar.grid(row=5, column=0, columnspan=3, padx=5, pady=5, sticky="ew")
 
    def browse_search_dir(self):
        directory = filedialog.askdirectory()
        self.search_dir.set(directory)
 
    def browse_output_dir(self):
        directory = filedialog.askdirectory()
        self.output_dir.set(directory)
 
    def run_organizer(self):
        self.log_display.delete(1.0, tk.END)
        self.commit_button.config(state="disabled")
        
        if not self.search_dir.get() or not self.output_dir.get():
            self.log("Please select both search and output directories.")
            return
 
        # Run the organizer in a separate thread to keep the UI responsive
        threading.Thread(target=self._run_organizer_thread, daemon=True).start()
 
    def _run_organizer_thread(self):
        search_directory = self.search_dir.get()
        output_directory = self.output_dir.get()
        is_dry_run = self.dry_run.get()
 
        date_code = datetime.datetime.now().strftime("%m%d%Y")
        self.log(f"Searching directory: {search_directory}")
        self.log(f"Output directory: {output_directory}")
        self.log(f"Dry run: {'Yes' if is_dry_run else 'No'}")
 
        categorized_files = defaultdict(list)
 
        # Scan the directory and categorize files
        total_files = 0
        for root, _, files in os.walk(search_directory):
            for file in files:
                total_files += 1
                file_path = os.path.join(root, file)
                _, ext = os.path.splitext(file)
                category = self.extension_mapping.get(ext.lower(), 'Misc')
                categorized_files[category].append(file_path)
 
        # Process files
        self.changes = []
        processed_files = 0
        for category, files in categorized_files.items():
            if files:
                target_folder = os.path.join(output_directory, f'{category}_{date_code}')
                self.log(f"Category: {category}")
                for file_path in files:
                    target_path = os.path.join(target_folder, os.path.basename(file_path))
                    self.changes.append((file_path, target_path))
                    # Extract the filename from the file path
                    filename = os.path.basename(file_path)  
                    # Extract the target folder name from the target path
                    target_folder_name = os.path.basename(target_folder)
                    # Format the log message with aligned filenames and target folders
                    self.log(f"  {filename:<30}{target_folder_name}")  
                    processed_files += 1
                    self.update_progress(processed_files, total_files)
 
        if not is_dry_run:
            self.run_commit_changes()
        else:
            self.log("\nDry run completed. Review the log and click 'Commit Changes' to apply.")
            self.master.after(0, lambda: self.commit_button.config(state="normal"))
 
    def update_progress(self, current, total):
        progress = (current / total) * 100
        self.progress_var.set(progress)
        self.master.update_idletasks()
 
    def run_commit_changes(self):
        threading.Thread(target=self._commit_changes_thread, daemon=True).start()
 
    def _commit_changes_thread(self):
        asyncio.run(self.commit_changes())
 
    async def commit_changes(self):
        if not self.changes:
            self.log("No changes to commit.")
            return
 
        total_files = len(self.changes)
        self.progress_var.set(0)
        self.progress_bar["maximum"] = total_files
 
        async def process_batch(batch):
            tasks = [self.async_move_file(source, target) for source, target in batch]
            results = await asyncio.gather(*tasks)
            for result in results:
                self.log(result)
                self.progress_var.set(self.progress_var.get() + 1)
                self.master.update_idletasks()
 
        # Process files in batches of 10
        batch_size = 10
        for i in range(0, len(self.changes), batch_size):
            batch = self.changes[i:i+batch_size]
            await process_batch(batch)
 
        self.log("All changes committed.")
        self.changes = []
        self.commit_button.config(state="disabled")
 
    async def async_move_file(self, source, target):
        try:
            # Create target directory if it doesn't exist
            os.makedirs(os.path.dirname(target), exist_ok=True)
 
            # Check if the file already exists in the target location
            if os.path.exists(target):
                base, extension = os.path.splitext(target)
                counter = 1
                while os.path.exists(f"{base}_{counter}{extension}"):
                    counter += 1
                target = f"{base}_{counter}{extension}"
 
            # Check file size
            file_size = os.path.getsize(source)
            if file_size > 100 * 1024 * 1024:  # 100 MB
                # For large files, use shutil.move with a callback for progress
                await asyncio.to_thread(self.move_large_file, source, target, file_size)
            else:
                # For smaller files, use aiofiles
                async with aiofiles.open(source, 'rb') as src, aiofiles.open(target, 'wb') as dst:
                    await dst.write(await src.read())
                await asyncio.to_thread(os.remove, source)
 
            return f"Moved: '{source}' to '{target}'"
        except PermissionError:
            return f"Permission denied: Unable to move '{source}' to '{target}'"
        except FileNotFoundError:
            return f"File not found: '{source}'"
        except OSError as e:
            if "on a different device" in str(e):
                # File is on a different device (e.g., network drive), use shutil.copy2 and then remove original
                await asyncio.to_thread(shutil.copy2, source, target)
                await asyncio.to_thread(os.remove, source)
                return f"Copied and removed: '{source}' to '{target}' (different device)"
            else:
                return f"Error moving '{source}' to '{target}': {str(e)}"
        except Exception as e:
            return f"Unexpected error moving '{source}' to '{target}': {str(e)}"
 
    def move_large_file(self, source, target, total_size):
        with open(source, 'rb') as src, open(target, 'wb') as dst:
            copied = 0
            while True:
                buf = src.read(8388608)  # 8MB chunks
                if not buf:
                    break
                dst.write(buf)
                copied += len(buf)
                progress = (copied / total_size) * 100
                self.master.after(0, lambda: self.progress_var.set(progress))
        os.remove(source)
 
    def log(self, message):
        self.log_display.insert(tk.END, message + "\n")
        self.log_display.see(tk.END)
 
if __name__ == "__main__":
    root = tk.Tk()
    app = AdvancedFileOrganizer(root)
    root.mainloop()