"""
STEP to Three.js Optimized Format
Compact but full-featured for WebGL rendering
"""

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

print("=" * 80)
print("STEP to Three.js Format Converter")
print("=" * 80)

try:
    import FreeCAD
    import Import
    print(f"✓ FreeCAD Version: {FreeCAD.Version()[0]}.{FreeCAD.Version()[1]}")
except ImportError as e:
    print(f"✗ ERROR: FreeCAD not found: {e}")
    sys.exit(1)

# ============================================================================
# CONFIGURATION
# ============================================================================
STEP_FILE = Path(r"D:\__sistem__\ProgramFiles\FreeCAD\bin\freecad.stp")
OUTPUT_FILE = None
MESH_TOLERANCE = 1.0

# Export options
EXPORT_FACES = True
EXPORT_EDGES = True
EXPORT_VERTICES = False  # Usually not needed for rendering

USE_GZIP = True  # Recommended for web delivery

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

def get_face_color(body, face_idx):
    """Get RGB color for face from ViewObject"""
    try:
        if hasattr(body, 'ViewObject'):
            vo = body.ViewObject
            
            # Try per-face colors first (DiffuseColor)
            if hasattr(vo, 'DiffuseColor') and vo.DiffuseColor:
                colors = vo.DiffuseColor
                if len(colors) > face_idx:
                    return colors[face_idx][:3]
                elif len(colors) > 0:
                    return colors[0][:3]
            
            # Fallback to ShapeColor
            if hasattr(vo, 'ShapeColor'):
                return vo.ShapeColor[:3]
    except:
        pass
    return (0.7, 0.7, 0.7)

def get_transparency(body):
    """Get transparency value (0=transparent, 1=opaque)"""
    try:
        if hasattr(body, 'ViewObject'):
            trans = getattr(body.ViewObject, 'Transparency', 0)
            return 1.0 - (trans / 100.0)
    except:
        pass
    return 1.0

def get_line_color(body):
    """Get line/edge color from ViewObject"""
    try:
        if hasattr(body, 'ViewObject'):
            vo = body.ViewObject
            
            # LineColor property
            if hasattr(vo, 'LineColor'):
                color = vo.LineColor
                return (color[0], color[1], color[2])
            
            # ShapeColor as fallback
            if hasattr(vo, 'ShapeColor'):
                color = vo.ShapeColor
                return (color[0], color[1], color[2])
    except:
        pass
    return (0.0, 0.0, 0.0)  # Black default

def get_line_width(body):
    """Get line width from ViewObject"""
    try:
        if hasattr(body, 'ViewObject'):
            return getattr(body.ViewObject, 'LineWidth', 1.0)
    except:
        pass
    return 1.0

def get_point_size(body):
    """Get point size from ViewObject"""
    try:
        if hasattr(body, 'ViewObject'):
            return getattr(body.ViewObject, 'PointSize', 2.0)
    except:
        pass
    return 2.0

def get_point_color(body):
    """Get point/vertex color from ViewObject"""
    try:
        if hasattr(body, 'ViewObject'):
            vo = body.ViewObject
            
            # PointColor property
            if hasattr(vo, 'PointColor'):
                color = vo.PointColor
                return (color[0], color[1], color[2])
            
            # ShapeColor as fallback
            if hasattr(vo, 'ShapeColor'):
                color = vo.ShapeColor
                return (color[0], color[1], color[2])
    except:
        pass
    return (1.0, 0.0, 0.0)  # Red default

def get_display_mode(body):
    """Get display mode (Flat, Shaded, Wireframe, Points)"""
    try:
        if hasattr(body, 'ViewObject'):
            return getattr(body.ViewObject, 'DisplayMode', 'Flat Lines')
    except:
        pass
    return 'Flat Lines'

def get_lighting(body):
    """Get lighting mode (One side, Two side)"""
    try:
        if hasattr(body, 'ViewObject'):
            return getattr(body.ViewObject, 'Lighting', 'Two side')
    except:
        pass
    return 'Two side'

def round_vec3(v, precision=4):
    """Round 3D vector"""
    return [round(v[i], precision) for i in range(3)]

def round_color(c, precision=3):
    """Round color values"""
    return [round(c[i], precision) for i in range(len(c))]

# ============================================================================
# GEOMETRY PROCESSING
# ============================================================================

