import struct
import sys
import os
from PIL import Image

def read_u32(f):
    data = f.read(4)
    if len(data) < 4:
        return None
    return struct.unpack('<I', data)[0]

def read_u16(f):
    data = f.read(2)
    if len(data) < 2:
        return None
    return struct.unpack('<H', data)[0]

def rgb1555_to_rgb(color):
    a = (color >> 15) & 0x1
    r = ((color >> 10) & 0x1F) << 3
    g = ((color >> 5) & 0x1F) << 3
    b = (color & 0x1F) << 3
    return (r, g, b, 255 if a else 0) if a else (r, g, b)

def process_texture_file(input_file):
    with open(input_file, 'rb') as f:
        # Читаем адреса текстурных блоков
        texture_addresses = []
        first_address = read_u32(f)
        if first_address is None:
            print("Error: Invalid file format (cannot read first address)")
            return
            
        texture_addresses.append(first_address)
        print(f"First texture block at: 0x{first_address:08X}")
        
        # Читаем адреса до достижения позиции first_address
        while f.tell() < first_address:
            addr = read_u32(f)
            if addr is None:
                break
            texture_addresses.append(addr)
        
        print(f"Found {len(texture_addresses)} texture blocks")
        
        # Создаем изображение 2048x512
        img = Image.new('RGBA', (2048, 512))
        pixels = img.load()
        
        for addr in texture_addresses:
            if addr >= os.path.getsize(input_file):
                print(f"Warning: Invalid texture address 0x{addr:08X}, skipping")
                continue
                
            f.seek(addr)
            
            # Читаем параметры текстуры
            clut_mode = read_u32(f)
            image_mode = read_u32(f)
            
            if clut_mode is None or image_mode is None:
                print(f"Warning: Invalid texture header at 0x{addr:08X}, skipping")
                continue
            
            print(f"Processing texture at 0x{addr:08X}, clut_mode={clut_mode}, image_mode={image_mode}")
            
            # Обрабатываем CLUT чанк
            clut_size = read_u32(f)
            if clut_size is None:
                print(f"Warning: Invalid CLUT size at 0x{addr:08X}, skipping")
                continue
                
            clut_x = read_u16(f) * 2  # Умножаем только позицию на 2
            clut_y = read_u16(f)
            clut_w = read_u16(f)
            clut_h = read_u16(f) 
            
            if None in (clut_x, clut_y, clut_w, clut_h):
                print(f"Warning: Invalid CLUT header at 0x{addr:08X}, skipping")
                continue
            
            print(f"CLUT: pos=({clut_x},{clut_y}), size=({clut_w},{clut_h})")
            
            clut_data_size = clut_size - 12
            clut_data = f.read(clut_data_size)
            if len(clut_data) < clut_data_size:
                print(f"Warning: CLUT data truncated at 0x{addr:08X}, skipping")
                continue
            
            # Создаем CLUT палитру
            clut = []
            if clut_mode == 16:
                for i in range(0, min(len(clut_data), clut_w*clut_h*2), 2):
                    color = struct.unpack('<H', clut_data[i:i+2])[0]
                    rgb = rgb1555_to_rgb(color)
                    clut.append(rgb)
                    
                    # Рисуем CLUT (x-координата уже умножена на 2)
                    px = clut_x + (i//2)
                    py = clut_y + ((i//2) // clut_w)
                    if 0 <= px < 2048 and 0 <= py < 512:
                        pixels[px, py] = rgb
            
            # Обрабатываем текстуру
            tex_size = read_u32(f)
            if tex_size is None:
                print(f"Warning: Invalid texture size at 0x{addr:08X}, skipping")
                continue
                
            tex_x = read_u16(f) * 2  # Умножаем только позицию на 2
            tex_y = read_u16(f)
            tex_w = read_u16(f) * 2
            tex_h = read_u16(f)
            
            if None in (tex_x, tex_y, tex_w, tex_h):
                print(f"Warning: Invalid texture header at 0x{addr:08X}, skipping")
                continue
            
            print(f"Texture: pos=({tex_x},{tex_y}), size=({tex_w},{tex_h})")
            
            tex_data_size = tex_size - 12
            tex_data = f.read(tex_data_size)
            if len(tex_data) < tex_data_size:
                print(f"Warning: Texture data truncated at 0x{addr:08X}, skipping")
                continue
            
            if image_mode == 9 and clut_mode == 16 and clut:
                for y in range(tex_h):
                    for x in range(tex_w):
                        idx = y * tex_w + x
                        if idx >= len(tex_data):
                            continue
                        color_idx = tex_data[idx]
                        if color_idx >= len(clut):
                            continue
                        # Координаты текстуры (x-позиция уже умножена на 2)
                        px = tex_x + x
                        py = tex_y + y
                        if 0 <= px < 2048 and 0 <= py < 512:
                            pixels[px, py] = clut[color_idx]
        
        # Сохраняем PNG
        output_file = os.path.splitext(input_file)[0] + '.png'
        img.save(output_file)
        print(f"Image saved to {output_file}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python texture_converter.py <input_file>")
        sys.exit(1)
    
    input_file = sys.argv[1]
    if not os.path.exists(input_file):
        print(f"Error: File {input_file} not found")
        sys.exit(1)
    
    process_texture_file(input_file)