import cv2
import numpy as np
import os
import glob

# Configuration
# Input files (Generated images)
INPUT_FILES = [
    {
        'path': '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/trees_adult_noir_1769706733583.png',
        'prefix': 'tree_adult',
        'output_folder': 'assets/DEMO_FAZA1/Trees'
    },
    {
        'path': '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/dead_nature_noir_1769706761105.png',
        'prefix': 'dead_nature',
        'output_folder': 'assets/DEMO_FAZA1/Environment'
    },
    {
        'path': '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/fence_sign_noir_1769706790281.png',
        'prefix': 'fence_sign',
        'output_folder': 'assets/DEMO_FAZA1/Environment'
    },
    {
        'path': '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/muddy_path_noir_v2_1769706866951.png',
        'prefix': 'path_tile',
        'output_folder': 'assets/DEMO_FAZA1/Ground',
        'mode': 'grid',
        'grid_size': 256
    },
    {
        'path': '/Users/davidkotnik/.gemini/antigravity/brain/8233d64e-0c17-43b1-b8b5-fbc41754e56b/water_texture_noir_1769717283296.png',
        'prefix': 'water_tile',
        'output_folder': 'assets/DEMO_FAZA1/Environment',
        'mode': 'grid',
        'grid_size': 256
    }
]

BASE_DIR = '/Users/davidkotnik/repos/novafarma'

def process_and_slice(file_data):
    full_path = file_data['path']
    prefix = file_data['prefix']
    out_dir_rel = file_data['output_folder']
    mode = file_data.get('mode', 'contour') # Default to contour
    
    out_dir = os.path.join(BASE_DIR, out_dir_rel)

    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    print(f"Processing {prefix} from {full_path} in {mode} mode...")
    
    img = cv2.imread(full_path, cv2.IMREAD_COLOR)
    if img is None:
        print(f"Error loading {full_path}")
        return

    h, w = img.shape[:2]

    if mode == 'grid':
        # Split into fixed grid cells
        tile_size = file_data.get('grid_size', 256)
        rows = h // tile_size
        cols = w // tile_size
        
        count = 0
        for r in range(rows):
            for c in range(cols):
                y = r * tile_size
                x = c * tile_size
                
                roi = img[y:y+tile_size, x:x+tile_size]
                
                # Check if tile is empty (all white/uniform?)
                # Calculate variance or average color.
                if np.mean(roi) > 250: # Mostly white
                    continue
                
                # We can enable alpha here if needed, but for ground tiles usually we want Opaque?
                # User wants "Tile Palette". Tiles usually have transparency if they are "Path on Grass".
                # Let's clean white background using same threshold logic.
                
                roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
                _, roi_mask = cv2.threshold(roi_gray, 240, 255, cv2.THRESH_BINARY_INV)
                
                roi_bgra = cv2.cvtColor(roi, cv2.COLOR_BGR2BGRA)
                roi_bgra[:, :, 3] = roi_mask
                
                # Save
                filename = f"{prefix}_{count}.png"
                out_path = os.path.join(out_dir, filename)
                cv2.imwrite(out_path, roi_bgra)
                print(f"  Saved Grid Tile {filename}")
                count += 1
                
    else: # Contour mode (old logic)
        # 1. Remove Background (White) -> Alpha
        # Assuming white background.
        # Convert to Gray for thresholding
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Threshold: Close to white (240+) is background
        _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
        # mask: White (255) is OBJECT, Black (0) is BACKGROUND

        # Clean mask (remove noise)
        kernel = np.ones((3,3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        # Dilate slightly to include outlines?
        # mask = cv2.dilate(mask, kernel, iterations=1)

        # 2. Find Contours (Separate Objects)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        count = 0
        for i, c in enumerate(contours):
            area = cv2.contourArea(c)
            if area < 500: # Filter small noise
                continue
            
            # Bounding box
            x, y, w, h = cv2.boundingRect(c)
            
            # Crop
            roi = img[y:y+h, x:x+w]
            
            # Create alpha channel for ROI
            # We need a mask for this specific object in the ROI
            roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
            # Re-threshold ROI to get clean alpha
            _, roi_mask = cv2.threshold(roi_gray, 240, 255, cv2.THRESH_BINARY_INV)
            
            # Convert ROI to BGRA
            roi_bgra = cv2.cvtColor(roi, cv2.COLOR_BGR2BGRA)
            roi_bgra[:, :, 3] = roi_mask

            # Save
            filename = f"{prefix}_{count}.png"
            out_path = os.path.join(out_dir, filename)
            cv2.imwrite(out_path, roi_bgra)
            print(f"  Saved {filename}")
            count += 1

if __name__ == "__main__":
    for f in INPUT_FILES:
        process_and_slice(f)