import os
import shutil
import tkinter as tk
from tkinter import filedialog, messagebox, simpledialog
import re
from collections import defaultdict

# Initialize global variables
folder_path = ""
groups = {}
display_to_group_map = {}
pending_renames = {}  # Stores pending renames: {original_group_name: new_group_name}
pending_merges = []   # Stores pending merges: list of group_names to merge into new_group_name
merge_history = []    # Stack to keep track of merges for undo functionality
rename_in_progress = False  # Flag to prevent multiple rename operations

def clean_filename(filename, redundant_words):
    """Clean filename according to specified rules (remove numbers, repeated words)."""
    # 1. Ignore text before and including '：'.
    filename = filename.split('：', 1)[-1]

    # 2. Remove file extension.
    filename = os.path.splitext(filename)[0]

    # 3. Replace delimiters with spaces, and transform to lowercase.
    words = filename.replace('_', ' ').lower().split()

    # Check if the entire filename is numeric
    if all(part.isdigit() for part in words):
        return ' '.join(words)

    # 4. Replace numeric groups with spaces
    words = [re.sub(r'\d+', ' ', word).strip() for word in words]

    # Flatten the words split further on spaces
    words = ' '.join(words).split()

    # 5. Remove redundant words/groups of characters.
    unique_words = [w for w in words if w not in redundant_words]

    # 6. Remove words that are less than 3 characters and ensure uniqueness across filename
    seen = set()
    cleaned_words = [word for word in unique_words if len(word) > 2 and not (word in seen or seen.add(word))]

    # Join cleaned words back into a string
    cleaned_filename = ' '.join(cleaned_words)

    # Use a placeholder for files that end up without significant content
    return cleaned_filename if cleaned_filename else "misc_files"

def find_redundant_words(files):
    """Identify groups of characters that appear in all filenames."""
    if not files:
        return []

    # Split filenames on '：' and take the relevant part
    processed_files = [file.split('：', 1)[-1].replace('_', ' ').lower().split() for file in files]

    # Find common words across all processed filenames
    common_words = set(processed_files[0])
    for file_words in processed_files[1:]:
        common_words.intersection_update(file_words)

    redundant_words = list(common_words)
    return redundant_words

def extract_sort_key(filename):
    """Extract numerical suffix from the filename for sorting."""
    # Extract part after 's_' and before file extension
    # Example: 'Ifrete_knit_s_13.jpg' -> 13
    match = re.search(r's_(\d+)', filename.lower())
    if match:
        return int(match.group(1))
    return 0  # Default if pattern not found

def prepare_groups(files, grouping_enabled):
    """Prepare groups of files by cleaned filenames or treat all as one group."""
    if not files:
        return {}

    if grouping_enabled:
        # Identify redundant words/groups of characters
        redundant_words = find_redundant_words(files)

        temp_groups = defaultdict(list)

        for file in files:
            cleaned_name = clean_filename(file, redundant_words)
            temp_groups[cleaned_name].append(file)

        # Sort files within each group based on numerical suffix
        for group in temp_groups:
            temp_groups[group].sort(key=extract_sort_key)

        return temp_groups
    else:
        # Treat all files as a single group named "All Files"
        return {"All Files": sorted(files)}

def split_into_subgroups(files, file_num):
    """Split files into subgroups based on file_num and remaining files rule."""
    subgroups = []
    total_files = len(files)
    index = 0

    while index < total_files:
        remaining = total_files - index
        if remaining <= 2 * file_num:
            subgroups.append(files[index:])
            break
        else:
            subgroups.append(files[index:index + file_num])
            index += file_num

    return subgroups

def capitalize_folder_name(folder_name):
    """Capitalize folder name excluding certain short words."""
    short_words = {'of', 'is', 'a', 'an'}
    words = folder_name.split()
    capitalized = [word.capitalize() if word.lower() not in short_words else word for word in words]
    modified_name = ' '.join(capitalized)
    return modified_name

def get_preview_groups():
    """Create a preview of groups by applying pending renames and merges."""
    # Start with a copy of the current groups
    preview = groups.copy()

    # Apply pending renames
    for original_group, new_group in pending_renames.items():
        if original_group in preview:
            preview[new_group] = preview.pop(original_group)

    # Apply pending merges
    for merge_op in pending_merges:
        groups_to_merge = merge_op['groups']
        new_group_name = merge_op['new_group']
        merged_files = []

        for group in groups_to_merge:
            merged_files.extend(preview.get(group, []))
            if group in preview:
                del preview[group]

        if new_group_name in preview:
            preview[new_group_name].extend(merged_files)
        else:
            preview[new_group_name] = merged_files

    return preview