def process_faces(shape, body, body_idx):
    """Extract face geometry with per-vertex colors"""
    vertices = []
    indices = []
    vertex_colors = []  # Color for EACH vertex
    face_metadata = []  # Metadata per face
    
    for face_idx, face in enumerate(shape.Faces):
        try:
            tess = face.tessellate(MESH_TOLERANCE)
            if not tess or len(tess) < 2:
                continue
            
            verts, faces = tess
            if len(faces) == 0:
                continue
            
            # Get face color from STEP file
            rgb = get_face_color(body, face_idx)
            alpha = get_transparency(body)
            color = round_color([rgb[0], rgb[1], rgb[2], alpha])
            
            # Face metadata
            face_meta = {
                "face_id": face_idx,
                "area": round(face.Area, 6),
                "vertex_start": len(vertices),
                "triangle_start": len(indices),
                "triangle_count": len(faces)
            }
            face_metadata.append(face_meta)
            
            # Add vertices with their colors
            vert_offset = len(vertices)
            for v in verts:
                vertices.append(round_vec3(v))
                vertex_colors.append(color)  # Each vertex gets face color
            
            # Add triangles
            for tri in faces:
                if len(tri) >= 3:
                    indices.append([
                        tri[0] + vert_offset,
                        tri[1] + vert_offset,
                        tri[2] + vert_offset
                    ])
        
        except:
            continue
    
    if len(indices) == 0:
        return None
    
    return {
        "type": "faces",
        "vertices": vertices,
        "indices": indices,
        "vertex_colors": vertex_colors,  # Per-vertex colors
        "face_info": face_metadata  # Per-face metadata
    }

def process_edges(shape, body):
    """Extract edge geometry as line segments with FreeCAD visual properties"""
    vertices = []
    indices = []
    edge_metadata = []  # Metadata per edge
    
    # Get edge visual properties from FreeCAD
    line_color = get_line_color(body)
    line_width = get_line_width(body)
    alpha = get_transparency(body)
    
    for edge_idx, edge in enumerate(shape.Edges):
        try:
            # Discretize edge into points
            points = edge.discretize(Number=20)
            
            if len(points) < 2:
                continue
            
            vert_offset = len(vertices)
            
            # Edge metadata
            edge_meta = {
                "edge_id": edge_idx,
                "length": round(edge.Length, 6),
                "vertex_start": vert_offset,
                "segment_count": len(points) - 1
            }
            edge_metadata.append(edge_meta)
            
            # Add vertices
            for p in points:
                vertices.append(round_vec3([p.x, p.y, p.z]))
            
            # Add line segments
            for i in range(len(points) - 1):
                indices.append([vert_offset + i, vert_offset + i + 1])
        
        except:
            continue
    
    if len(indices) == 0:
        return None
    
    return {
        "type": "edges",
        "vertices": vertices,
        "indices": indices,
        "color": round_color([line_color[0], line_color[1], line_color[2], alpha]),
        "line_width": line_width,
        "edge_info": edge_metadata
    }

def process_vertices(shape, body):
    """Extract vertex geometry as points with FreeCAD visual properties"""
    vertices = []
    
    # Get point visual properties from FreeCAD
    point_color = get_point_color(body)
    point_size = get_point_size(body)
    alpha = get_transparency(body)
    
    for vertex in shape.Vertexes:
        try:
            p = vertex.Point
            vertices.append(round_vec3([p.x, p.y, p.z]))
        except:
            continue
    
    if len(vertices) == 0:
        return None
    
    return {
        "type": "vertices",
        "vertices": vertices,
        "color": round_color([point_color[0], point_color[1], point_color[2], alpha]),
        "point_size": point_size
    }

def process_body(body, body_idx):
    """Process single body with all geometry types and visual properties"""
    if not hasattr(body, 'Shape') or not body.Shape:
        return None
    
    shape = body.Shape
    if len(shape.Faces) == 0:
        return None
    
    body_label = getattr(body, 'Label', f'Body_{body_idx}')
    
    print(f"\n[Body {body_idx}] {body_label}")
    print(f"  Faces: {len(shape.Faces)}, Edges: {len(shape.Edges)}, Vertices: {len(shape.Vertexes)}")
    
    # Extract visual properties
    visual_props = {
        "display_mode": get_display_mode(body),
        "lighting": get_lighting(body),
        "transparency": round(get_transparency(body), 3),
        "line_width": get_line_width(body),
        "point_size": get_point_size(body),
        "line_color": round_color(list(get_line_color(body)) + [1.0]),
        "point_color": round_color(list(get_point_color(body)) + [1.0])
    }
    
    print(f"  Visual: {visual_props['display_mode']}, Transparency: {visual_props['transparency']}")
    
    body_data = {
        "id": body_idx,
        "name": body_label,
        "visual": visual_props,
        "geometry": []
    }
    
    # Process faces
    if EXPORT_FACES:
        faces = process_faces(shape, body, body_idx)
        if faces:
            body_data["geometry"].append(faces)
            print(f"  ✓ Faces: {len(faces['vertices'])} verts, {len(faces['indices'])} triangles")
    
    # Process edges
    if EXPORT_EDGES:
        edges = process_edges(shape, body)
        if edges:
            body_data["geometry"].append(edges)
            print(f"  ✓ Edges: {len(edges['vertices'])} verts, {len(edges['indices'])} lines, width: {edges['line_width']}")
    
    # Process vertices
    if EXPORT_VERTICES:
        vertices = process_vertices(shape, body)
        if vertices:
            body_data["geometry"].append(vertices)
            print(f"  ✓ Vertices: {len(vertices['vertices'])} points, size: {vertices['point_size']}")
    
    if len(body_data["geometry"]) == 0:
        return None
    
    return body_data

