import os
import shutil
import re

# Configuration
SOURCE_ROOT = "MOJE_SLIKE_KONCNA"
DEST_ROOT = "assets/slike"

# Mapping User's requested folders
# Key = Folder in MOJE_SLIKE_KONCNA, Value = Folder in assets/slike
MAPPING = {
    "zombiji": "animations",     # "vse zombije in premike"
    "liki": "liki",              # "Kai, Ana in vse njune verzije"
    "teren": "teren",            # "vsa tla, trava in zemlja"
    "okolje": "teren",           # Merging Trees/Rocks into teren (closest match to "ground/earth")
    "predmeti": "MASTER_REFS",   # Items -> Refs (Assuming these are references/tools)
    "ostalo": "MASTER_REFS"      # Dustbin for everything else
}

def restore_assets():
    if not os.path.exists(DEST_ROOT):
        os.makedirs(DEST_ROOT)

    print(f"🚀 Restoring assets from {SOURCE_ROOT} to {DEST_ROOT}...")

    stats = {k: 0 for k in set(MAPPING.values())}

    for root, dirs, files in os.walk(SOURCE_ROOT):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                src_path = os.path.join(root, file)
                
                # Determine category based on parent folder in SOURCE
                # MOJE_SLIKE_KONCNA/zombiji/... -> category = zombiji
                rel_path = os.path.relpath(src_path, SOURCE_ROOT)
                top_folder = rel_path.split(os.sep)[0]
                
                # If top_folder matches mapping, use it, else default
                target_sub = MAPPING.get(top_folder, "MASTER_REFS")
                
                # Special Check: User mentioned "animations" for "moves"
                # If file has "walk", "run", "idle" -> force animations?
                if re.search(r"(walk|run|idle|attack|move)", file.lower()):
                    if top_folder == "liki" or top_folder == "zombiji":
                         # Characters moving -> likely animation, but user said "KAI -> Liki".
                         # "animations" -> "vse zombije in premike".
                         # I'll stick to: Zombies -> Animation. Kai -> Liki.
                         pass
                
                target_dir = os.path.join(DEST_ROOT, target_sub)
                
                # Preserve structure? No, user said "Vse nazaj točno tako kot je bilo prej... te mapah".
                # Implies flat or minimal nesting.
                # However, if animations have 100 frames, dumping them in one folder is chaos.
                # I will preserve the *immediate parent* if it looks like an animation sequence.
                # e.g. zombiji/walk/01.png -> assets/slike/animations/walk/01.png
                
                parent_name = os.path.basename(os.path.dirname(src_path))
                if parent_name.lower() not in ["zombiji", "liki", "teren", "okolje", "predmeti", "ostalo"]:
                    # Likely a sequence folder
                    target_dir = os.path.join(target_dir, parent_name)
                    
                if not os.path.exists(target_dir):
                    os.makedirs(target_dir)
                    
                dest_path = os.path.join(target_dir, file)
                
                try:
                    shutil.copy2(src_path, dest_path)
                    stats[target_sub] += 1
                except Exception as e:
                    print(f"Error copying {file}: {e}")

    print("="*40)
    print("✅ RESTORE COMPLETE")
    print("="*40)
    for folder, count in stats.items():
        print(f"  - assets/slike/{folder}/: {count} images")
    print("="*40)

if __name__ == "__main__":
    restore_assets()