def update_folder_info(preview_groups):
    """Update the informational labels with folder details."""
    # Clear previous info
    total_files_label.config(text=f"Total Files: {sum(len(files) for files in preview_groups.values())}")

def show_folder_contents(event=None):
    """Display files or subfolders in the selected group (Preview)."""
    try:
        selection = folder_listbox.curselection()
        if not selection:
            content_listbox.delete(0, tk.END)
            return

        folder_display_name = folder_listbox.get(selection[0])
        original_group_name = display_to_group_map.get(folder_display_name, folder_display_name)

        # Get the preview groups
        preview_groups = get_preview_groups()

        # Update folder info
        update_folder_info(preview_groups)

        # Clear previous contents
        content_listbox.delete(0, tk.END)

        # Get the list of files in the selected group, considering pending operations
        files_in_group = preview_groups.get(original_group_name, [])

        if not files_in_group:
            content_listbox.insert(tk.END, "No files in this group.")
            return

        if group_files_var.get():
            file_num_str = file_number_var.get().strip()
            if file_num_str:
                try:
                    file_num = int(file_num_str)
                    if file_num <= 0:
                        raise ValueError
                except ValueError:
                    messagebox.showerror("Invalid Input", "Please enter a valid positive integer for file grouping.")
                    content_listbox.delete(0, tk.END)  # Clear contents if invalid input
                    return

                if len(files_in_group) >= 2 * file_num:
                    subgroups = split_into_subgroups(files_in_group, file_num)
                    for idx, subgroup in enumerate(subgroups, 1):
                        subfolder_display_name = f"{folder_display_name} - {idx}"
                        content_listbox.insert(tk.END, f"{subfolder_display_name}:")
                        for file in subgroup:
                            content_listbox.insert(tk.END, f"  - {file}")
                else:
                    for file in files_in_group:
                        content_listbox.insert(tk.END, file)
            else:
                for file in files_in_group:
                    content_listbox.insert(tk.END, file)
        else:
            for file in files_in_group:
                content_listbox.insert(tk.END, file)

        if enable_grouping_var.get():
            file_number_info.config(text=f"Files in this Group: {len(files_in_group)}")
        else:
            file_number_info.config(text="")
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while displaying folder contents:\n{str(e)}")

def load_directory():
    """Load the selected directory, prepare groups, and update the displayed folders."""
    global folder_path, groups, pending_renames, pending_merges, merge_history
    folder_path = filedialog.askdirectory(title="Select the folder containing your files")
    if not folder_path:
        return
    parent_folder_name = os.path.basename(folder_path)
    parent_name_var.set(parent_folder_name)
    try:
        files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
        groups = prepare_groups(files, enable_grouping_var.get())
        pending_renames.clear()
        pending_merges.clear()
        merge_history.clear()
        update_displayed_folders()
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while loading the directory:\n{str(e)}")

def apply_options(*args):
    """Apply the selected options and update the preview."""
    if not folder_path:
        return
    try:
        update_displayed_folders()
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while applying options:\n{str(e)}")

def toggle_grouping(*args):
    """Handle changes in the folder grouping setting."""
    if not folder_path:
        return
    try:
        files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
        groups.clear()  # Reset any current groups
        groups.update(prepare_groups(files, enable_grouping_var.get()))
        update_displayed_folders()
        show_folder_contents()
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while toggling grouping:\n{str(e)}")

def update_displayed_folders():
    """Update the folder_listbox based on current groups and selected options."""
    global display_to_group_map
    parent_prefix = parent_name_var.get().strip()
    add_prefix = prefix_var.get()
    capitalize_names = capitalize_var.get()

    folder_listbox.delete(0, tk.END)
    display_to_group_map.clear()

    # Get the preview groups
    preview_groups = get_preview_groups()

    # Update folder info
    update_folder_info(preview_groups)

    for folder_name in sorted(preview_groups.keys()):
        display_name = folder_name

        if enable_grouping_var.get():
            # When grouping is enabled
            if folder_name != "All Files":
                modified_folder_name = capitalize_folder_name(display_name) if capitalize_names else display_name
                modified_folder_name = f"{parent_prefix} - {modified_folder_name}" if add_prefix else modified_folder_name
            else:
                modified_folder_name = parent_prefix
        else:
            # When grouping is disabled
            modified_folder_name = parent_prefix

        folder_listbox.insert(tk.END, modified_folder_name)
        display_to_group_map[modified_folder_name] = folder_name

    if folder_listbox.size() > 0:
        folder_listbox.select_set(0)
        show_folder_contents()

