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

import bpy
import struct
from bpy_extras.io_utils import ImportHelper
from bpy.props import FloatProperty, EnumProperty, BoolProperty

class ImportP2E(bpy.types.Operator, ImportHelper):
    bl_idname = "import_scene.p2e"
    bl_label = "Import Persona2 PSX Scene"
    bl_description = "Import Persona 2 Eternal Punishment PSX 3D scenes"
    bl_options = {'REGISTER', 'UNDO', 'PRESET'}
    
    filename_ext = ".p2e"
    
    filter_glob: bpy.props.StringProperty(
        default="*.p2e",
        options={'HIDDEN'},
        maxlen=255,
    )
    
    coord_scale: FloatProperty(
        name="Coordinate Scale",
        description="PSX fixed-point scale multiplier",
        default=0.01,
        min=0.00001,
        max=1.0,
        step=0.01,
        precision=5
    )
    
    scale_preset: EnumProperty(
        name="PSX Scale Preset",
        description="Standard PSX fixed-point scales",
        items=[
            ('CUSTOM', "Custom", "Custom scale"),
            ('100', "1/100", "Life-like scale"),
            ('4096', "1/4096 (12.4)", "High precision"),
            ('256', "1/256 (8.8)", "Alternative scale"),
            ('16', "1/16 (12.4)", "Typical PSX vertex scale"), 
            ('1', "1/1 (Raw)", "No scaling"),
        ],
        default='100'
    )
    
    show_vertices: BoolProperty(
        name="Show Vertices",
        description="Show vertex points for debugging",
        default=False
    )
    
    vertex_size: FloatProperty(
        name="Vertex Size",
        description="Size of debug vertex markers",
        default=0.1,
        min=0.001,
        max=1.0
    )
    
    def draw(self, context):
        layout = self.layout
        
        box = layout.box()
        box.label(text="PSX Import Settings", icon='SETTINGS')
        
        box.prop(self, "scale_preset")
        
        if self.scale_preset != 'CUSTOM':
            preset_scales = {
                '100': 1/100.0,
                '4096': 1/4096.0,
                '256': 1/256.0,
                '16': 1/16.0,
                '1': 1.0
            }
            self.coord_scale = preset_scales[self.scale_preset]
        
        if self.scale_preset == 'CUSTOM':
            box.prop(self, "coord_scale")
        
        box.separator()
        box.label(text="Debug Options", icon='MODIFIER')
        box.prop(self, "show_vertices")
        if self.show_vertices:
            box.prop(self, "vertex_size")
        
        box.separator()
        box.label(text=f"Current scale: {self.coord_scale:.6f}", icon='INFO')
    
    def execute(self, context):
        # Читаем файл
        vertices, objects_data = self.read_p2e_file(self.filepath)
        
        if not vertices:
            self.report({'ERROR'}, "No vertices found")
            return {'CANCELLED'}
        
        # Создаем меши для каждого объекта
        self.create_objects_from_data(vertices, objects_data)
        
        # Опционально показываем вершины для отладки
        if self.show_vertices:
            self.create_debug_vertices(vertices)
        
        self.report({'INFO'}, f"Imported {len(vertices)} vertices, {len(objects_data)} objects")
        return {'FINISHED'}
    
    def read_p2e_file(self, filepath):
        """Чтение .p2e файла: вершины и полигоны объектов"""
        try:
            with open(filepath, 'rb') as f:
                # Пропускаем заголовок (16 байт sourceName)
                f.seek(0x10)
                
                # Читаем заголовок сцены
                data = f.read(12)
                if len(data) < 12:
                    self.report({'ERROR'}, "File too small")
                    return [], []
                
                objects_cnt, some_obj_ctr, vertex_ptr, shift_ptr = struct.unpack('<HHII', data)
                
                # Вычисляем количество вершин
                vertex_count = (shift_ptr - vertex_ptr) // 8
                
                print(f"=== P2E File Info ===")
                print(f"Objects: {objects_cnt}")
                print(f"Vertices: {vertex_count}")
                print(f"Vertex pointer: 0x{vertex_ptr:08X}")
                print(f"Shift pointer: 0x{shift_ptr:08X}")
                
                # 1. Читаем вершины
                vertices = []
                f.seek(vertex_ptr)
                
                for i in range(vertex_count):
                    vertex_data = f.read(8)
                    if len(vertex_data) < 8:
                        break
                    
                    x, y, z, d = struct.unpack('<4h', vertex_data)
                    # PSX координаты: X=X, Y=Z, Z=-Y
                    vertex = (
                        x * self.coord_scale,
                        z * self.coord_scale,
                        -y * self.coord_scale
                    )
                    vertices.append(vertex)
                
                print(f"Read {len(vertices)} vertices")
                
                # 2. Читаем объекты
                objects_data = []
                f.seek(0x10 + 12)  # После заголовка сцены начинаются объекты
                
                for obj_idx in range(objects_cnt):
                    # Читаем указатель на объект
                    obj_ptr_data = f.read(4)
                    if len(obj_ptr_data) < 4:
                        break
                    
                    obj_ptr = struct.unpack('<I', obj_ptr_data)[0]
                    
                    # Сохраняем текущую позицию
                    current_pos = f.tell()
                    
                    # Переходим к данным объекта
                    f.seek(obj_ptr)
                    
                    # Читаем gpuPacket заголовок
                    packet_data = f.read(8)  # u16 globalChunkCnt, u16 localChunkCnt, u32 objectId
                    if len(packet_data) < 8:
                        f.seek(current_pos)
                        continue
                    
                    global_chunk_cnt, local_chunk_cnt, object_id = struct.unpack('<HHI', packet_data)
                    
                    print(f"\nObject {obj_idx}:")
                    print(f"  Pointer: 0x{obj_ptr:08X}")
                    print(f"  ID: {object_id}")
                    print(f"  Chunks: {global_chunk_cnt}")
                    
                    # Собираем полигоны этого объекта
                    obj_polygons = []
                    
                    for chunk_idx in range(global_chunk_cnt):
                        # Читаем заголовок chunk
                        chunk_header = f.read(4)
                        if len(chunk_header) < 4:
                            break
                        
                        cmd_cnt, chunk_size_words, header_fuck = struct.unpack('<BBH', chunk_header)
                        
                        # print(f"  Chunk {chunk_idx}: cmd={cmd_cnt}, size={chunk_size_words} words")
                        
                        # Определяем тип полигона по cmd_cnt
                        poly_type = self.get_polygon_type(cmd_cnt)
                        
                        if poly_type is None:
                            # Пропускаем неизвестный chunk
                            bytes_to_skip = chunk_size_words * 4  # chunk_size в словах (4 байта)
                            f.read(bytes_to_skip)
                            continue
                        
                        # Читаем ВЕСЬ chunk (включая индексы вершин)
                        # chunk_size_words - количество 4-байтных слов данных ПОСЛЕ заголовка
                        chunk_data_size = chunk_size_words * 4
                        poly_data = f.read(chunk_data_size)
                        if len(poly_data) < chunk_data_size:
                            break
                        
                        # Извлекаем индексы вершин из конца chunk
                        vertex_indices = self.extract_vertex_indices(poly_type, poly_data)
                        
                        if vertex_indices:
                            obj_polygons.append({
                                'type': poly_type,
                                'cmd': cmd_cnt,
                                'vertices': vertex_indices,
                                'chunk_size': chunk_size_words
                            })
                    
                    # Сохраняем данные объекта
                    if obj_polygons:
                        objects_data.append({
                            'id': object_id,
                            'ptr': obj_ptr,
                            'polygons': obj_polygons,
                            'vertex_count': len(obj_polygons) * 4  # Примерно
                        })
                    
                    # Возвращаемся к списку объектов
                    f.seek(current_pos)
                
                print(f"\n=== Summary ===")
                print(f"Total objects with polygons: {len(objects_data)}")
                
                return vertices, objects_data
                
        except Exception as e:
            self.report({'ERROR'}, f"Error reading file: {str(e)}")
            import traceback
            traceback.print_exc()
            return [], []
    
    def get_polygon_type(self, cmd_cnt):
        """Определяем тип полигона по команде"""
        poly_types = {
            7: 'shaded_textured_quad',      # shadedTexturedPoly
            6: 'textured_quad',             # texturedPoly
            5: 'shaded_quad',               # shadedPoly
            3: 'shaded_textured_tri',       # shadedTextured3Poly
            2: 'textured_tri',              # textured3Poly
        }
        
        return poly_types.get(cmd_cnt, None)
    
    def extract_vertex_indices(self, poly_type, poly_data):
        """Извлекаем индексы вершин из конца chunk данных"""
        try:
            if poly_type in ['shaded_textured_quad', 'textured_quad', 'shaded_quad']:
                # Квады: 4 индекса uint32 в конце (16 байт)
                if len(poly_data) >= 16:
                    # Последние 16 байт: vertexId1-4
                    v1, v2, v3, v4 = struct.unpack_from('<4I', poly_data, len(poly_data) - 16)
                    return [v1, v2, v4, v3]
            
            elif poly_type in ['shaded_textured_tri', 'textured_tri']:
                # Треугольники: 3 индекса uint32 в конце (12 байт)
                if len(poly_data) >= 12:
                    # Последние 12 байт: vertexId1-3
                    v1, v2, v3 = struct.unpack_from('<3I', poly_data, len(poly_data) - 12)
                    return [v1, v2, v3]
            
        except Exception as e:
            print(f"Error extracting indices from {poly_type}: {e}")
        
        return []
    
    def create_objects_from_data(self, vertices, objects_data):
        """Создание отдельных мешей для каждого объекта"""
        if not objects_data:
            return
        
        # Создаем коллекцию для импортированных объектов
        collection_name = "PSX_Objects"
        if collection_name not in bpy.data.collections:
            collection = bpy.data.collections.new(collection_name)
            bpy.context.scene.collection.children.link(collection)
        else:
            collection = bpy.data.collections[collection_name]
        
        total_polygons = 0
        successful_objects = 0
        
        for obj_idx, obj_data in enumerate(objects_data):
            # Собираем все полигоны этого объекта
            all_faces = []
            
            for poly in obj_data['polygons']:
                vertex_ids = poly['vertices']
                
                # Проверяем, что индексы в пределах массива вершин
                valid_ids = []
                for vid in vertex_ids:
                    if vid < len(vertices):
                        valid_ids.append(vid)
                    else:
                        print(f"Warning: Vertex index {vid} out of range (max {len(vertices)-1})")
                
                if len(valid_ids) >= 3:  # Нужно минимум 3 вершины для полигона
                    all_faces.append(valid_ids)
                    total_polygons += 1
            
            if not all_faces:
                print(f"Object {obj_idx} has no valid polygons")
                continue
            
            # Создаем меш
            mesh_name = f"PSX_Obj_{obj_idx:03d}_ID{obj_data['id']}"
            mesh = bpy.data.meshes.new(mesh_name)
            
            # Создаем вершины и полигоны
            mesh.from_pydata(vertices, [], all_faces)
            
            # Обновляем меш
            mesh.update()
            
            # Создаем объект
            obj = bpy.data.objects.new(mesh_name, mesh)
            
            # Добавляем кастомные свойства
            obj["psx_object_id"] = obj_data['id']
            obj["psx_object_ptr"] = obj_data['ptr']
            obj["psx_poly_count"] = len(all_faces)
            obj["psx_poly_types"] = ','.join(str(p['cmd']) for p in obj_data['polygons'])
            
            # Добавляем в коллекцию
            collection.objects.link(obj)
            successful_objects += 1
            
            print(f"  Created object {obj_idx} with {len(all_faces)} polygons")
        
        print(f"\nCreated {successful_objects} objects with {total_polygons} total polygons")
    
    def create_debug_vertices(self, vertices):
        """Создание мета-шаров для отладки вершин"""
        if not vertices:
            return
        
        meta = bpy.data.metaballs.new("Debug_Vertices")
        meta_obj = bpy.data.objects.new("PSX_Debug_Vertices", meta)
        
        meta.resolution = 0.1
        meta.threshold = 0.05
        
        for i, vertex in enumerate(vertices):
            elem = meta.elements.new()
            elem.co = vertex
            elem.radius = self.vertex_size
        
        bpy.context.collection.objects.link(meta_obj)
        
        print(f"Created {len(vertices)} debug vertices")

def menu_func_import(self, context):
    self.layout.operator(ImportP2E.bl_idname, text="Persona2 PSX Scene (.p2e)")

classes = (ImportP2E,)

def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    
    bpy.types.TOPBAR_MT_file_import.append(menu_func_import)

def unregister():
    for cls in classes:
        bpy.utils.unregister_class(cls)
    
    bpy.types.TOPBAR_MT_file_import.remove(menu_func_import)

if __name__ == "__main__":
    register()