import cv2
import numpy as np
import os

def remove_pink_and_resize(image_path, target_size=(256, 256)):
    img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
    if img is None:
        print(f"Failed to load {image_path}")
        return None

    # Handle Alpha channel
    if img.shape[2] == 3:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    # Remove Pink (#FF00FF and range)
    # Convert to HSV might be safer, but exact #FF00FF is (255, 0, 255) in BGR
    # Let's target the magenta color.
    # BGR: 255, 0, 255
    lower_pink = np.array([250, 0, 250])
    upper_pink = np.array([255, 10, 255])
    
    mask = cv2.inRange(img[:, :, :3], lower_pink, upper_pink)
    
    # Set alpha to 0 where mask is true
    img[mask > 0] = [0, 0, 0, 0]

    # Resize to fit 256x256
    # For a tileset grid, we usually want it to fill the tile or be centered?
    # User said "polagal čez plato". Let's simply resize to 256x256 for simplicity of the grid.
    # Distortion might occur but these are nature assets (grass).
    # Ideally preserve aspect ratio but for "Tileset Image" grid alignment, 256x256 is safest.
    resized = cv2.resize(img, target_size, interpolation=cv2.INTER_AREA)
    
    return resized

def create_full_tileset():
    # Sources
    grass_path = 'main/assets/tla_trava_tekstura.png'
    water_path = 'main/assets/stream_water.png'
    high_grass_path = 'assets/DEMO_FAZA1/Vegetation/visoka_trava.png'
    dense_grass_path = 'assets/DEMO_FAZA1/Vegetation/grass_cluster_dense.png'
    
    # Output
    output_png = 'main/assets/ground_tileset.png' # We update the existing one
    
    # Process
    # 1. Base Grass (assumed clean)
    grass = cv2.imread(grass_path, cv2.IMREAD_UNCHANGED)
    if grass.shape[2] == 3: grass = cv2.cvtColor(grass, cv2.COLOR_BGR2BGRA)
    grass = cv2.resize(grass, (256, 256))

    # 2. Water (assumed clean)
    water = cv2.imread(water_path, cv2.IMREAD_UNCHANGED)
    if water.shape[2] == 3: water = cv2.cvtColor(water, cv2.COLOR_BGR2BGRA)
    water = cv2.resize(water, (256, 256))

    # 3. High Grass (Clean Pink)
    high_grass = remove_pink_and_resize(high_grass_path)
    
    # 4. Dense Grass (Clean Pink)
    dense_grass = remove_pink_and_resize(dense_grass_path)

    if high_grass is None or dense_grass is None:
        print("Error processing grass images.")
        return

    # Combine: [Grass] [Water] [HighGrass] [DenseGrass]
    tileset = np.hstack((grass, water, high_grass, dense_grass))
    
    # Save
    cv2.imwrite(output_png, tileset)
    print(f"Saved combined tileset to {output_png} with shape {tileset.shape}")

    # Generate .tsx content
    tsx_content = f"""<?xml version="1.0" encoding="UTF-8"?>
<tileset version="1.10" tiledversion="1.10.2" name="ground_tileset" tilewidth="256" tileheight="256" tilecount="4" columns="4">
 <image source="ground_tileset.png" width="{tileset.shape[1]}" height="{tileset.shape[0]}"/>
</tileset>
"""
    with open('main/assets/ground_tileset.tsx', 'w') as f:
        f.write(tsx_content)
    print("Saved ground_tileset.tsx")

if __name__ == "__main__":
    create_full_tileset()