def reset_program():
    """Reset the program to its initial state after sorting."""
    global folder_path, groups, display_to_group_map, pending_renames, pending_merges, merge_history, rename_in_progress

    folder_path = ""
    groups.clear()
    display_to_group_map.clear()
    pending_renames.clear()
    pending_merges.clear()
    merge_history.clear()
    rename_in_progress = False

    folder_listbox.delete(0, tk.END)
    content_listbox.delete(0, tk.END)
    parent_name_var.set("")
    total_files_label.config(text="Total Files: 0")
    file_number_info.config(text="")

def execute_sorting():
    """Execute the sorting: move files into organized folders and subfolders based on selected options."""
    if not folder_path or not groups:
        messagebox.showerror("Error", "No directory loaded or no files grouped. Please load a directory first.")
        return

    try:
        parent_prefix = parent_name_var.get().strip()
        add_prefix = prefix_var.get()
        group_files = group_files_var.get()
        capitalize_names = capitalize_var.get()
        grouping_enabled = enable_grouping_var.get()
        file_num_str = file_number_var.get().strip()

        file_num = None
        if group_files and file_num_str:
            try:
                file_num = int(file_num_str)
                if file_num <= 0:
                    raise ValueError
            except ValueError:
                messagebox.showerror("Invalid Input", "Please enter a valid positive integer for file grouping.")
                return

        for original_group, new_group in pending_renames.items():
            if original_group in groups:
                groups[new_group] = groups.pop(original_group)

        pending_renames.clear()
        pending_merges.clear()
        merge_history.clear()

        for merge_operation in pending_merges:
            groups_to_merge = merge_operation['groups']
            new_group_name = merge_operation['new_group']
            merged_files = []

            for group in groups_to_merge:
                if group in groups:
                    merged_files.extend(groups.pop(group, []))

            if new_group_name in groups:
                groups[new_group_name].extend(merged_files)
            else:
                groups[new_group_name] = merged_files

        for folder_display_name, group_name in display_to_group_map.items():
            files = groups.get(group_name, [])

            if group_files and file_num and len(files) >= 2 * file_num:
                subgroups = split_into_subgroups(files, file_num)
                for idx, subgroup in enumerate(subgroups, 1):
                    subfolder_name = f"{parent_prefix} - {idx}"
                    subfolder_path = os.path.join(folder_path, subfolder_name)
                    os.makedirs(subfolder_path, exist_ok=True)
                    for file in subgroup:
                        src = os.path.join(folder_path, file)
                        dest = os.path.join(subfolder_path, file)
                        shutil.move(src, dest)

            elif grouping_enabled and group_name != "All Files":
                base_folder_name = capitalize_folder_name(group_name) if capitalize_names else group_name
                modified_folder_name = f"{parent_prefix} - {base_folder_name}" if add_prefix else base_folder_name
                new_folder_path = os.path.join(folder_path, modified_folder_name)

                if os.path.abspath(new_folder_path) == os.path.abspath(folder_path):
                    messagebox.showerror("Error", f"Cannot create a folder with the same name as the parent folder: {modified_folder_name}")
                    continue

                os.makedirs(new_folder_path, exist_ok=True)
                for file in files:
                    src = os.path.join(folder_path, file)
                    dest = os.path.join(new_folder_path, file)
                    shutil.move(src, dest)

            elif not grouping_enabled or group_name == "All Files":
                subfolder_name = f"{parent_prefix}"
                subfolder_path = os.path.join(folder_path, subfolder_name)
                if os.path.abspath(subfolder_path) != os.path.abspath(folder_path):
                    os.makedirs(subfolder_path, exist_ok=True)
                    for file in files:
                        src = os.path.join(folder_path, file)
                        dest = os.path.join(subfolder_path, file)
                        shutil.move(src, dest)

        messagebox.showinfo("Success", "Files sorted and organized successfully!")
        reset_program()
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred during sorting:\n{str(e)}")

