import os
import shutil

BASE_DIR = os.path.abspath("assets/slike/glavna_referenca")
ANIMALS_DIR = os.path.join(BASE_DIR, "Animals")

# Subcategories
CATS = {
    "Domestic": ["cow", "pig", "chicken", "hen", "rooster", "dog", "cat", "horse", "sheep", "goat", "donkey", "ox", "bull", "domaca", "krava", "pujs", "kokos", "pes", "macka", "konj", "ovca", "koza"],
    "Wild": ["deer", "bear", "wolf", "rabbit", "hare", "fox", "boar", "bird", "eagle", "hawk", "owl", "snake", "frog", "toad", "lizard", "srna", "medved", "vuk", "zajec", "lisica", "divja", "ptica", "sokol", "sova", "kaca", "zaba"],
    "Insects": ["bee", "butterfly", "moth", "fly", "dragonfly", "firefly", "worm", "spider", "ant", "beetle", "ladybug", "wasp", "hornet", "mosquito", "cricket", "cicada", "mantis", "cebela", "metulj", "muha", "crv", "pajek", "mravlja", "hrosc"],
    "Infected": ["zombie", "mutant", "infected", "skeleton", "ghost", "monster", "beast", "undead", "boss", "guar", "dragon", "hydra", "yeti", "chimera", "basilisk", "kitsune", "kelpie", "roc", "medusa", "mermaid", "quetzalcoatl", "banshee", "cerberus", "gnome", "golem", "dinosaur", "trex", "raptor", "chupacabra", "kraken", "zombi", "okuzena", "pošast"]
}

def refine_animals():
    print("🦁 Refining 'Animals' folder...")
    
    # Ensure subdirs exist
    for sub in CATS:
        path = os.path.join(ANIMALS_DIR, sub)
        if not os.path.exists(path):
            os.makedirs(path)

    count = 0
    moved_out = 0
    
    # Iterate files in Animals root
    for filename in os.listdir(ANIMALS_DIR):
        if filename.startswith("."): continue
        src = os.path.join(ANIMALS_DIR, filename)
        if os.path.isdir(src): continue
        
        lower = filename.lower()
        target = None
        
        # 1. Check Keywords
        for category, keywords in CATS.items():
            if any(k in lower for k in keywords):
                target = category
                break
        
        if target:
            dst = os.path.join(ANIMALS_DIR, target, filename)
            # Collision
            if os.path.exists(dst):
                 name, ext = os.path.splitext(filename)
                 c=1
                 while os.path.exists(os.path.join(ANIMALS_DIR, target, f"{name}_{c}{ext}")):
                     c+=1
                 dst = os.path.join(ANIMALS_DIR, target, f"{name}_{c}{ext}")
            
            shutil.move(src, dst)
            count += 1
        else:
            # 2. Check for Impostors (Logic: If it's not an animal, kick it out?)
            # Or just leave it for manual check. User said: "kaj ne spada v to mapa poravilno dodaj v mapo katero pase"
            # Let's try to identify if it belongs to UI, Items etc.
            
            # Simple heuristic: move unknown to 'glavna_referenca/Items' if it sounds like an item? 
            # Or just move to root 'glavna_referenca' so Ultimate Sort can catch it again?
            # Let's move to root for safety.
            
            # For now, I will leave UNKNOWN animals in Animals root, 
            # BUT if I see obvious keywords like "tool", "house", "kai", I move them.
            
            IMPOSTORS = {
                "Items": ["tool", "axe", "seed", "potion", "sword", "gun", "armor"],
                "Environment": ["house", "tree", "rock", "wall", "grass"],
                "Characters": ["kai", "ana", "gronk", "human", "man", "woman"],
                "UI": ["button", "menu"]
            }
            
            pushed_out = False
            for cat, kws in IMPOSTORS.items():
                if any(k in lower for k in kws):
                    # Move to proper main folder
                    real_dst_dir = os.path.join(BASE_DIR, cat)
                    real_dst = os.path.join(real_dst_dir, filename)
                    shutil.move(src, real_dst)
                    moved_out += 1
                    pushed_out = True
                    print(f"   👋 Kicked impostor {filename} to {cat}")
                    break
            
            if not pushed_out:
                # Still in Animals root (Maybe an animal I missed)
                pass

    print(f"✨ Sorted {count} animals into subfolders.")
    print(f"👋 Kicked {moved_out} impostors out of Animals.")

if __name__ == "__main__":
    refine_animals()