import os
from rembg import remove, new_session
from PIL import Image
import numpy as np

# --- CONFIG ---
GRONK_INPUT_DIR = "assets/slike 🟢/demo 🔴/characters 🔴/gronk/"
GRONK_OUTPUT_FILE = "assets/characters/gronk_pro_sheet.png"
CANNABIS_PATHS = [
     "assets/PHASE_PACKS/0_DEMO/crops/cannabis/growth_stages/cannabis_stage4_ready.png",
     "assets/PHASE_PACKS/1_FAZA_1/crops/cannabis/growth_stages/cannabis_stage4_ready.png",
     "references/crops/cannabis/growth_stages/cannabis_stage4_ready.png"
]

FRAME_SIZE = 128
SHEET_COLS = 4
SHEET_ROWS = 5 # 0:WalkDown, 1:WalkRight, 2:WalkUp, 3:Vape, 4:Idle

# Initialize rembg session
session = new_session()

def clean_and_center(img_path, target_size=(128, 128)):
    print(f"Processing {os.path.basename(img_path)}...")
    
    # 1. Load and Remove Background
    original = Image.open(img_path).convert("RGBA")
    
    # Use alpha matting for better edge detection (no white halos)
    cleaned = remove(original, session=session, alpha_matting=True, alpha_matting_foreground_threshold=240, alpha_matting_background_threshold=10)
    
    # 2. Find Bounding Box (Non-transparent pixels)
    bbox = cleaned.getbbox()
    if not bbox:
        return cleaned.resize(target_size)
    
    # Crop to content
    content = cleaned.crop(bbox)
    w, h = content.size
    
    # 3. Create Canvas (128x128)
    canvas = Image.new("RGBA", target_size, (0, 0, 0, 0))
    
    # 4. Center on Feet logic
    # We want the bottom-center of the content to align with the bottom-center of the tile
    # Target X center: 64
    # Target Bottom Y: 115 (leave some margin)
    
    target_x = (target_size[0] - w) // 2
    target_y = target_size[1] - h - 10 # 10px from bottom margin
    
    # Paste content onto canvas
    canvas.paste(content, (target_x, target_y), content)
    
    return canvas

def process_gronk():
    print("🧛‍♂️ STARTING GRONK PRO PROCESSING...")
    # Expected ordering based on file names or hardcoded list to ensure animation sync
    # We need 19 images.
    # Logic: Sort files properly. 
    files = sorted([f for f in os.listdir(GRONK_INPUT_DIR) if f.lower().endswith(".png")])
    
    # Mapping based on typical file naming (assuming 1..19 or similar)
    # If filenames are random, we might have issues. Let's assume user provided them in order or standard naming.
    # Previous script used sorted list.
    
    # Create big spritesheet
    sheet = Image.new("RGBA", (SHEET_COLS * FRAME_SIZE, SHEET_ROWS * FRAME_SIZE))
    
    # Layout Mapping (Row -> Image Indices)
    # Row 0 (Down): 0,1,2,3
    # Row 1 (Right): 4,5,6,7
    # Row 2 (Up): 8,9,10,11
    # Row 3 (Vape): 12,13,14,15
    # Row 4 (Idle): 16,17,18 (19th unused)
    
    for i, filename in enumerate(files):
        if i >= 19: break
        
        input_path = os.path.join(GRONK_INPUT_DIR, filename)
        processed_frame = clean_and_center(input_path)
        
        # Calculate grid position
        # Frames 0-15 map nicely. Frames 16-18 go to Row 4.
        row = i // SHEET_COLS
        col = i % SHEET_COLS
        
        if row >= SHEET_ROWS: break
        
        x = col * FRAME_SIZE
        y = row * FRAME_SIZE
        
        sheet.paste(processed_frame, (x, y))
        
    output_dir = os.path.dirname(GRONK_OUTPUT_FILE)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    sheet.save(GRONK_OUTPUT_FILE)
    print(f"✅ Gronk Pro Sheet Saved: {GRONK_OUTPUT_FILE}")

def process_cannabis():
    print("🌿 CLEANING CANNABIS...")
    for path in CANNABIS_PATHS:
        full_path = path # Assuming relative to cwd
        if os.path.exists(full_path):
            try:
                img = Image.open(full_path)
                # Rembg with aggressive cutout
                out = remove(img, session=session, alpha_matting=True)
                out.save(full_path)
                print(f"✨ Cleaned: {path}")
            except Exception as e:
                print(f"❌ Error cleaning {path}: {e}")
        else:
            print(f"⚠️ Not found: {path}")

if __name__ == "__main__":
    process_cannabis()
    try:
        process_gronk()
    except Exception as e:
        print(f"❌ Gronk Error: {e}")