import struct
import os
import sys
from PIL import Image

def import_texture_to_pack(pack_filename, texture_filename, output_pack_filename=None):
    """
    Заменяет пиксельные данные текстуры в pack-файле на данные из PNG файла
    """
    
    if output_pack_filename is None:
        output_pack_filename = pack_filename
    
    # Получаем ID текстуры из имени файла
    texture_id = extract_texture_id(texture_filename)
    if texture_id is None:
        print("Не удалось определить ID текстуры из имени файла!")
        return False
    
    print(f"Импортируем текстуру ID: {texture_id} из {texture_filename}")
    
    # Читаем pack-файл
    with open(pack_filename, 'rb') as f:
        pack_data = bytearray(f.read())
    
    # Находим смещения текстур
    tx_offsets = find_texture_offsets(pack_data)
    if texture_id >= len(tx_offsets) - 1:
        print(f"Ошибка: текстура с ID {texture_id} не найдена в pack-файле!")
        return False
    
    # Находим позицию пиксельных данных в pack-файле
    pixel_data_pos = find_pixel_data_position(pack_data, tx_offsets, texture_id)
    if pixel_data_pos is None:
        print(f"Не удалось найти позицию пиксельных данных для текстуры {texture_id}!")
        return False
    
    # Получаем информацию о размерах текстуры из pack-файла
    texture_info = get_texture_info(pack_data, tx_offsets, texture_id)
    if texture_info is None:
        print(f"Не удалось получить информацию о текстуре {texture_id}!")
        return False
    
    image_w, image_h, image_size = texture_info
    
    # Загружаем PNG файл
    try:
        img = Image.open(texture_filename)
    except Exception as e:
        print(f"Ошибка загрузки PNG файла: {e}")
        return False
    
    # Проверяем размеры
    if img.width != image_w or img.height != image_h:
        print(f"Ошибка: размеры не совпадают!")
        print(f"Ожидается: {image_w}x{image_h}")
        print(f"PNG файл: {img.width}x{img.height}")
        return False
    
    # Конвертируем в индексированное изображение если нужно
    if img.mode != 'P':
        print("Предупреждение: PNG не в индексированном режиме. Конвертируем...")
        img = img.convert('P')
    
    # Извлекаем пиксельные данные
    pixel_data = []
    for y in range(image_h):
        for x in range(image_w):
            try:
                pixel_index = img.getpixel((x, y))
                pixel_data.append(pixel_index & 0xFF)  #确保是字节
            except:
                pixel_data.append(0)
    
    # Проверяем размер данных
    expected_pixel_count = image_w * image_h
    actual_pixel_count = len(pixel_data)
    
    if actual_pixel_count != expected_pixel_count:
        print(f"Ошибка: количество пикселей не совпадает!")
        print(f"Ожидается: {expected_pixel_count}")
        print(f"Получено: {actual_pixel_count}")
        return False
    
    # Заменяем данные в pack-файле
    bytes_replaced = replace_pixel_data(pack_data, pixel_data_pos, pixel_data, image_size - 12)
    
    # Сохраняем измененный pack-файл
    try:
        with open(output_pack_filename, 'wb') as f:
            f.write(pack_data)
        print(f"Успешно заменено {bytes_replaced} байт!")
        print(f"Сохранено в: {output_pack_filename}")
        return True
    except Exception as e:
        print(f"Ошибка сохранения файла: {e}")
        return False

def extract_texture_id(filename):
    """Извлекает ID текстуры из имени файла"""
    basename = os.path.basename(filename)
    # Ищем паттерн texture_XXX.png
    if basename.startswith('texture_') and basename.endswith('.png'):
        try:
            id_str = basename[8:-4]  # Убираем 'texture_' и '.png'
            return int(id_str)
        except:
            pass
    return None

def find_texture_offsets(pack_data):
    """Находит смещения текстур в pack-файле"""
    tx_offsets = []
    first_addr = struct.unpack('<I', pack_data[0:4])[0]
    tx_offsets.append(first_addr)
    
    reader = 4
    while reader < first_addr:
        addr = struct.unpack('<I', pack_data[reader:reader+4])[0]
        if addr == 0:
            break
        tx_offsets.append(addr)
        reader += 4
    
    return tx_offsets

def find_pixel_data_position(pack_data, tx_offsets, texture_id):
    """Находит позицию пиксельных данных для указанной текстуры"""
    try:
        reader = tx_offsets[texture_id]
        
        # Пропускаем режимы CLUT и изображения
        reader += 8
        
        # Читаем размер CLUT и пропускаем CLUT данные
        clut_size = struct.unpack('<I', pack_data[reader:reader+4])[0]
        reader += clut_size  # Пропускаем весь CLUT блок
        
        # Пропускаем заголовок изображения
        reader += 12
        
        # Теперь reader указывает на начало пиксельных данных
        return reader
        
    except Exception as e:
        print(f"Ошибка поиска позиции данных: {e}")
        return None

def get_texture_info(pack_data, tx_offsets, texture_id):
    """Получает информацию о размерах текстуры"""
    try:
        reader = tx_offsets[texture_id]
        
        # Пропускаем режимы
        reader += 8
        
        # Пропускаем CLUT
        clut_size = struct.unpack('<I', pack_data[reader:reader+4])[0]
        reader += clut_size
        
        # Читаем информацию об изображении
        image_size = struct.unpack('<I', pack_data[reader:reader+4])[0]
        image_x = struct.unpack('<H', pack_data[reader+4:reader+6])[0] * 2
        image_y = struct.unpack('<H', pack_data[reader+6:reader+8])[0] - 256
        image_w = struct.unpack('<H', pack_data[reader+8:reader+10])[0] * 2
        image_h = struct.unpack('<H', pack_data[reader+10:reader+12])[0]
        
        return image_w, image_h, image_size
        
    except Exception as e:
        print(f"Ошибка получения информации о текстуре: {e}")
        return None

def replace_pixel_data(pack_data, start_pos, new_pixel_data, expected_size):
    """Заменяет пиксельные данные в pack-файле"""
    bytes_replaced = 0
    
    for i, pixel_byte in enumerate(new_pixel_data):
        if start_pos + i < len(pack_data) and i < expected_size:
            pack_data[start_pos + i] = pixel_byte
            bytes_replaced += 1
        else:
            break
    
    return bytes_replaced

def main():
    if len(sys.argv) < 3:
        print("Использование: python import.py <pack_file> <texture_file> [output_file]")
        print("Пример: python import.py game.pack textures/texture_005.png")
        print("Пример: python import.py game.pack textures/texture_010.png game_modified.pack")
        sys.exit(1)
    
    pack_filename = sys.argv[1]
    texture_filename = sys.argv[2]
    output_filename = sys.argv[3] if len(sys.argv) > 3 else pack_filename
    
    if not os.path.exists(pack_filename):
        print(f"Pack-файл не найден: {pack_filename}")
        sys.exit(1)
    
    if not os.path.exists(texture_filename):
        print(f"Текстура не найдена: {texture_filename}")
        sys.exit(1)
    
    print(f"Pack файл: {pack_filename}")
    print(f"Текстура: {texture_filename}")
    print(f"Выходной файл: {output_filename}")
    print("-" * 50)
    
    success = import_texture_to_pack(pack_filename, texture_filename, output_filename)
    
    if success:
        print("Импорт завершен успешно!")
    else:
        print("Импорт завершен с ошибками!")
        sys.exit(1)

if __name__ == "__main__":
    main()