bl_info = {
    "name": "Persona2 PSX Scene Exporter",
    "author": "Sergey Shemet",
    "version": (1, 0),
    "blender": (4, 5, 0),
    "location": "File > Export",
    "description": "Export Persona 2 Eternal Punishment PSX 3D scenes (.p2e)",
    "category": "Import-Export",
}

import bpy
import struct
import math
import binascii
from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty, BoolProperty

class ExportP2E(bpy.types.Operator, ExportHelper):
    bl_idname = "export_scene.p2e"
    bl_label = "Export Persona2 PSX Scene"
    bl_description = "Export Persona 2 Eternal Punishment PSX 3D scenes"
    bl_options = {'REGISTER', 'UNDO', 'PRESET'}
    
    filename_ext = ".p2e"
    
    filter_glob: StringProperty(
        default="*.p2e",
        options={'HIDDEN'},
        maxlen=255,
    )
    
    def execute(self, context):
        # 1. Сбор объектов из коллекции PSX_Export
        export_collection = bpy.data.collections.get("PSX_Export")
        if not export_collection:
            self.report({'ERROR'}, "Collection 'PSX_Export' not found")
            return {'CANCELLED'}
        
        scene_data = self.prepare_scene_data(export_collection)
        if scene_data is None:  # Если prepare_scene_data провалился
            return {'CANCELLED'}
        
        self.write_p2e_file(self.filepath, scene_data)
        self.report({'INFO'}, f"Exported {len(scene_data['objects'])} objects")
        return {'FINISHED'}
    
    def prepare_scene_data(self, export_collection):
        """Подготовка всех данных сцены"""
        if not export_collection:
            self.report({'ERROR'}, "Export collection is None")
            return None

        objects_data = []
        global_vertices = []
        global_vertex_offset = 0

        for obj in export_collection.objects:
            if obj.type != 'MESH':
                continue

            # Данные объекта
            obj_data = {
                'id': len(objects_data),
                'custom_props': {k: v for k, v in obj.items() if k not in ['_RNA_UI']},
                'polygon_count': 0,
                'chunk_data': bytearray(),
                'local_vertices': [],
                'vertex_map': {}
            }

            # Обработка вершин меша
            mesh = obj.data
            mesh.calc_loop_triangles()  # Убедимся, что полигональные данные актуальны

            local_vertices = []
            for vert in mesh.vertices:
                world_co = obj.matrix_world @ vert.co
                psx_vert = (
                    int(world_co.x * 100),
                    -int(world_co.z * 100),
                    int(world_co.y * 100),
                    0
                )
                local_vertices.append(psx_vert)

            # Уникальные вершины
            unique_vertices = []
            vertex_map = {}
            for i, vert in enumerate(local_vertices):
                found = False
                for j, unique_vert in enumerate(unique_vertices):
                    if vert == unique_vert:
                        vertex_map[i] = j
                        found = True
                        break
                if not found:
                    vertex_map[i] = len(unique_vertices)
                    unique_vertices.append(vert)

            for vert in unique_vertices:
                global_vertices.append(vert)
                obj_data['local_vertices'].append(vert)
            obj_data['vertex_map'] = vertex_map

            # === ПРОВЕРКА UV СЛОЯ ===
            if not mesh.uv_layers.active:
                self.report({'ERROR'}, f"Object '{obj.name}' has no active UV map. Please unwrap UVs.")
                return None

            uv_layer = mesh.uv_layers.active.data
            if len(uv_layer) != len(mesh.loops):
                self.report({'WARNING'}, f"UV layer size mismatch for '{obj.name}': {len(uv_layer)} UVs vs {len(mesh.loops)} loops")

            # === ОБРАБОТКА ПОЛИГОНОВ ===
            for poly in mesh.polygons:
                vertex_count = len(poly.vertices)
                if vertex_count not in (3, 4):
                    self.report({'ERROR'}, f"Polygon in '{obj.name}' has {vertex_count} vertices. Only triangles (3) or quads (4) allowed.")
                    return None

                # Собираем UV координаты с проверкой индексов
                uvs = []
                for loop_idx in poly.loop_indices:
                    if loop_idx >= len(uv_layer):
                        self.report({'ERROR'}, f"Loop index {loop_idx} out of range in '{obj.name}'. UV data is broken or inconsistent.")
                        return None
                    uv = uv_layer[loop_idx].uv
                    uvs.append((uv.x, uv.y))

                # Получаем индексы вершин
                vertex_ids = [obj_data['vertex_map'][v] + global_vertex_offset for v in poly.vertices]

                # Получаем цвет затенения
                shade = obj.get("shade", (128, 128, 128))

                # Формируем чанк
                if vertex_count == 3:
                    chunk = self.create_textured_tri_chunk(uvs, shade)
                else:  # 4
                    chunk = self.create_textured_quad_chunk(uvs, shade)

                # тестовый разворот вершин
                vertex_ids = vertex_ids[::-1]
                

                # Меняем 3 и 4 вершины местами для квадов (если 4 вершины)
                if vertex_count == 4:
                    vertex_ids[2], vertex_ids[3] = vertex_ids[3], vertex_ids[2]

                # Добавляем индексы в чанк
                for vid in vertex_ids:
                    chunk += struct.pack('<I', vid)

                obj_data['chunk_data'] += chunk
                obj_data['polygon_count'] += 1

            # Обновляем смещение
            global_vertex_offset += len(unique_vertices)
            objects_data.append(obj_data)

        return {
            'objects': objects_data,
            'vertices': global_vertices,
            'vertex_count': len(global_vertices),
            'object_count': len(objects_data)
        }

    
    def calc_tex_params(self, uvs):
        """Расчет CLUT и TexPage по UV координатам"""
        # Среднее U координат
        avg_u = sum(u for u, v in uvs) / len(uvs)
        

        # Расчет CLUT (0-7)
        # CLUT - это индекс в палитре текстур, нам нужно определить в каком из блоков 1/8 текстуры находится среднее U
        #   Блоки CLUT расположены в VRAM линиями, начиная с 0x480. Поэтому нам нужно просто смещаться по Y по 1 для каждого блока
        clut_index = min(int(avg_u * 8), 7)
        clut = (clut_index + 480) << 6  # CLUT хранится как int16, YYYYYYYYYYXXXXX. Так как X всегда 0, то просто смещаем Y на 6 бит
        
        
        # Расчет TexPage
        tex_page_raw = int(avg_u * 4)
        tex_page = ((tex_page_raw & 15) + 10) | 16 | 128 #default p2 texpage shift (10) & enabling lower page (5th bit) & 8-bit mode (8th bit)
        
        
        # Расчет X координат - нам нужно сместить x относительно текстурной страницы
        tex_page_u_offset = (tex_page_raw & 3) * 0.25 # Коэффициент текстурной страницы в 0...1
        xs = []
        
        for u, v in uvs:
            # X = (U - смещение страницы) * 512.
            # 512 - это ширина всей текстуры. Отнимаем коэффициент от X, умножая на 512,
            #       чтобы получить положение X относительно начала страницы в координатах всей текстуры
            x = int((u - tex_page_u_offset) * 512) & 0xFF
            xs.append(x)
        
        return {
            'clut': clut,
            'tex_page': tex_page,
            'xs': xs
        }
    
    def create_textured_tri_chunk(self, uvs, shade):
        """Создание чанка для текстурированного треугольника"""
        params = self.calc_tex_params(uvs)
        
        chunk = bytearray()
        
        # Заглушка (4 байта)
        chunk += b'\x5c\x64\x61\x07'
        
        # Цвет и команда GPU
        chunk += struct.pack('<BBBB', shade[0], shade[1], shade[2], 0x24)
        
        # Вершина 1
        chunk += b'\x70\x5c\x30\x30'
        chunk += struct.pack('<BBH', 
            params['xs'][0],       # X1
            int((1 - uvs[0][1]) * 255),  # Y1 = V1 * 256
            params['clut']         # CLUT
        )
        
        # Вершина 2
        chunk += b'\x30\x32\x34\x2e'
        chunk += struct.pack('<BBH',
            params['xs'][1],       # X2
            int((1 - uvs[2][1]) * 255),  # Y2
            params['tex_page']     # TexPage
        )
        
        # Вершина 3
        chunk += b'\x1e\xc1\x00\x09'
        chunk += struct.pack('<BBH',
            params['xs'][2],       # X3
            int((1 - uvs[3][1]) * 255),  # Y3
            0                      # Заглушка
        )
        
        # Дублируем массив 2 раза (технический нюанс)
        full_chunk = chunk * 2
        
        # Добавляем заголовок чанка
        header = struct.pack('<BBH', 0x2, 0x13, 0)
        return header + full_chunk
    
    def create_textured_quad_chunk(self, uvs, shade):
        """Создание чанка для текстурированного квадра"""
        params = self.calc_tex_params(uvs)
        
        chunk = bytearray()

        # Магические байты
        chunk += b'\x5c\x64\x61\x09'
        
        # Цвет и команда
        chunk += struct.pack('<BBBB', shade[0], shade[1], shade[2], 0x2C)
        
        # Вершина 1
        chunk += b'\x70\x5c\x30\x30'
        chunk += struct.pack('<BBH',
            params['xs'][3],       # X1
            # params['xs'][0],       # X1
            int((1 - uvs[3][1]) * 255),  # Y1
            # int((1 - uvs[0][1]) * 255),  # Y1
            params['clut']         # CLUT
        )

        # Вершина 2
        chunk += b'\x30\x32\x34\x2e'
        chunk += struct.pack('<BBH',
            params['xs'][2],       # X2
            # params['xs'][1],       # X2
            int((1 - uvs[2][1]) * 255),  # Y2
            # int((1 - uvs[1][1]) * 255),  # Y2
            params['tex_page']     # TexPage
        )
        
        # Вершина 4 
        chunk += b'\x1e\xC1\x00\x09'
        chunk += struct.pack('<BBBB',
            params['xs'][0],       # X4
            # params['xs'][3],       # X4
            int((1 - uvs[0][1]) * 255),  # Y4 = V4 * 256
            # int((1 - uvs[3][1]) * 255),  # Y4 = V4 * 256
            84, 12                 # Заглушка
        )
        
        # Вершина 3
        chunk += b'\x5c\x64\x61\x07'
        chunk += struct.pack('<BBBB',
            params['xs'][1],       # X3
            # params['xs'][2],       # X3
            int((1 - uvs[1][1]) * 255),  # Y3 = V3 * 256
            # int((1 - uvs[2][1]) * 255),  # Y3 = V3 * 256
            128, 36                # Заглушка
        )
        
        # Дублируем массив 2 раза
        full_chunk = chunk * 2
        
        # Добавляем заголовок
        header = struct.pack('<BBH', 0x6, 0x18, 0)
        return header + full_chunk
    
    def write_p2e_file(self, filepath, scene_data):
        """Запись .p2e файла"""
        with open(filepath, 'wb') as f:
                    # 1. sourceName (16 байт)
            f.write(b'Hello, blya.\x00\x00\x00\x00')
            
            # 2. Заголовок сцены (пока заглушки)
            header_pos = f.tell()  # Должно быть 0x10
            f.write(struct.pack('<HHII',
                0,  # objectsCnt - заполним позже
                0,  # someObjCtr
                0,  # vertexArrayPtr - заполним позже  
                0   # shiftInfoPtr - заполним позже
            ))
            
            # 3. Таблица указателей на объекты (начинается с 0x1C)
            # Каждый objectPtr - 4 байта
            object_ptr_table_pos = f.tell()  # Должно быть 0x1C
            
            # Резервируем места для указателей
            object_ptr_slots = []
            for i in range(scene_data['object_count']):
                slot_pos = f.tell()
                object_ptr_slots.append(slot_pos)
                f.write(b'\x00\x00\x00\x00')  # Заполним позже
            
            # 4. Запись gpuPacket объектов
            object_data_addresses = []
            
            for obj_idx, obj_data in enumerate(scene_data['objects']):
                # Запоминаем адрес начала данных объекта
                object_data_pos = f.tell()
                object_data_addresses.append(object_data_pos)
                
                # Записываем gpuPacket заголовок
                global_chunk_cnt = obj_data['polygon_count']
                f.write(struct.pack('<HHI', 
                    global_chunk_cnt,  # globalChunkCnt
                    0,                 # localChunkCnt (обычно 1)
                    obj_data['id']     # objectId
                ))
                
                # Записываем чанки полигонов
                f.write(obj_data['chunk_data'])
            
            # 5. Обновляем указатели в таблице
            current_pos = f.tell()
            for i, slot_pos in enumerate(object_ptr_slots):
                f.seek(slot_pos)
                f.write(struct.pack('<I', object_data_addresses[i]))
            f.seek(current_pos)


            # 6. Запись массива вершин
            vertex_ptr = f.tell()
            for x, y, z, d in scene_data['vertices']:
                f.write(struct.pack('<4h', x, y, z, d))
            
            # 7. Запись таблицы смещений
            shift_ptr = f.tell()
            self.write_shift_table(f, scene_data['object_count'])
            
            # 8. Обновляем заголовок сцены
            f.seek(header_pos)
            f.write(struct.pack('<HHII',
                scene_data['object_count'],  # objectsCnt
                0,                           # someObjCtr
                vertex_ptr,                  # vertexArrayPtr  
                shift_ptr                    # shiftInfoPtr
            ))

    def write_shift_table(self, f, object_count):
        """Запись таблицы смещений (objShiftInfo)"""
        for i in range(object_count):
            
            f.write(struct.pack('<B', 0))      # 0
            f.write(struct.pack('<B', 16))      # 0
            f.write(struct.pack('<B', 0))     # 16
            f.write(struct.pack('<B', 0))      # 0
            
            f.write(struct.pack('<i', 0))
           
            f.write(struct.pack('<B', 0))      # 0
            f.write(struct.pack('<B', 16))     # 16
            f.write(struct.pack('<B', 0))      # 0
            f.write(struct.pack('<B', 0))      # 0
           
            f.write(struct.pack('<i', 0))
           
            f.write(struct.pack('<B', 0))      # 0
            f.write(struct.pack('<B', 16))     # 16
            f.write(struct.pack('<B', 0))      # 0
            f.write(struct.pack('<B', 0))      # 0
            
            # shiftX, shiftY, shiftZ (3 x int32)
            # (0,0,0,0)  # комментарий о координатах инициализации
            f.write(struct.pack('<3i', 0, 0, 0))
            
            # # Последние 5 x u32 hz2
            f.write(struct.pack('<i', 0) * 5)

# Регистрация
# Регистрация плагина
def menu_func_export(self, context):
    self.layout.operator(ExportP2E.bl_idname, text="Persona2 PSX Scene (.p2e)")

def register():
    bpy.utils.register_class(ExportP2E)
    bpy.types.TOPBAR_MT_file_export.append(menu_func_export)

def unregister():
    bpy.utils.unregister_class(ExportP2E)
    bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)

if __name__ == "__main__":
    register()