import struct
import math
import os

def transform_stl(input_path, output_path, matrix):
    # matrix is a flat 16 float array from Three.js (column-major format)
    # [m0, m1, m2, m3,  m4, m5, m6, m7,  m8, m9, m10, m11,  m12, m13, m14, m15]
    
    with open(input_path, 'rb') as f:
        header = f.read(80)
        num_faces = struct.unpack('<I', f.read(4))[0]
        
        faces = f.read()

    # Precompute for speed
    def apply_m(_x, _y, _z):
        w = _x * matrix[3] + _y * matrix[7] + _z * matrix[11] + matrix[15]
        nx = (_x * matrix[0] + _y * matrix[4] + _z * matrix[8] + matrix[12]) / w
        ny = (_x * matrix[1] + _y * matrix[5] + _z * matrix[9] + matrix[13]) / w
        nz = (_x * matrix[2] + _y * matrix[6] + _z * matrix[10] + matrix[14]) / w
        return nx, ny, nz

    def apply_rot(_nx, _ny, _nz):
        # Normals don't need translation, only upper 3x3
        # Assuming isotropic scaling for normals, otherwise needs inverse transpose
        x = _nx * matrix[0] + _ny * matrix[4] + _nz * matrix[8]
        y = _nx * matrix[1] + _ny * matrix[5] + _nz * matrix[9]
        z = _nx * matrix[2] + _ny * matrix[6] + _nz * matrix[10]
        length = math.sqrt(x*x + y*y + z*z)
        if length > 0.000001:
            return x/length, y/length, z/length
        return 0.0, 0.0, 0.0

    out = bytearray(80 + 4 + num_faces * 50)
    out[0:80] = header
    out[80:84] = struct.pack('<I', num_faces)
    
    offset = 84
    src_offset = 0
    for _ in range(num_faces):
        face_data = faces[src_offset:src_offset+50]
        if len(face_data) < 50:
            break
            
        nx, ny, nz, v1x, v1y, v1z, v2x, v2y, v2z, v3x, v3y, v3z, attr = struct.unpack('<12fH', face_data)
        
        # Transform normal
        nnx, nny, nnz = apply_rot(nx, ny, nz)
        # Transform vertices
        nv1x, nv1y, nv1z = apply_m(v1x, v1y, v1z)
        nv2x, nv2y, nv2z = apply_m(v2x, v2y, v2z)
        nv3x, nv3y, nv3z = apply_m(v3x, v3y, v3z)
        
        struct.pack_into('<12fH', out, offset, nnx, nny, nnz, nv1x, nv1y, nv1z, nv2x, nv2y, nv2z, nv3x, nv3y, nv3z, attr)
        offset += 50
        src_offset += 50
        
    with open(output_path, 'wb') as f:
        f.write(out)

    return num_faces

def merge_stls(input_files, output_path):
    total_faces = 0
    meshes_data = []

    for path, matrix in input_files:
        with open(path, 'rb') as f:
            f.read(80)
            faces = struct.unpack('<I', f.read(4))[0]
            data = f.read(faces * 50)
            
            # transform data
            # To speed things up, we could just do the transform in python like above
            # For this simple prototype, let's process each file into memory:
            
            def apply_m(_x, _y, _z):
                w = _x * matrix[3] + _y * matrix[7] + _z * matrix[11] + matrix[15]
                nx = (_x * matrix[0] + _y * matrix[4] + _z * matrix[8] + matrix[12]) / w
                ny = (_x * matrix[1] + _y * matrix[5] + _z * matrix[9] + matrix[13]) / w
                nz = (_x * matrix[2] + _y * matrix[6] + _z * matrix[10] + matrix[14]) / w
                return nx, ny, nz

            def apply_rot(_nx, _ny, _nz):
                x = _nx * matrix[0] + _ny * matrix[4] + _nz * matrix[8]
                y = _nx * matrix[1] + _ny * matrix[5] + _nz * matrix[9]
                z = _nx * matrix[2] + _ny * matrix[6] + _nz * matrix[10]
                length = math.sqrt(x*x + y*y + z*z)
                if length > 0.000001:
                    return x/length, y/length, z/length
                return 0.0, 0.0, 0.0
            
            new_data = bytearray(faces * 50)
            src_offset = 0
            dst_offset = 0
            for _ in range(faces):
                n_x, n_y, n_z, v1x, v1y, v1z, v2x, v2y, v2z, v3x, v3y, v3z, attr = struct.unpack_from('<12fH', data, src_offset)
                
                nnx, nny, nnz = apply_rot(n_x, n_y, n_z)
                nv1x, nv1y, nv1z = apply_m(v1x, v1y, v1z)
                nv2x, nv2y, nv2z = apply_m(v2x, v2y, v2z)
                nv3x, nv3y, nv3z = apply_m(v3x, v3y, v3z)
                
                struct.pack_into('<12fH', new_data, dst_offset, nnx, nny, nnz, nv1x, nv1y, nv1z, nv2x, nv2y, nv2z, nv3x, nv3y, nv3z, attr)
                src_offset += 50
                dst_offset += 50
                
            meshes_data.append(new_data)
            total_faces += faces
            
    # write merged
    with open(output_path, 'wb') as f:
        f.write(b'\0' * 80)
        f.write(struct.pack('<I', total_faces))
        for d in meshes_data:
            f.write(d)
            
