"""
STEP to GLB Converter - Complete Working Version
All features included, proper mesh handling
"""

import sys
import json
from pathlib import Path
from datetime import datetime

print("=" * 80)
print("STEP to GLB Converter - Professional Version")
print("=" * 80)

# Import FreeCAD modules
try:
    print("Importing FreeCAD modules...")
    import FreeCAD
    import Import
    import Mesh
    import Part
    print(f"✓ FreeCAD Version: {FreeCAD.Version()}")
except ImportError as e:
    print(f"✗ ERROR: FreeCAD modules not found: {e}")
    sys.exit(1)

# ============================================================================
# CONFIGURATION
# ============================================================================
STEP_FILE = Path("D:\\__sistem__\\ProgramFiles\\FreeCAD\\bin\\freecad.stp")
OUTPUT_FILE = None  # None = auto-generate name

# Mesh Quality Settings
MESH_TOLERANCE = 0.1  # Lower = finer mesh (0.01-1.0)
MESH_ANGULAR_DEFLECTION = 0.5  # Degrees (0.1-1.0)
MIN_FACETS_THRESHOLD = 3  # Skip faces with fewer facets

# Export Options
EXPORT_FACES = True
EXPORT_EDGES = True
EXPORT_VERTICES = False

# Visual Settings
DEFAULT_COLOR = (0.7, 0.7, 0.7, 1.0)  # RGBA
EDGE_COLOR = (0.0, 0.0, 0.0, 1.0)  # Black edges
VERTEX_COLOR = (1.0, 0.0, 0.0, 1.0)  # Red vertices
VERTEX_SIZE = 0.5  # Vertex sphere radius

# Metadata
INCLUDE_METADATA = True
METADATA_FILE = True  # Save JSON separately

# ============================================================================
# UTILITIES
# ============================================================================

def log_section(title):
    """Print section header"""
    print("\n" + "=" * 80)
    print(f" {title}")
    print("=" * 80)


def get_color_from_body(body, face_idx):
    """Extract color from body's visual properties"""
    try:
        if hasattr(body, 'ViewObject'):
            vo = body.ViewObject
            
            # Try DiffuseColor (per-face colors)
            if hasattr(vo, 'DiffuseColor') and vo.DiffuseColor:
                colors = vo.DiffuseColor
                if len(colors) > face_idx:
                    rgb = colors[face_idx]
                    return (rgb[0], rgb[1], rgb[2])
                elif len(colors) > 0:
                    rgb = colors[0]
                    return (rgb[0], rgb[1], rgb[2])
            
            # Try ShapeColor (single color)
            if hasattr(vo, 'ShapeColor'):
                rgb = vo.ShapeColor
                return (rgb[0], rgb[1], rgb[2])
    except:
        pass
    
    return (DEFAULT_COLOR[0], DEFAULT_COLOR[1], DEFAULT_COLOR[2])


def get_transparency_from_body(body):
    """Extract transparency from body"""
    try:
        if hasattr(body, 'ViewObject'):
            trans = getattr(body.ViewObject, 'Transparency', 0)
            return 1.0 - (trans / 100.0)
    except:
        pass
    return 1.0


def apply_vertex_colors_to_mesh(mesh, color_rgba):
    """Apply uniform color to all mesh vertices"""
    try:
        num_points = mesh.CountPoints
        colors = [color_rgba] * num_points
        mesh.VertexColors = colors
        return True
    except AttributeError:
        return False
    except Exception as e:
        print(f"      Warning: Color application failed: {e}")
        return False


# ============================================================================
# MESH CREATION
# ============================================================================

def tessellate_face_to_mesh(face, body, body_idx, face_idx, body_label):
    """Convert a face to mesh with proper tessellation"""
    try:
        # Tessellate the face
        triangles = face.tessellate(MESH_TOLERANCE)
        
        # triangles returns (vertices, face_indices)
        if not triangles or len(triangles) < 2:
            return None
        
        vertices, face_indices = triangles
        
        # Check if we have enough data
        if len(face_indices) < MIN_FACETS_THRESHOLD:
            return None
        
        # Create mesh from tessellation
        mesh = Mesh.Mesh()
        
        for tri in face_indices:
            if len(tri) >= 3:
                v1 = FreeCAD.Vector(vertices[tri[0]])
                v2 = FreeCAD.Vector(vertices[tri[1]])
                v3 = FreeCAD.Vector(vertices[tri[2]])
                mesh.addFacet(v1, v2, v3)
        
        if mesh.CountFacets == 0:
            return None
        
        # Apply colors
        rgb = get_color_from_body(body, face_idx)
        alpha = get_transparency_from_body(body)
        color_rgba = (rgb[0], rgb[1], rgb[2], alpha)
        
        apply_vertex_colors_to_mesh(mesh, color_rgba)
        
        return mesh
        
    except Exception as e:
        print(f"      Error tessellating face {face_idx}: {e}")
        return None


