import os
import sqlite3
import shutil

# --- CONFIGURATION ---
# The root directory where the original images are stored.
# IMPORTANT: This must match the IMAGE_DIR in your app.py
IMAGE_DIR = "/mnt/secret-items/sd-outputs/Sorted/Images/Portrait"

# The directory where the sorted images will be moved.
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "swiper-output")

# The path to the selections database.
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "image_selections.db")

def sort_swiped_images():
    """Reads the database, moves swiped files into action-based folders, and cleans the DB."""
    print(f"Connecting to database at {DB_PATH}...")
    if not os.path.exists(DB_PATH):
        print(f"Error: Database file not found at {DB_PATH}")
        return

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    try:
        # Get all selections from the database
        cursor.execute("SELECT id, image_path, action FROM image_selections")
        selections = cursor.fetchall()
    except sqlite3.OperationalError as e:
        print(f"Error querying database: {e}")
        print("This might mean the 'image_selections' table doesn't exist or is empty.")
        conn.close()
        return

    if not selections:
        print("No swiped images found in the database to sort.")
        conn.close()
        return

    print(f"Found {len(selections)} swiped images to sort.")
    moved_count = 0

    for selection_id, image_path, action in selections:
        # The image_path from the DB might have a leading '/images/' which we need to remove.
        # It's safer to handle both cases.
        clean_image_path = image_path.replace('/images/', '', 1) if image_path.startswith('/images/') else image_path
        
        source_path = os.path.join(IMAGE_DIR, clean_image_path)
        
        if not os.path.exists(source_path):
            print(f"Warning: Source file not found, skipping: {source_path}")
            # We should still remove this dangling DB record
            cursor.execute("DELETE FROM image_selections WHERE id = ?", (selection_id,))
            cursor.execute("DELETE FROM image_metadata WHERE path = ?", (clean_image_path,))
            print(f"Removed dangling database entry for {clean_image_path}")
            continue

        # Create the destination folder based on the action
        destination_folder = os.path.join(OUTPUT_DIR, action)
        os.makedirs(destination_folder, exist_ok=True)
        
        destination_path = os.path.join(destination_folder, os.path.basename(clean_image_path))

        try:
            print(f"Moving '{clean_image_path}' to '{action}' folder...")
            shutil.move(source_path, destination_path)
            moved_count += 1

            # If move was successful, remove the records from the database
            cursor.execute("DELETE FROM image_selections WHERE id = ?", (selection_id,))
            cursor.execute("DELETE FROM image_metadata WHERE path = ?", (clean_image_path,))
            print(f"Successfully moved and removed DB entries for {clean_image_path}")

        except Exception as e:
            print(f"Error moving file {source_path}: {e}")

    # Commit all changes to the database
    conn.commit()
    conn.close()

    print(f"\nSorting complete. Moved {moved_count} images.")

if __name__ == "__main__":
    sort_swiped_images()