import cv2
import numpy as np
import os
import shutil

def clean_ui_assets():
    # Paths to the newly generated UI artifacts (Solid Magenta #FF00FF usually)
    # I have to hunt them down by pattern again
    artifacts_dir = '/Users/davidkotnik/.gemini/antigravity/brain/07019d04-a214-43ab-9565-86f4e8f17e5b'
    
    mapping = {
        'ui_health_bar': 'health_bar.png',
        'ui_inventory_slot': 'inventory_slot.png',
        'ui_action_button': 'action_btn.png',
        'ui_dialog_box': 'dialog_panel.png'
    }
    
    # Destination in new project
    target_repo_dir = '/Users/davidkotnik/repos/novafarma/main/assets'
    target_project_dir = '/Users/davidkotnik/nova farma/main/assets'
    
    # Also copy to DEMO_FAZA1/UI if it exists?
    demo_ui_dir = '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/UI'
    if not os.path.exists(demo_ui_dir):
        os.makedirs(demo_ui_dir, exist_ok=True)

    for key_pattern, dest_name in mapping.items():
        found_path = None
        # Find latest file matching pattern
        candidates = []
        for f in os.listdir(artifacts_dir):
            if key_pattern in f and f.endswith('.png'):
                candidates.append(os.path.join(artifacts_dir, f))
        
        if not candidates:
            print(f"Skipping {key_pattern} (not found)")
            continue
            
        # Sort by mtime to get latest
        candidates.sort(key=os.path.getmtime, reverse=True)
        found_path = candidates[0]
        
        print(f"Processing {found_path}")
        img = cv2.imread(found_path)
        if img is None:
            print("Failed to load")
            continue
            
        # Convert to BGRA
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
        
        # Remove Magenta #FF00FF
        # OpenCV uses BGR: Magenta is (255, 0, 255) in BGR order
        # With tolerance just in case
        target = np.array([255, 0, 255]) 
        tol = 60 # strict-ish
        
        lower = np.clip(target - tol, 0, 255)
        upper = np.clip(target + tol, 0, 255)
        
        mask = cv2.inRange(img[:,:,:3], lower, upper)
        
        # Transparent wherever mask is true
        img[mask > 0, 3] = 0
        
        # Crop to content (non-transparent pixels)
        # Find bounding box of non-zero alpha
        alpha = img[:,:,3]
        coords = cv2.findNonZero(alpha)
        if coords is not None:
            x, y, w, h = cv2.boundingRect(coords)
            img = img[y:y+h, x:x+w]
        
        # Resize if massive? Generated images are likely 1024x1024
        # UI elements need to be manageable.
        # Health Bar: width 300?
        # Slot: 64x64 or 128x128?
        # Button: 128?
        # Panel: 400?
        
        h, w = img.shape[:2]
        new_w, new_h = w, h
        
        if 'health_bar' in dest_name and w > 400:
            scale = 400 / w
            new_w, new_h = int(w*scale), int(h*scale)
        elif 'inventory' in dest_name and w > 128:
            scale = 128 / w
            new_w, new_h = int(w*scale), int(h*scale)
        elif 'action' in dest_name and w > 140:
            scale = 128 / w
            new_w, new_h = int(w*scale), int(h*scale)
        elif 'dialog' in dest_name and w > 500:
            scale = 500 / w
            new_w, new_h = int(w*scale), int(h*scale)
            
        if new_w != w:
             img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

        # Save to Main Assets
        dest_path = os.path.join(target_repo_dir, dest_name)
        cv2.imwrite(dest_path, img)
        print(f"Saved {dest_name} to {dest_path}")
        
        # Sync to demo folder
        shutil.copy2(dest_path, os.path.join(demo_ui_dir, dest_name))
        print(f"Synced to {demo_ui_dir}")
        
        # Sync to Project Folder
        shutil.copy2(dest_path, os.path.join(target_project_dir, dest_name))
        print(f"Synced to {target_project_dir}")

if __name__ == "__main__":
    clean_ui_assets()