import cv2
import numpy as np
import os
import shutil

def process_character_sheet():
    # Find latest generated sheet
    artifacts_dir = '/Users/davidkotnik/.gemini/antigravity/brain/07019d04-a214-43ab-9565-86f4e8f17e5b'
    sheet_src = None
    candidates = []
    for f in os.listdir(artifacts_dir):
        if 'kai_walk_sheet' in f and f.endswith('.png'):
            candidates.append(os.path.join(artifacts_dir, f))
    
    if not candidates:
        print("Sheet not found.")
        return

    candidates.sort(key=os.path.getmtime, reverse=True)
    sheet_src = candidates[0]
    
    print(f"Processing {sheet_src}")
    img = cv2.imread(sheet_src)
    
    # Remove BG
    img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
    target = np.array([255, 0, 255]) 
    tol = 60
    lower = np.clip(target - tol, 0, 255)
    upper = np.clip(target + tol, 0, 255)
    mask = cv2.inRange(img[:,:,:3], lower, upper)
    img[mask > 0, 3] = 0
    
    # Save info about dimensions
    h, w = img.shape[:2]
    print(f"Total Sheet Dimensions: {w}x{h}")
    # 4x4 grid
    frame_w = w // 4
    frame_h = h // 4
    print(f"Calculated Frame Size: {frame_w}x{frame_h}")

    # Save
    dest_name = 'kai_walk_sheet.png'
    targets = [
        '/Users/davidkotnik/repos/novafarma/main/assets',
        '/Users/davidkotnik/repos/novafarma/assets/characters',
        '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Characters' # New location?
    ]
    
    # Also ensure dir exists
    if not os.path.exists('/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Characters'):
        os.makedirs('/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Characters')

    for t in targets:
        if not os.path.exists(t): os.makedirs(t, exist_ok=True)
        final_path = os.path.join(t, dest_name)
        cv2.imwrite(final_path, img)
        print(f"Saved to {final_path}")

if __name__ == "__main__":
    process_character_sheet()