#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
3D Tiles Tileset to PLY Converter
将整个3D Tiles tileset转换为单个PLY文件
"""

import json
import struct
import os
import sys
from pathlib import Path
import numpy as np

try:
    import DracoPy
except ImportError:
    print("警告: DracoPy库未安装，无法处理Draco压缩的数据")
    print("请运行: pip install DracoPy")
    DracoPy = None

class TilesetToPLYConverter:
    def __init__(self):
        self.all_vertices = []
        self.vertex_count = 0
        
    def multiply_matrix_vector(self, matrix, vector):
        """4x4矩阵与4D向量相乘"""
        # matrix是16个元素的列表，按列主序排列
        # 转换为4x4矩阵（行主序）
        m = [
            [matrix[0], matrix[4], matrix[8], matrix[12]],
            [matrix[1], matrix[5], matrix[9], matrix[13]],
            [matrix[2], matrix[6], matrix[10], matrix[14]],
            [matrix[3], matrix[7], matrix[11], matrix[15]]
        ]
        
        # 向量扩展为齐次坐标 [x, y, z, 1]
        v = [vector[0], vector[1], vector[2], 1.0]
        
        # 矩阵乘法
        result = [
            m[0][0]*v[0] + m[0][1]*v[1] + m[0][2]*v[2] + m[0][3]*v[3],
            m[1][0]*v[0] + m[1][1]*v[1] + m[1][2]*v[2] + m[1][3]*v[3],
            m[2][0]*v[0] + m[2][1]*v[1] + m[2][2]*v[2] + m[2][3]*v[3]
        ]
        
        return result
    
    def multiply_matrices(self, m1, m2):
        """两个4x4矩阵相乘"""
        # 将16元素列表转换为4x4矩阵
        def list_to_matrix(lst):
            return [
                [lst[0], lst[4], lst[8], lst[12]],
                [lst[1], lst[5], lst[9], lst[13]],
                [lst[2], lst[6], lst[10], lst[14]],
                [lst[3], lst[7], lst[11], lst[15]]
            ]
        
        def matrix_to_list(mat):
            return [
                mat[0][0], mat[1][0], mat[2][0], mat[3][0],
                mat[0][1], mat[1][1], mat[2][1], mat[3][1],
                mat[0][2], mat[1][2], mat[2][2], mat[3][2],
                mat[0][3], mat[1][3], mat[2][3], mat[3][3]
            ]
        
        a = list_to_matrix(m1)
        b = list_to_matrix(m2)
        
        result = [[0 for _ in range(4)] for _ in range(4)]
        for i in range(4):
            for j in range(4):
                for k in range(4):
                    result[i][j] += a[i][k] * b[k][j]
        
        return matrix_to_list(result)
    
    def apply_transform_to_vertices(self, vertices, transform_matrix):
        """对顶点应用变换矩阵"""
        if not transform_matrix:
            return vertices
        
        transformed_vertices = []
        for vertex in vertices:
            transformed = self.multiply_matrix_vector(transform_matrix, vertex)
            transformed_vertices.append(transformed)
        
        return transformed_vertices
        
    def parse_tileset_json(self, tileset_path, parent_transform=None):
        """解析tileset.json文件，收集B3DM文件和变换矩阵"""
        try:
            with open(tileset_path, 'r', encoding='utf-8') as f:
                tileset_data = json.load(f)
            
            b3dm_files = []
            
            def process_node(node, base_path, accumulated_transform):
                # 获取当前节点的变换矩阵
                current_transform = node.get('transform')
                
                # 计算累积变换矩阵
                if current_transform and accumulated_transform:
                    # 矩阵相乘：accumulated_transform * current_transform
                    final_transform = self.multiply_matrices(accumulated_transform, current_transform)
                elif current_transform:
                    final_transform = current_transform
                else:
                    final_transform = accumulated_transform
                
                if 'content' in node and 'uri' in node['content']:
                    uri = node['content']['uri']
                    if uri.endswith('.b3dm'):
                        full_path = os.path.join(base_path, uri)
                        if os.path.exists(full_path):
                            b3dm_files.append((full_path, final_transform))
                    elif uri.endswith('.json'):
                        # 递归处理子tileset
                        sub_tileset_path = os.path.join(base_path, uri)
                        if os.path.exists(sub_tileset_path):
                            sub_files = self.parse_tileset_json(sub_tileset_path, final_transform)
                            b3dm_files.extend(sub_files)
                
                if 'children' in node:
                    for child in node['children']:
                        process_node(child, base_path, final_transform)
            
            base_path = os.path.dirname(tileset_path)
            if 'root' in tileset_data:
                process_node(tileset_data['root'], base_path, parent_transform)
            
            return b3dm_files
            
        except Exception as e:
            print(f"解析tileset.json时出错: {e}")
            return []
    
    def parse_b3dm_file(self, file_path):
        """解析B3DM文件"""
        try:
            with open(file_path, 'rb') as f:
                # 读取B3DM头部
                magic = f.read(4)
                if magic != b'b3dm':
                    print(f"警告: {file_path} 不是有效的B3DM文件")
                    return None
                
                version = struct.unpack('<I', f.read(4))[0]
                byte_length = struct.unpack('<I', f.read(4))[0]
                feature_table_json_byte_length = struct.unpack('<I', f.read(4))[0]
                feature_table_binary_byte_length = struct.unpack('<I', f.read(4))[0]
                batch_table_json_byte_length = struct.unpack('<I', f.read(4))[0]
                batch_table_binary_byte_length = struct.unpack('<I', f.read(4))[0]
                
                # 跳过feature table和batch table
                f.seek(28 + feature_table_json_byte_length + feature_table_binary_byte_length + 
                       batch_table_json_byte_length + batch_table_binary_byte_length)
                
                # 读取glTF数据
                gltf_data = f.read()
                return self.parse_gltf_data(gltf_data)
                
        except Exception as e:
            print(f"解析B3DM文件 {file_path} 失败: {e}")
            return None
    
    def parse_gltf_data(self, gltf_data):
        """解析glTF数据"""
        try:
            # 检查是否为GLB格式
            if gltf_data[:4] == b'glTF':
                return self.parse_glb_data(gltf_data)
            else:
                # 尝试作为JSON解析
                gltf_json = json.loads(gltf_data.decode('utf-8'))
                return self.extract_vertices_from_gltf(gltf_json, None)
        except Exception as e:
            print(f"解析glTF数据失败: {e}")
            return None
    
    def parse_glb_data(self, glb_data):
        """解析GLB格式的glTF数据"""
        try:
            # GLB头部: magic(4) + version(4) + length(4)
            magic = glb_data[:4]
            if magic != b'glTF':
                return None
            
            version = struct.unpack('<I', glb_data[4:8])[0]
            total_length = struct.unpack('<I', glb_data[8:12])[0]
            
            offset = 12
            json_data = None
            binary_data = None
            
            # 读取chunks
            while offset < len(glb_data):
                if offset + 8 > len(glb_data):
                    break
                    
                chunk_length = struct.unpack('<I', glb_data[offset:offset+4])[0]
                chunk_type = glb_data[offset+4:offset+8]
                chunk_data = glb_data[offset+8:offset+8+chunk_length]
                
                if chunk_type == b'JSON':
                    json_data = json.loads(chunk_data.decode('utf-8'))
                elif chunk_type == b'BIN\x00':
                    binary_data = chunk_data
                
                offset += 8 + chunk_length
                # 对齐到4字节边界
                offset = (offset + 3) & ~3
            
            if json_data:
                return self.extract_vertices_from_gltf(json_data, binary_data)
            
        except Exception as e:
            print(f"解析GLB数据失败: {e}")
            return None
    
    def extract_vertices_from_gltf(self, gltf_json, binary_data):
        """从glTF JSON中提取顶点数据"""
        vertices = []
        
        try:
            # 检查是否使用了Draco压缩
            if 'extensionsUsed' in gltf_json and 'KHR_draco_mesh_compression' in gltf_json['extensionsUsed']:
                if DracoPy is None:
                    print("警告: 检测到Draco压缩但DracoPy未安装")
                    return vertices
                return self.extract_draco_vertices(gltf_json, binary_data)
            
            # 处理标准glTF格式
            if 'meshes' not in gltf_json:
                return vertices
            
            for mesh in gltf_json['meshes']:
                for primitive in mesh['primitives']:
                    if 'attributes' in primitive and 'POSITION' in primitive['attributes']:
                        position_accessor_index = primitive['attributes']['POSITION']
                        
                        if 'accessors' in gltf_json and position_accessor_index < len(gltf_json['accessors']):
                            accessor = gltf_json['accessors'][position_accessor_index]
                            buffer_view_index = accessor['bufferView']
                            
                            if 'bufferViews' in gltf_json and buffer_view_index < len(gltf_json['bufferViews']):
                                buffer_view = gltf_json['bufferViews'][buffer_view_index]
                                buffer_index = buffer_view['buffer']
                                byte_offset = buffer_view.get('byteOffset', 0) + accessor.get('byteOffset', 0)
                                
                                if binary_data and buffer_index == 0:
                                    # 从二进制数据中读取顶点
                                    component_type = accessor['componentType']
                                    count = accessor['count']
                                    
                                    if component_type == 5126:  # FLOAT
                                        vertex_data = struct.unpack(f'<{count*3}f', 
                                                                   binary_data[byte_offset:byte_offset+count*12])
                                        for i in range(0, len(vertex_data), 3):
                                            vertices.append([vertex_data[i], vertex_data[i+1], vertex_data[i+2]])
            
        except Exception as e:
            print(f"提取顶点数据失败: {e}")
        
        return vertices
    
    def extract_draco_vertices(self, gltf_json, binary_data):
        """提取Draco压缩的顶点数据"""
        vertices = []
        
        try:
            if 'meshes' not in gltf_json:
                return vertices
            
            for mesh in gltf_json['meshes']:
                for primitive in mesh['primitives']:
                    if 'extensions' in primitive and 'KHR_draco_mesh_compression' in primitive['extensions']:
                        draco_ext = primitive['extensions']['KHR_draco_mesh_compression']
                        buffer_view_index = draco_ext['bufferView']
                        
                        if 'bufferViews' in gltf_json and buffer_view_index < len(gltf_json['bufferViews']):
                            buffer_view = gltf_json['bufferViews'][buffer_view_index]
                            byte_offset = buffer_view.get('byteOffset', 0)
                            byte_length = buffer_view['byteLength']
                            
                            if binary_data:
                                draco_data = binary_data[byte_offset:byte_offset+byte_length]
                                
                                # 使用DracoPy解码
                                mesh_data = DracoPy.decode(draco_data)
                                if hasattr(mesh_data, 'points'):
                                    points = mesh_data.points
                                    for point in points:
                                        vertices.append([float(point[0]), float(point[1]), float(point[2])])
        
        except Exception as e:
            print(f"解码Draco数据失败: {e}")
        
        return vertices
    
    def save_ply_file(self, output_path):
        """保存PLY文件"""
        try:
            with open(output_path, 'w') as f:
                # 写入PLY头部
                f.write("ply\n")
                f.write("format ascii 1.0\n")
                f.write(f"element vertex {len(self.all_vertices)}\n")
                f.write("property float x\n")
                f.write("property float y\n")
                f.write("property float z\n")
                f.write("end_header\n")
                
                # 写入顶点数据
                for vertex in self.all_vertices:
                    f.write(f"{vertex[0]} {vertex[1]} {vertex[2]}\n")
            
            print(f"PLY文件已保存: {output_path}")
            print(f"总顶点数: {len(self.all_vertices)}")
            
        except Exception as e:
            print(f"保存PLY文件失败: {e}")
    
    def convert_tileset_to_ply(self, tileset_path, output_path):
        """将整个tileset转换为PLY文件"""
        print(f"开始处理tileset: {tileset_path}")
        
        # 解析主tileset.json
        tileset_data = self.parse_tileset_json(tileset_path)
        if not tileset_data:
            print("无法解析tileset.json文件")
            return False
        
        # 获取基础路径
        base_path = os.path.dirname(tileset_path)
        
        # 提取所有b3dm文件和变换矩阵
        b3dm_files = self.parse_tileset_json(tileset_path)
        print(f"找到 {len(b3dm_files)} 个B3DM文件")
        
        if not b3dm_files:
            print("未找到任何B3DM文件")
            return False
        
        # 处理每个b3dm文件
        processed_count = 0
        for i, (b3dm_file, transform_matrix) in enumerate(b3dm_files):
            print(f"处理文件 {i+1}/{len(b3dm_files)}: {os.path.basename(b3dm_file)}")
            
            vertices = self.parse_b3dm_file(b3dm_file)
            if vertices:
                # 应用变换矩阵
                if transform_matrix:
                    vertices = self.apply_transform_to_vertices(vertices, transform_matrix)
                    print(f"  应用了变换矩阵")
                
                self.all_vertices.extend(vertices)
                processed_count += 1
                print(f"  提取到 {len(vertices)} 个顶点")
        
        print(f"\n成功处理 {processed_count}/{len(b3dm_files)} 个文件")
        print(f"总计提取 {len(self.all_vertices)} 个顶点")
        
        if self.all_vertices:
            self.save_ply_file(output_path)
            return True
        else:
            print("未提取到任何顶点数据")
            return False

def main():
    if len(sys.argv) < 2:
        print("用法: python tileset_to_ply.py <tileset.json路径> [输出PLY文件路径]")
        print("示例: python tileset_to_ply.py tileset.json output.ply")
        return
    
    tileset_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else "merged_tileset.ply"
    
    if not os.path.exists(tileset_path):
        print(f"错误: 文件不存在 {tileset_path}")
        return
    
    converter = TilesetToPLYConverter()
    success = converter.convert_tileset_to_ply(tileset_path, output_path)
    
    if success:
        print("\n转换完成!")
    else:
        print("\n转换失败!")

if __name__ == "__main__":
    main()