import os
import struct
from PIL import Image

# Path to folder with files
folder = 'images/TITLE'
palette_file = os.path.join(folder, 'TITLE.AP_01.pal')

# Load palette
with open(palette_file, 'rb') as f:
    data = f.read()
palette = [tuple(data[i:i+3]) for i in range(0, len(data), 3)]

def find_palette_index(color):
    try:
        return palette.index(color)
    except ValueError:
        return 0

def encode_image_to_rle(image_path, output_path):
    image = Image.open(image_path).convert('RGB')
    width, height = image.size
    pixels = image.load()

    encoded_data = bytearray()
    encoded_data += struct.pack('<HH', width, height)
    encoded_body = bytearray()

    for y in range(height):
        black_count = 0
        color_indices = []

        # Count black pixels at the beginning
        x = 0
        while x < width and pixels[x, y] == (0, 0, 0):
            black_count += 1
            x += 1

        # Skip trailing black pixels on the right
        end = width - 1
        while end >= x and pixels[end, y] == (0, 0, 0):
            end -= 1

        # Collect remaining pixel values
        while x <= end:
            pixel = pixels[x, y]
            color_indices.append(find_palette_index(pixel))
            x += 1

        # Split into blocks of 255
        remaining = color_indices[:]
        first_chunk = True
        while len(remaining) > 0:
            chunk = remaining[:255]
            remaining = remaining[255:]
            encoded_body += struct.pack('<BB', black_count if first_chunk else 0, len(chunk))
            encoded_body += bytes(chunk)
            first_chunk = False

    encoded_data += struct.pack('<H', len(encoded_body))
    encoded_data += encoded_body

    with open(output_path, 'wb') as f:
        f.write(encoded_data)

# Process all PNG files in folder
for filename in os.listdir(folder):
    if filename.lower().endswith('.png'):
        image_path = os.path.join(folder, filename)
        base, _ = os.path.splitext(filename)
        output_path = os.path.join(folder, base)
        encode_image_to_rle(image_path, output_path)

print("All PNG files successfully packed into RLE format.")