import os
from PIL import Image

fragment_width = 16
fragment_height = 16

def raw_to_png(raw_file, output_png, palette_offset=0x10, palette_size=0x300, data_offset=0x310):
    if not os.path.exists(raw_file):
        print("Error: RAW file not found.")
        return

    with open(raw_file, 'rb') as f:
        # Read the first 4 bytes
        first_byte = f.read(1)[0]
        second_byte = f.read(1)[0]
        third_byte = f.read(1)[0]
        fourth_byte = f.read(1)[0]

        # Read palette
        f.seek(palette_offset)
        palette_data = f.read(palette_size)
        if len(palette_data) != palette_size:
            print("Error: Unable to read palette from RAW file.")
            return

        # Read image data
        f.seek(data_offset)
        raw_data = f.read()

    # Compute full image dimensions
    image_width = fragment_width * first_byte
    image_height = fragment_height * second_byte

    # Build indexed image
    image = Image.frombytes('P', (fragment_width, len(raw_data) // fragment_width), raw_data)
    image.putpalette(palette_data)

    # Split into fragments
    fragments = split_image(image)

    # Arrange fragments
    if first_byte != fourth_byte:
        print(f"{os.path.basename(raw_file)}: non-standard layout (1st ≠ 4th byte)")
        result = build_custom_layout(fragments, first_byte, second_byte)
    else:
        result = build_standard_layout(fragments, first_byte, second_byte)

    result.save(output_png)
    print(f"Created: {output_png}")


def split_image(image):
    fragments = []
    img_width, img_height = image.size
    for y in range(0, img_height, fragment_height):
        for x in range(0, img_width, fragment_width):
            fragments.append(image.crop((x, y, x + fragment_width, y + fragment_height)))
    return fragments


def build_standard_layout(fragments, cols, rows):
    """Standard layout: fragments placed row by row."""
    result = Image.new('RGB', (fragment_width * cols, fragment_height * rows))
    idx = 0
    for y in range(rows):
        for x in range(cols):
            if idx < len(fragments):
                result.paste(fragments[idx], (x * fragment_width, y * fragment_height))
                idx += 1
    return result


def build_custom_layout(fragments, first_byte, second_byte):
    """
    Custom layout when 1st byte != 4th byte:
    - Left section: 20 × (second_byte)
    - Right section: (first_byte - 20) × (second_byte), placed left-to-right
    """
    left_cols = 20
    right_cols = first_byte - 20
    total_rows = second_byte

    total_width = fragment_width * first_byte
    total_height = fragment_height * total_rows
    result = Image.new('RGB', (total_width, total_height))

    idx = 0

    # Left block
    for y in range(total_rows):
        for x in range(left_cols):
            if idx >= len(fragments):
                return result
            result.paste(fragments[idx], (x * fragment_width, y * fragment_height))
            idx += 1

    # Right block
    for y in range(total_rows):
        for x in range(left_cols, left_cols + right_cols):
            if idx >= len(fragments):
                return result
            result.paste(fragments[idx], (x * fragment_width, y * fragment_height))
            idx += 1

    return result


# Process all RAW files in directory
directory = 'images'
raw_files = [file for file in os.listdir(directory) if file.endswith('1.raw')]

for raw_file in raw_files:
    output_png = os.path.splitext(raw_file)[0] + '.png'
    raw_to_png(os.path.join(directory, raw_file), os.path.join(directory, output_png))