def rename_selected_folder():
    """Rename the selected folder directly in the listbox."""
    global rename_in_progress
    selected_indices = folder_listbox.curselection()
    if len(selected_indices) != 1:
        messagebox.showinfo("Rename Folder", "Please select exactly one folder to rename.")
        return

    if rename_in_progress:
        messagebox.showinfo("Rename in Progress", "Please complete the current rename operation before starting a new one.")
        return

    rename_in_progress = True

    index = selected_indices[0]
    old_display_name = folder_listbox.get(index)
    original_group_name = display_to_group_map.get(old_display_name, old_display_name)

    entry = tk.Entry(frame, width=40)
    entry.insert(0, old_display_name)
    entry.select_range(0, tk.END)
    entry.focus()

    bbox = folder_listbox.bbox(index)
    if bbox:
        entry.place(x=folder_listbox.winfo_x() + bbox[0], y=folder_listbox.winfo_y() + bbox[1], width=bbox[2], height=bbox[3])

    def on_enter(event=None):
        global rename_in_progress
        new_display_name = entry.get().strip()
        if not new_display_name:
            messagebox.showerror("Invalid Name", "Folder name cannot be empty.")
            rename_in_progress = False
            entry.destroy()
            root.unbind("<Button-1>")
            return
        if new_display_name in display_to_group_map and new_display_name != old_display_name:
            messagebox.showerror("Duplicate Name", "A folder with this name already exists.")
            rename_in_progress = False
            entry.destroy()
            root.unbind("<Button-1>")
            return

        pending_renames[original_group_name] = new_display_name

        folder_listbox.delete(index)
        folder_listbox.insert(index, new_display_name)
        display_to_group_map[new_display_name] = original_group_name
        del display_to_group_map[old_display_name]

        entry.destroy()
        rename_in_progress = False
        root.unbind("<Button-1>")
        update_displayed_folders()

    entry.bind("<Return>", on_enter)

    def on_click_anywhere(event):
        if event.widget != entry:
            on_enter()

    root.bind("<Button-1>", on_click_anywhere)
    entry.bind("<FocusOut>", on_enter)

def merge_selected_folders():
    """Merge the selected folders into a single folder."""
    selected_indices = folder_listbox.curselection()
    if len(selected_indices) < 2:
        messagebox.showinfo("Merge Folders", "Please select two or more folders to merge.")
        return

    selected_display_names = [folder_listbox.get(i) for i in selected_indices]
    selected_group_names = [display_to_group_map.get(name, name) for name in selected_display_names]

    default_new_folder_name = selected_display_names[0]
    new_folder_name = simpledialog.askstring("Merge Folders", "Enter the name for the merged folder:", initialvalue=default_new_folder_name)
    if not new_folder_name:
        messagebox.showerror("Invalid Name", "Folder name cannot be empty.")
        return

    if (new_folder_name in display_to_group_map and
            display_to_group_map[new_folder_name] not in selected_group_names and
            new_folder_name not in [merge['new_group'] for merge in pending_merges]):
        messagebox.showerror("Duplicate Name", "A folder with this name already exists.")
        return

    for merge_op in pending_merges:
        if new_folder_name == merge_op['new_group']:
            messagebox.showerror("Duplicate Name", "A pending merge to this name already exists.")
            return

    pending_merges.append({'groups': selected_group_names, 'new_group': new_folder_name})
    merge_history.append({'groups': selected_group_names, 'new_group': new_folder_name})

    for i in reversed(selected_indices):
        display_name = folder_listbox.get(i)
        folder_listbox.delete(i)
        del display_to_group_map[display_name]

    folder_listbox.insert(tk.END, new_folder_name)
    display_to_group_map[new_folder_name] = new_folder_name
    update_displayed_folders()

def undo_merge():
    """Undo the last merge operation."""
    if not merge_history:
        messagebox.showinfo("Undo Merge", "No merge operations to undo.")
        return

    last_merge = merge_history.pop()
    groups_to_split = last_merge['groups']
    new_group_name = last_merge['new_group']

    pending_merges[:] = [merge for merge in pending_merges if merge['new_group'] != new_group_name]

    for group in groups_to_split:
        original_group_name = group
        pending_renames.pop(original_group_name, None)

    folder_listbox.delete(tk.END)
    for group in groups_to_split:
        folder_listbox.insert(tk.END, group)
        display_to_group_map[group] = group

    if new_group_name in display_to_group_map:
        del display_to_group_map[new_group_name]

    update_displayed_folders()

def create_context_menu(event):
    """Create and display a context menu for the folder listbox."""
    try:
        context_menu = tk.Menu(root, tearoff=0)
        context_menu.add_command(label="Rename", command=rename_selected_folder)

        selected_indices = folder_listbox.curselection()
        if len(selected_indices) > 1:
            context_menu.add_command(label="Merge Selected Folders", command=merge_selected_folders)

        if merge_history:
            context_menu.add_separator()
            context_menu.add_command(label="Undo Last Merge", command=undo_merge)

        context_menu.post(event.x_root, event.y_root)
    except Exception as e:
        messagebox.showerror("Error", f"An error occurred while creating the context menu:\n{str(e)}")