def create_edge_mesh(edge, color_rgba):
    """Create mesh representation of an edge (as thin cylinder)"""
    try:
        # Discretize edge
        points = edge.discretize(Number=20)
        
        if len(points) < 2:
            return None
        
        mesh = Mesh.Mesh()
        
        # Create line segments as very thin triangles
        for i in range(len(points) - 1):
            p1 = points[i]
            p2 = points[i + 1]
            
            # Create a tiny offset for visibility
            offset = FreeCAD.Vector(0.01, 0.01, 0.01)
            p3 = p1 + offset
            
            mesh.addFacet(p1, p2, p3)
        
        if mesh.CountFacets > 0:
            apply_vertex_colors_to_mesh(mesh, color_rgba)
            return mesh
        
        return None
        
    except Exception as e:
        return None


def create_vertex_mesh(vertex, color_rgba):
    """Create small sphere at vertex location"""
    try:
        point = vertex.Point
        sphere = Part.makeSphere(VERTEX_SIZE, point)
        
        mesh = Mesh.Mesh()
        mesh.addMesh(Mesh.Mesh(sphere.tessellate(0.1)))
        
        if mesh.CountFacets > 0:
            apply_vertex_colors_to_mesh(mesh, color_rgba)
            return mesh
        
        return None
        
    except Exception as e:
        return None


# ============================================================================
# MAIN PROCESSING
# ============================================================================

def process_bodies(doc):
    """Process all bodies in document and generate meshes"""
    
    all_meshes = []
    metadata = {
        "version": "1.0",
        "timestamp": datetime.now().isoformat(),
        "source_file": STEP_FILE.name,
        "bodies": []
    }
    
    stats = {
        "bodies_processed": 0,
        "faces_meshed": 0,
        "edges_meshed": 0,
        "vertices_meshed": 0,
        "total_facets": 0
    }
    
    log_section("Processing Bodies")
    
    for body_idx, body in enumerate(doc.Objects):
        # Skip objects without geometry
        if not hasattr(body, 'Shape') or not body.Shape:
            continue
        
        shape = body.Shape
        body_label = getattr(body, 'Label', f'Body_{body_idx}')
        
        # Skip if no faces
        if len(shape.Faces) == 0:
            continue
        
        print(f"\n[Body {body_idx + 1}] {body_label}")
        print(f"  Faces: {len(shape.Faces)}, Edges: {len(shape.Edges)}, Vertices: {len(shape.Vertexes)}")
        
        body_metadata = {
            "index": body_idx + 1,
            "label": body_label,
            "face_count": len(shape.Faces),
            "edge_count": len(shape.Edges),
            "vertex_count": len(shape.Vertexes)
        }
        
        body_mesh_count = 0
        
        # Process Faces
        if EXPORT_FACES:
            print(f"  Processing {len(shape.Faces)} faces...")
            for face_idx, face in enumerate(shape.Faces):
                mesh = tessellate_face_to_mesh(face, body, body_idx, face_idx, body_label)
                
                if mesh:
                    all_meshes.append(mesh)
                    stats["faces_meshed"] += 1
                    stats["total_facets"] += mesh.CountFacets
                    body_mesh_count += 1
                    
                    if face_idx % 50 == 0 and face_idx > 0:
                        print(f"    → {face_idx}/{len(shape.Faces)} faces processed")
        
        # Process Edges
        if EXPORT_EDGES:
            print(f"  Processing {len(shape.Edges)} edges...")
            for edge_idx, edge in enumerate(shape.Edges):
                mesh = create_edge_mesh(edge, EDGE_COLOR)
                
                if mesh:
                    all_meshes.append(mesh)
                    stats["edges_meshed"] += 1
        
        # Process Vertices
        if EXPORT_VERTICES:
            print(f"  Processing {len(shape.Vertexes)} vertices...")
            for vertex_idx, vertex in enumerate(shape.Vertexes):
                mesh = create_vertex_mesh(vertex, VERTEX_COLOR)
                
                if mesh:
                    all_meshes.append(mesh)
                    stats["vertices_meshed"] += 1
        
        if body_mesh_count > 0:
            stats["bodies_processed"] += 1
            print(f"  ✓ Body complete: {body_mesh_count} meshes created")
        
        metadata["bodies"].append(body_metadata)
    
    metadata["stats"] = stats
    
    return all_meshes, metadata, stats


def export_meshes(meshes, output_path, metadata):
    """Export meshes to GLB file"""
    
    log_section("Exporting Meshes")
    
    print(f"Total meshes: {len(meshes)}")
    print(f"Output path: {output_path}")
    
    if len(meshes) == 0:
        print("\n✗ ERROR: No meshes to export!")
        return False
    
    try:
        # Method 1: Direct export
        print("\nAttempting direct export...")
        Import.export(meshes, str(output_path))
        
        if output_path.exists():
            size_kb = output_path.stat().st_size / 1024
            print(f"✓ Export successful: {size_kb:.2f} KB")
            return True
        
    except Exception as e:
        print(f"✗ Direct export failed: {e}")
    
    try:
        # Method 2: Merge all meshes first
        print("\nTrying merged export...")
        merged = Mesh.Mesh()
        
        for i, mesh in enumerate(meshes):
            merged.addMesh(mesh)
            if i % 100 == 0 and i > 0:
                print(f"  Merged {i}/{len(meshes)} meshes...")
        
        print(f"Final merged mesh: {merged.CountFacets} facets")
        
        Import.export([merged], str(output_path))
        
        if output_path.exists():
            size_kb = output_path.stat().st_size / 1024
            print(f"✓ Merged export successful: {size_kb:.2f} KB")
            return True
        
    except Exception as e:
        print(f"✗ Merged export failed: {e}")
    
    return False