# ============================================================================
# FILE WRITING
# ============================================================================

def write_threejs_file(bodies, output_path, metadata):
    """Write Three.js optimized format"""
    
    data = {
        "version": "1.0",
        "format": "threejs_mesh",
        "metadata": {
            "source": STEP_FILE.name,
            "date": datetime.now().isoformat(),
            "tolerance": MESH_TOLERANCE,
            "generator": "FreeCAD STEP Converter"
        },
        "stats": metadata["stats"],
        "bodies": bodies
    }
    
    # Convert to JSON (compact)
    json_str = json.dumps(data, separators=(',', ':'))
    
    if USE_GZIP:
        # Write compressed
        with gzip.open(output_path, 'wt', encoding='utf-8') as f:
            f.write(json_str)
        print(f"\n✓ File written (gzipped): {output_path}")
    else:
        # Write plain
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(json_str)
        print(f"\n✓ File written: {output_path}")
    
    # Show sizes
    actual_size = output_path.stat().st_size
    uncompressed_size = len(json_str.encode('utf-8'))
    
    print(f"\nFile Sizes:")
    print(f"  Compressed: {actual_size / 1024:.2f} KB")
    print(f"  Uncompressed: {uncompressed_size / 1024:.2f} KB")
    
    if USE_GZIP:
        ratio = (1 - actual_size / uncompressed_size) * 100
        print(f"  Compression: {ratio:.1f}%")
    
    # Write usage example
    example_path = output_path.parent / "load_example.js"
    with open(example_path, 'w') as f:
        f.write("""
// Three.js Loading Example - Per-Vertex Colors
import * as THREE from 'three';
import pako from 'pako'; // For gzip decompression

async function loadMesh(url) {
    // Fetch and decompress
    const response = await fetch(url);
    const compressed = await response.arrayBuffer();
    const decompressed = pako.ungzip(compressed, { to: 'string' });
    const data = JSON.parse(decompressed);
    
    const group = new THREE.Group();
    
    // Load each body
    for (const body of data.bodies) {
        const bodyGroup = new THREE.Group();
        bodyGroup.name = body.name;
        
        // Process each geometry type
        for (const geom of body.geometry) {
            if (geom.type === 'faces') {
                // Create mesh with vertex colors
                const vertices = new Float32Array(geom.vertices.flat());
                const indices = new Uint32Array(geom.indices.flat());
                const colors = new Float32Array(geom.vertex_colors.flat());
                
                const geometry = new THREE.BufferGeometry();
                geometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3));
                geometry.setAttribute('color', new THREE.BufferAttribute(colors, 4)); // RGBA
                geometry.setIndex(new THREE.BufferAttribute(indices, 1));
                geometry.computeVertexNormals();
                
                // Material with vertex colors enabled
                const material = new THREE.MeshStandardMaterial({
                    vertexColors: true,  // Use vertex colors!
                    transparent: true,   // Enable transparency
                    side: THREE.DoubleSide,
                    metalness: 0.1,
                    roughness: 0.8
                });
                
                const mesh = new THREE.Mesh(geometry, material);
                mesh.name = body.name + '_faces';
                bodyGroup.add(mesh);
            }
            
            if (geom.type === 'edges') {
                // Create line segments
                const vertices = new Float32Array(geom.vertices.flat());
                const indices = new Uint32Array(geom.indices.flat());
                
                const geometry = new THREE.BufferGeometry();
                geometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3));
                geometry.setIndex(new THREE.BufferAttribute(indices, 1));
                
                const material = new THREE.LineBasicMaterial({
                    color: new THREE.Color(geom.color[0], geom.color[1], geom.color[2]),
                    opacity: geom.color[3],
                    transparent: geom.color[3] < 1.0,
                    linewidth: geom.line_width || 1  // Use FreeCAD line width
                });
                
                const lines = new THREE.LineSegments(geometry, material);
                lines.name = body.name + '_edges';
                bodyGroup.add(lines);
            }
            
            if (geom.type === 'vertices') {
                // Create points
                const vertices = new Float32Array(geom.vertices.flat());
                
                const geometry = new THREE.BufferGeometry();
                geometry.setAttribute('position', new THREE.BufferAttribute(vertices, 3));
                
                const material = new THREE.PointsMaterial({
                    color: new THREE.Color(geom.color[0], geom.color[1], geom.color[2]),
                    size: geom.point_size || 2.0,  // Use FreeCAD point size
                    sizeAttenuation: true
                });
                
                const points = new THREE.Points(geometry, material);
                points.name = body.name + '_vertices';
                bodyGroup.add(points);
            }
        }
        
        group.add(bodyGroup);
    }
    
    return group;
}

// Usage Example
loadMesh('freecad.mesh.gz').then(group => {
    scene.add(group);
    
    // Center and scale
    const box = new THREE.Box3().setFromObject(group);
    const center = box.getCenter(new THREE.Vector3());
    group.position.sub(center);
    
    const size = box.getSize(new THREE.Vector3());
    const maxDim = Math.max(size.x, size.y, size.z);
    const scale = 10 / maxDim;
    group.scale.setScalar(scale);
    
    // Access specific bodies
    const anaGovde = group.getObjectByName('ANA GOVDE');
    const uzunYarim = group.getObjectByName('UZUN YARIM');
    
    // Toggle edges visibility
    group.traverse(obj => {
        if (obj.name.includes('_edges')) {
            obj.visible = true; // Show/hide edges
        }
    });
    
    // Change body visibility
    anaGovde.visible = false;
});
""")
    
    print(f"✓ Usage example: {example_path}")

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

