"""
Mesh Converter Utility
Converts Nastran BDF/OP2 files to GLTF for web visualization
"""

import json
import numpy as np
from pathlib import Path
from typing import Optional, Dict, Any
import trimesh
from pyNastran.bdf.bdf import BDF
from pyNastran.op2.op2 import OP2

def convert_study_mesh(study_dir: Path) -> Optional[Path]:
    """
    Convert the mesh and results of a study to GLTF format.
    
    Args:
        study_dir: Path to the study directory
        
    Returns:
        Path to the generated GLTF file, or None if conversion failed
    """
    try:
        # Locate files
        setup_dir = study_dir / "1_setup" / "model"
        results_dir = study_dir / "2_results"
        vis_dir = study_dir / "3_visualization"
        vis_dir.mkdir(parents=True, exist_ok=True)
        
        # Find BDF/DAT file
        bdf_files = list(setup_dir.glob("*.dat")) + list(setup_dir.glob("*.bdf"))
        if not bdf_files:
            # Fallback: Generate placeholder if no BDF found
            return _generate_placeholder_mesh(vis_dir)
            
        bdf_path = bdf_files[0]
        
        # Find OP2 file (optional)
        op2_files = list(results_dir.glob("*.op2"))
        op2_path = op2_files[0] if op2_files else None
        
        # Load BDF
        model = BDF()
        model.read_bdf(bdf_path, xref=False)
        
        # Extract nodes and elements
        # This is a simplified extraction for shell/solid elements
        # A full implementation would handle all element types
        nodes = []
        node_ids = []
        for nid, node in model.nodes.items():
            nodes.append(node.get_position())
            node_ids.append(nid)
            
        nodes = np.array(nodes)
        node_map = {nid: i for i, nid in enumerate(node_ids)}
        
        faces = []
        
        # Process CQUAD4/CTRIA3 elements
        for eid, element in model.elements.items():
            if element.type == 'CQUAD4':
                n = [node_map[nid] for nid in element.nodes]
                faces.append([n[0], n[1], n[2]])
                faces.append([n[0], n[2], n[3]])
            elif element.type == 'CTRIA3':
                n = [node_map[nid] for nid in element.nodes]
                faces.append([n[0], n[1], n[2]])
                
        if not faces:
             # Fallback if no compatible elements found
            return _generate_placeholder_mesh(vis_dir)
            
        # Create mesh
        mesh = trimesh.Trimesh(vertices=nodes, faces=faces)
        
        # Map results if OP2 exists
        if op2_path:
            op2 = OP2()
            op2.read_op2(op2_path)
            
            # Example: Map displacement magnitude to vertex colors
            if 1 in op2.displacements:
                disp = op2.displacements[1]
                # Get last timestep
                t3 = disp.data[-1, :, :3] # Translation x,y,z
                mag = np.linalg.norm(t3, axis=1)
                
                # Normalize to 0-1 for coloring
                if mag.max() > mag.min():
                    norm_mag = (mag - mag.min()) / (mag.max() - mag.min())
                else:
                    norm_mag = np.zeros_like(mag)
                    
                # Apply colormap (simple blue-to-red)
                colors = np.zeros((len(nodes), 4))
                colors[:, 0] = norm_mag # R
                colors[:, 2] = 1 - norm_mag # B
                colors[:, 3] = 1.0 # Alpha
                
                mesh.visual.vertex_colors = colors
                
        # Export to GLTF
        output_path = vis_dir / "model.gltf"
        mesh.export(output_path)
        
        # Save metadata
        metadata = {
            "node_count": len(nodes),
            "element_count": len(faces),
            "has_results": op2_path is not None
        }
        with open(vis_dir / "model.json", 'w') as f:
            json.dump(metadata, f, indent=2)
            
        return output_path
        
    except Exception as e:
        print(f"Mesh conversion error: {e}")
        # Fallback on error
        return _generate_placeholder_mesh(vis_dir)

def _generate_placeholder_mesh(output_dir: Path) -> Path:
    """Generate a simple box mesh for testing"""
    mesh = trimesh.creation.box(extents=[10, 10, 10])
    output_path = output_dir / "model.gltf"
    mesh.export(output_path)
    
    with open(output_dir / "model.json", 'w') as f:
        json.dump({"placeholder": True}, f)
        
    return output_path