def save_metadata(metadata, output_path):
    """Save metadata to JSON file"""
    if not METADATA_FILE:
        return
    
    try:
        json_path = output_path.with_suffix('.json')
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2)
        
        print(f"✓ Metadata saved: {json_path}")
    except Exception as e:
        print(f"⚠ Could not save metadata: {e}")


# ============================================================================
# MAIN
# ============================================================================

def main():
    """Main execution function"""
    
    log_section("Configuration")
    
    print(f"Input File: {STEP_FILE}")
    print(f"  Exists: {STEP_FILE.exists()}")
    
    if STEP_FILE.exists():
        size_mb = STEP_FILE.stat().st_size / (1024 * 1024)
        print(f"  Size: {size_mb:.2f} MB")
    
    print(f"\nMesh Settings:")
    print(f"  Tolerance: {MESH_TOLERANCE}")
    print(f"  Angular Deflection: {MESH_ANGULAR_DEFLECTION}°")
    print(f"\nExport Options:")
    print(f"  Faces: {EXPORT_FACES}")
    print(f"  Edges: {EXPORT_EDGES}")
    print(f"  Vertices: {EXPORT_VERTICES}")
    
    # Validate input
    if not STEP_FILE.exists():
        print(f"\n✗ ERROR: File not found!")
        return 1
    
    if STEP_FILE.suffix.lower() not in ['.stp', '.step']:
        print(f"\n✗ ERROR: Not a STEP file!")
        return 1
    
    # Determine output path
    output_path = OUTPUT_FILE if OUTPUT_FILE else STEP_FILE.with_suffix('.glb')
    
    doc = None
    
    try:
        # Open STEP file
        log_section("Loading STEP File")
        
        print("Opening document...")
        Import.open(str(STEP_FILE))
        doc = FreeCAD.ActiveDocument
        
        if not doc:
            print("✗ Failed to open document")
            return 1
        
        print(f"✓ Document: {doc.Name}")
        print(f"✓ Objects: {len(doc.Objects)}")
        
        # List objects
        print("\nDocument contents:")
        for i, obj in enumerate(doc.Objects):
            obj_type = obj.TypeId
            has_shape = hasattr(obj, 'Shape') and obj.Shape
            
            if has_shape:
                face_count = len(obj.Shape.Faces)
                print(f"  [{i+1:3d}] {obj.Label:30s} | {obj_type:30s} | Faces: {face_count}")
            else:
                print(f"  [{i+1:3d}] {obj.Label:30s} | {obj_type:30s} | No geometry")
        
        # Process all bodies
        meshes, metadata, stats = process_bodies(doc)
        
        # Display statistics
        log_section("Statistics")
        print(f"Bodies processed: {stats['bodies_processed']}")
        print(f"Faces meshed: {stats['faces_meshed']}")
        print(f"Edges meshed: {stats['edges_meshed']}")
        print(f"Vertices meshed: {stats['vertices_meshed']}")
        print(f"Total facets: {stats['total_facets']:,}")
        print(f"Total meshes: {len(meshes)}")
        
        if len(meshes) == 0:
            print("\n✗ ERROR: No geometry was converted!")
            print("\nTroubleshooting:")
            print("  1. Increase MESH_TOLERANCE (try 0.5 or 1.0)")
            print("  2. Check if STEP file contains actual 3D geometry")
            print("  3. Try opening in FreeCAD GUI to verify content")
            return 1
        
        # Export
        success = export_meshes(meshes, output_path, metadata)
        
        if not success:
            print("\n✗ ERROR: Export failed!")
            return 1
        
        # Save metadata
        if INCLUDE_METADATA:
            save_metadata(metadata, output_path)
        
        # Final summary
        log_section("SUCCESS")
        print(f"✓ Output file: {output_path}")
        print(f"✓ File size: {output_path.stat().st_size / 1024:.2f} KB")
        print(f"✓ Meshes exported: {len(meshes)}")
        print(f"✓ Total facets: {stats['total_facets']:,}")
        
        return 0
        
    except Exception as e:
        print(f"\n✗ FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        return 1
        
    finally:
        if doc:
            try:
                FreeCAD.closeDocument(doc.Name)
                print("\n✓ Document closed")
            except:
                pass


# ============================================================================
# ENTRY POINT
# ============================================================================

if __name__ == "__main__":
    try:
        exit_code = main()
        sys.exit(exit_code)
    except KeyboardInterrupt:
        print("\n\n✗ Interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n✗ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)