def main():
    print("\n" + "=" * 80)
    print("Configuration")
    print("=" * 80)
    print(f"Input: {STEP_FILE}")
    print(f"Tolerance: {MESH_TOLERANCE}")
    print(f"Export Faces: {EXPORT_FACES}")
    print(f"Export Edges: {EXPORT_EDGES}")
    print(f"Export Vertices: {EXPORT_VERTICES}")
    print(f"Compression: {USE_GZIP}")
    
    if not STEP_FILE.exists():
        print("\n✗ File not found!")
        return 1
    
    # Output path
    if OUTPUT_FILE:
        output_path = OUTPUT_FILE
    else:
        ext = '.mesh.gz' if USE_GZIP else '.mesh.json'
        output_path = STEP_FILE.with_suffix(ext)
    
    print(f"Output: {output_path}")
    
    # Load STEP
    print("\n" + "=" * 80)
    print("Loading STEP File")
    print("=" * 80)
    
    Import.open(str(STEP_FILE))
    doc = FreeCAD.ActiveDocument
    
    if not doc:
        print("✗ Failed to open")
        return 1
    
    print(f"✓ Document: {doc.Name}")
    print(f"✓ Objects: {len(doc.Objects)}")
    
    # Process bodies
    print("\n" + "=" * 80)
    print("Processing Bodies")
    print("=" * 80)
    
    bodies = []
    stats = {
        "bodies": 0,
        "total_vertices": 0,
        "total_triangles": 0,
        "total_edges": 0
    }
    
    for i, obj in enumerate(doc.Objects):
        body_data = process_body(obj, i)
        
        if body_data:
            bodies.append(body_data)
            stats["bodies"] += 1
            
            for geom in body_data["geometry"]:
                if geom["type"] == "faces":
                    stats["total_vertices"] += len(geom["vertices"])
                    stats["total_triangles"] += len(geom["indices"])
                elif geom["type"] == "edges":
                    stats["total_edges"] += len(geom["indices"])
    
    if len(bodies) == 0:
        print("\n✗ No geometry found!")
        FreeCAD.closeDocument(doc.Name)
        return 1
    
    # Write file
    print("\n" + "=" * 80)
    print("Writing Output")
    print("=" * 80)
    
    metadata = {"stats": stats}
    write_threejs_file(bodies, output_path, metadata)
    
    # Summary
    print("\n" + "=" * 80)
    print("SUCCESS")
    print("=" * 80)
    print(f"✓ Bodies: {stats['bodies']}")
    print(f"✓ Vertices: {stats['total_vertices']:,}")
    print(f"✓ Triangles: {stats['total_triangles']:,}")
    print(f"✓ Edges: {stats['total_edges']:,}")
    print(f"✓ Output: {output_path}")
    
    FreeCAD.closeDocument(doc.Name)
    return 0

if __name__ == "__main__":
    # Setup logging
    log_file = STEP_FILE.parent / "conversion_log.txt"
    
    class Logger:
        def __init__(self, filename):
            self.terminal = sys.stdout
            self.log = open(filename, "w", encoding="utf-8")
        
        def write(self, message):
            self.terminal.write(message)
            self.log.write(message)
            self.log.flush()
        
        def flush(self):
            pass
    
    sys.stdout = Logger(log_file)
    
    try:
        exit_code = main()
        print(f"\nLog: {log_file}")
        sys.exit(exit_code)
    except Exception as e:
        print(f"\n✗ ERROR: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)