root = tk.Tk()
root.title("File Sorter")
root.geometry("900x600")
root.resizable(False, False)

frame = tk.Frame(root)
frame.pack(pady=20, padx=20, fill=tk.BOTH, expand=True)

folder_listbox_width = 40
content_listbox_width = 70

parent_name_label = tk.Label(frame, text="Parent Folder Name:")
parent_name_label.grid(row=0, column=0, padx=5, sticky='w')

parent_name_var = tk.StringVar()
parent_name_entry = tk.Entry(frame, width=50, textvariable=parent_name_var)
parent_name_entry.grid(row=0, column=1, columnspan=2, padx=5, sticky='ew')
parent_name_var.trace_add('write', lambda *args: apply_options())

total_files_label = tk.Label(frame, text="Total Files: 0")
total_files_label.grid(row=0, column=3, padx=5, sticky='w')

prefix_var = tk.IntVar(value=0)  # Checkbox unchecked by default
prefix_check = tk.Checkbutton(frame, text="Add Parent Name Prefix", variable=prefix_var, command=apply_options)
prefix_check.grid(row=1, column=0, columnspan=2, pady=5, sticky='w')

group_files_var = tk.IntVar(value=1)  # Checkbox checked by default
group_files_check = tk.Checkbutton(frame, text="Group by file number", variable=group_files_var, command=apply_options)
group_files_check.grid(row=2, column=0, pady=5, sticky='w')

file_number_label = tk.Label(frame, text="Files per subfolder:")
file_number_label.grid(row=2, column=1, padx=0, sticky='w')

file_number_var = tk.StringVar(value="30")  # Default for file grouping
file_number_entry = tk.Entry(frame, width=5, textvariable=file_number_var)
file_number_entry.grid(row=2, column=1, sticky='w', padx=(110, 0))
file_number_var.trace_add('write', lambda *args: apply_options())

capitalize_var = tk.IntVar(value=1)  # Checkbox checked by default
capitalize_check = tk.Checkbutton(frame, text="Capitalize Folder Names", variable=capitalize_var, command=apply_options)
capitalize_check.grid(row=3, column=0, columnspan=2, pady=5, sticky='w')

enable_grouping_var = tk.IntVar(value=1)  # Initially disabled
enable_grouping_check = tk.Checkbutton(frame, text="Enable Folder Grouping", variable=enable_grouping_var, command=apply_options)
enable_grouping_check.grid(row=4, column=0, columnspan=2, pady=5, sticky='w')
enable_grouping_var.trace_add('write', toggle_grouping)

folder_listbox_label = tk.Label(frame, text="Grouped Folders (Preview):", font=('Century Gothic', 10, 'bold'))
folder_listbox_label.grid(row=5, column=0, sticky='w')

folder_listbox = tk.Listbox(frame, width=folder_listbox_width, height=15, exportselection=False, selectmode=tk.EXTENDED)
folder_listbox.grid(row=6, column=0, columnspan=2, pady=5, sticky='nsew')
folder_listbox.bind("<<ListboxSelect>>", show_folder_contents)
folder_listbox.bind("<Double-Button-1>", lambda event: rename_selected_folder())

folder_listbox.bind("<Button-3>", create_context_menu)
folder_listbox.bind("<Button-2>", create_context_menu)

content_listbox_label = tk.Label(frame, text="Folder Contents (Preview):", font=('Century Gothic', 10, 'bold'))
content_listbox_label.grid(row=5, column=2, sticky='w')

content_listbox = tk.Listbox(frame, width=content_listbox_width, height=15, exportselection=False)
content_listbox.grid(row=6, column=2, columnspan=2, pady=5, sticky='nsew')

frame.grid_rowconfigure(6, weight=1)
frame.grid_columnconfigure(1, weight=1)

file_number_info = tk.Label(frame, text="", font=('Century Gothic', 8))
file_number_info.grid(row=7, column=3, pady=(0,20), sticky='ne')

browse_button = tk.Button(frame, text="Browse Directory", command=load_directory)
browse_button.grid(row=8, column=0, pady=(10, 0), sticky='w')

execute_button = tk.Button(frame, text="Execute Sorting", command=execute_sorting)
execute_button.grid(row=8, column=3, pady=(10, 0), sticky='e')

root.mainloop()

When Folder Grouping is DISABLED and Group by file number is ENABLE,  I would like to display all the subfolders in the Folder Listbox.