import os
import xml.etree.ElementTree as ET
from PIL import Image, ImageChops
import numpy as np

def generate_water_animation():
    # Paths
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    src_path = os.path.join(base_dir, 'assets', 'grounds', 'water.png')
    out_dir = os.path.join(base_dir, 'assets', 'maps', 'tilesets')
    out_img_path = os.path.join(out_dir, 'Water_Animated.png')
    out_tsx_path = os.path.join(out_dir, 'Water_Animated.tsx')

    if not os.path.exists(src_path):
        print(f"Error: Source image not found at {src_path}")
        return

    print(f"Processing: {src_path}")
    
    # Load Image
    img = Image.open(src_path).convert('RGBA')
    width, height = img.size
    tile_w, tile_h = 32, 32
    
    frames = []
    
    # Frame 1: Original
    frames.append(img)
    
    # Function to add subtle glimmer (No shifting!)
    def create_glimmer(image, intensity=30, seed=0):
        # Convert to numpy
        arr = np.array(image, dtype=np.int16) # Use int16 to avoid overflow
        
        # Create random noise
        np.random.seed(seed)
        noise = np.random.randint(-intensity, intensity, (height, width), dtype=np.int16)
        
        # Apply noise only to RGB channels, leave Alpha alone
        # arr shape is (H, W, 4)
        
        # Only apply where alpha > 0 (existing content)
        alpha = arr[:, :, 3]
        mask = alpha > 0
        
        for c in range(3): # R, G, B
            channel = arr[:, :, c]
            # Add noise
            channel[mask] = channel[mask] + noise[mask]
            # Clip
            arr[:, :, c] = np.clip(channel, 0, 255)
            
        return Image.fromarray(arr.astype(np.uint8))

    # Generate 3 variations
    print("Generating Glimmer Frames...")
    f2 = create_glimmer(img, intensity=20, seed=42)
    f3 = create_glimmer(img, intensity=20, seed=101)
    f4 = create_glimmer(img, intensity=20, seed=777)
    
    frames.append(f2)
    frames.append(f3)
    frames.append(f4)
    
    # Combine into vertical strip
    total_height = height * 4
    combined = Image.new('RGBA', (width, total_height))
    
    for i, f in enumerate(frames):
        combined.paste(f, (0, i * height))
        
    # Save Image
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    combined.save(out_img_path)
    print(f"Saved animated spritesheet: {out_img_path}")
    
    # Generate TSX
    cols = width // tile_w
    rows = height // tile_h
    total_tiles_per_frame = cols * rows
    total_tiles_all = total_tiles_per_frame * 4
    
    root = ET.Element("tileset")
    root.set("version", "1.10")
    root.set("tiledversion", "1.11.0")
    root.set("name", "Water_Animated")
    root.set("tilewidth", str(tile_w))
    root.set("tileheight", str(tile_h))
    root.set("tilecount", str(total_tiles_all))
    root.set("columns", str(cols))
    
    image_node = ET.SubElement(root, "image")
    image_node.set("source", "Water_Animated.png")
    image_node.set("width", str(width))
    image_node.set("height", str(total_height))
    
    # Generate Animations for the first frame's tiles (0 to total_tiles_per_frame-1)
    for i in range(total_tiles_per_frame):
        tile_node = ET.SubElement(root, "tile")
        tile_node.set("id", str(i))
        
        anim_node = ET.SubElement(tile_node, "animation")
        
        # 4 Frames loop
        # Frame 1 (Original)
        f1 = ET.SubElement(anim_node, "frame")
        f1.set("tileid", str(i))
        f1.set("duration", "150")
        
        # Frame 2
        f2 = ET.SubElement(anim_node, "frame")
        f2.set("tileid", str(i + total_tiles_per_frame))
        f2.set("duration", "150")
        
        # Frame 3
        f3 = ET.SubElement(anim_node, "frame")
        f3.set("tileid", str(i + 2 * total_tiles_per_frame))
        f3.set("duration", "150")
        
        # Frame 4
        f4 = ET.SubElement(anim_node, "frame")
        f4.set("tileid", str(i + 3 * total_tiles_per_frame))
        f4.set("duration", "150")
        
    # Formatting
    tree = ET.ElementTree(root)
    ET.indent(tree, space=" ", level=0)
    tree.write(out_tsx_path, encoding='UTF-8', xml_declaration=True)
    print(f"Saved TSX definition: {out_tsx_path}")

if __name__ == "__main__":
    generate_water_animation()