import os
import shutil
import json
import numpy as np
from PIL import Image
import colorsys

WORKSPACE_DIR = os.path.dirname(os.path.abspath(__file__))
ASSETS_DIR = os.path.join(WORKSPACE_DIR, "assets")
CLUSTER_DIR = os.path.join(WORKSPACE_DIR, "cluster_data")

os.makedirs(ASSETS_DIR, exist_ok=True)
os.makedirs(CLUSTER_DIR, exist_ok=True)

# Define the sources for the generated artifacts
BRAIN_DIR = r"C:\Users\Jeff Towers\.gemini\antigravity-cli\brain\79991b8b-c84e-4a18-b8f2-d7b22a1c42d4"
files = {
    "ironleaf_kush_seed": os.path.join(BRAIN_DIR, "ironleaf_kush_seed_1782674802209.jpg"),
    "ironleaf_kush_seedling": os.path.join(BRAIN_DIR, "ironleaf_kush_seedling_1782674811761.jpg"),
    "ironleaf_kush_mature": os.path.join(BRAIN_DIR, "ironleaf_kush_mature_1782674820400.jpg")
}

for name, src_path in files.items():
    if not os.path.exists(src_path):
        print(f"Missing {src_path}")
        continue
        
    # Open the image
    img = Image.open(src_path).convert("RGBA")
    
    # 1. Save base png
    base_png_path = os.path.join(ASSETS_DIR, f"{name}.png")
    img.save(base_png_path)
    
    # 2. Save neutral png (grayscale to act as luminosity map)
    gray_img = img.convert("L").convert("RGBA")
    neutral_png_path = os.path.join(ASSETS_DIR, f"{name}_neutral.png")
    gray_img.save(neutral_png_path)
    
    # 3. Create clustering labels
    pixels = np.array(img)
    labels = np.zeros((pixels.shape[0], pixels.shape[1]), dtype=int)
    
    for y in range(pixels.shape[0]):
        for x in range(pixels.shape[1]):
            r, g, b, a = pixels[y, x]
            h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
            
            # Simple hue-based mask for green (approx 60 to 180 degrees)
            if 0.16 < h < 0.5 and s > 0.1 and v > 0.1:
                labels[y, x] = 1  # GreenZone
            else:
                labels[y, x] = 0  # NonGreen
                
    # Flatten and save labels
    labels_flat = labels.reshape(-1)
    np.save(os.path.join(CLUSTER_DIR, f"{name}_labels.npy"), labels_flat)
    
    # Save mapping
    mapping = {
        "1": "GreenZone",
        "0": "NonGreen"
    }
    with open(os.path.join(CLUSTER_DIR, f"{name}_mapping.json"), 'w') as f:
        json.dump(mapping, f, indent=4)
        
    print(f"Processed {name}")

print("Assets and cluster data setup complete.")