import trimesh
import numpy as np
import ezdxf
from scipy.spatial import ConvexHull
from scipy.interpolate import splprep, splev

def create_layered_curves_dxf(input_file, output_file, num_layers=20, num_points_per_layer=30):
    """
    Creates a DXF file containing simplified, layered profiles from a 3D model,
    yielding progress updates along the way.
    """
    try:
        # 1. Load and Auto-Orient the Model
        yield {"status": "processing", "progress": 5, "message": f"Loading model..."}
        mesh = trimesh.load(input_file)
        
        # After loading, the result might be a Scene, so we extract the geometry.
        # This is a robust way to handle various mesh file formats.
        if isinstance(mesh, trimesh.Scene):
            # If the scene has no geometry, it's an invalid mesh
            if not mesh.geometry:
                 raise ValueError("Input file loaded as an empty scene.")
            # If there are multiple geometries, combine them into one
            if len(mesh.geometry) > 1:
                mesh = trimesh.util.concatenate(list(mesh.geometry.values()))
            else:
                mesh = list(mesh.geometry.values())[0]

        # Ensure we have a valid mesh with vertices
        if not isinstance(mesh, trimesh.Trimesh) or not mesh.vertices.shape[0] > 0:
            raise ValueError("Input file could not be loaded as a valid mesh with vertices.")

        yield {"status": "processing", "progress": 15, "message": f"Analyzing orientation..."}
        extents = mesh.extents
        longest_axis_index = np.argmax(extents)
        if longest_axis_index != 2:
            yield {"status": "processing", "progress": 25, "message": f"Re-orienting model..."}
            transform = trimesh.transformations.rotation_matrix(
                angle=trimesh.transformations.angle_between_vectors(np.eye(3)[longest_axis_index], [0, 0, 1]),
                direction=np.cross(np.eye(3)[longest_axis_index], [0, 0, 1])
            )
            mesh.apply_transform(transform)
            mesh.vertices -= mesh.center_mass

        # 2. Generate Aligned "Fish Net" Profiles
        yield {"status": "processing", "progress": 40, "message": f"Generating {num_layers} profiles..."}
        bounds = mesh.bounds
        z_min, z_max = bounds[:, 2]
        z_height = z_max - z_min
        z_levels = np.linspace(z_min + z_height * 0.01, z_max - z_height * 0.01, num_layers)
        
        all_profiles_3d = []
        for i, z in enumerate(z_levels):
            # Report progress for each slice
            progress = 40 + int((i / num_layers) * 40)
            yield {"status": "processing", "progress": progress, "message": f"Slicing layer {i+1} of {num_layers}"}
            
            try:
                slice_3d = mesh.section(plane_origin=[0, 0, z], plane_normal=[0, 0, 1])
                if slice_3d is None: continue

                slice_vertices_2d = np.vstack([path for path in slice_3d.to_planar()[0].discrete])
                if len(slice_vertices_2d) < 3: continue
                
                hull = ConvexHull(slice_vertices_2d)
                hull_vertices = slice_vertices_2d[hull.vertices]

                tck, u = splprep([hull_vertices[:, 0], hull_vertices[:, 1]], s=0, per=True)
                u_new = np.linspace(u.min(), u.max(), num_points_per_layer)
                x_new, y_new = splev(u_new, tck)
                standardized_points = np.array(list(zip(x_new, y_new)))
                
                seam_start_index = np.argmax(standardized_points[:, 0])
                aligned_points_2d = np.roll(standardized_points, -seam_start_index, axis=0)
                
                profile_3d = np.column_stack((aligned_points_2d, np.full(len(aligned_points_2d), z)))
                all_profiles_3d.append(profile_3d)
            except Exception:
                # Silently skip layers that fail, the UI will just move to the next percentage
                continue

        if not all_profiles_3d:
            raise ValueError("Could not generate any valid profiles. The model may be too complex or thin.")

        # 3. Create a new DXF document and add polylines
        yield {"status": "processing", "progress": 85, "message": f"Creating DXF structure..."}
        doc = ezdxf.new()
        msp = doc.modelspace()
        
        for profile in all_profiles_3d:
            msp.add_polyline3d(points=profile, close=True)

        # 4. Save the DXF file
        yield {"status": "processing", "progress": 95, "message": f"Saving to file..."}
        doc.saveas(output_file)

    finally:
        # In a real app, you might clean up the input file here
        pass


create_layered_curves_dxf('input_model.stl', 'output_profiles.dxf')