"""
NX Journal: Comprehensive Part Introspection Tool
===================================================

This journal performs deep introspection of an NX .prt file and extracts:
- All expressions (user and internal, with values, units, formulas)
- Mass properties (mass, volume, surface area, center of gravity)
- Material properties (name, density, all material attributes)
- Body information (solid bodies, sheet bodies, body attributes)
- Part attributes (all user-defined attributes)
- Groups (all groups and their members)
- Features (all features in the part)
- References (linked parts, assembly components)
- Datum planes, coordinate systems
- Units system

Usage:
    run_journal.exe introspect_part.py <prt_file_path> [output_dir]

Output:
    _temp_introspection.json - Comprehensive JSON with all extracted data

Author: Atomizer
Created: 2025-12-19
Version: 1.0
"""

import sys
import os
import json
import NXOpen
import NXOpen.UF


def get_expressions(part):
    """Extract all expressions from the part."""
    expressions = {
        'user': [],
        'internal': [],
        'total_count': 0,
        'user_count': 0
    }

    try:
        for expr in part.Expressions:
            try:
                expr_data = {
                    'name': expr.Name,
                    'value': expr.Value,
                    'rhs': expr.RightHandSide if hasattr(expr, 'RightHandSide') else None,
                    'units': expr.Units.Name if expr.Units else None,
                    'type': str(expr.Type) if hasattr(expr, 'Type') else 'Unknown',
                }

                # Determine if internal (p0, p1, p123, etc.)
                name = expr.Name
                is_internal = False
                if name.startswith('p') and len(name) > 1:
                    rest = name[1:].replace('.', '').replace('_', '')
                    if rest.isdigit():
                        is_internal = True

                if is_internal:
                    expressions['internal'].append(expr_data)
                else:
                    expressions['user'].append(expr_data)

            except Exception as e:
                pass

        expressions['total_count'] = len(expressions['user']) + len(expressions['internal'])
        expressions['user_count'] = len(expressions['user'])

    except Exception as e:
        expressions['error'] = str(e)

    return expressions


def get_all_solid_bodies(part):
    """Get all solid bodies from the part."""
    bodies = []
    try:
        for body in part.Bodies:
            if body.IsSolidBody:
                bodies.append(body)
    except Exception as e:
        pass
    return bodies


def get_mass_properties(part, bodies):
    """Extract mass properties using MeasureManager."""
    results = {
        'mass_kg': 0.0,
        'mass_g': 0.0,
        'volume_mm3': 0.0,
        'surface_area_mm2': 0.0,
        'center_of_gravity_mm': [0.0, 0.0, 0.0],
        'num_bodies': len(bodies),
        'success': False
    }

    if not bodies:
        return results

    try:
        measureManager = part.MeasureManager
        bodyArray = list(bodies)

        # Build mass_units array
        uc = part.UnitCollection
        mass_units = [
            uc.GetBase("Area"),
            uc.GetBase("Volume"),
            uc.GetBase("Mass"),
            uc.GetBase("Length")
        ]

        measureBodies = measureManager.NewMassProperties(mass_units, 0.99, bodyArray)

        if measureBodies:
            try:
                results['mass_kg'] = measureBodies.Mass
                results['mass_g'] = results['mass_kg'] * 1000.0
            except:
                pass

            try:
                results['volume_mm3'] = measureBodies.Volume
            except:
                pass

            try:
                results['surface_area_mm2'] = measureBodies.Area
            except:
                pass

            try:
                cog = measureBodies.Centroid
                if cog:
                    results['center_of_gravity_mm'] = [cog.X, cog.Y, cog.Z]
            except:
                pass

            results['success'] = True

            try:
                measureBodies.Dispose()
            except:
                pass

    except Exception as e:
        results['error'] = str(e)

    return results


def get_materials(part, bodies):
    """Extract all materials from the part."""
    materials = {
        'assigned': [],
        'available': [],
        'library': []
    }

    # Get materials assigned to bodies
    for body in bodies:
        try:
            phys_mat = body.GetPhysicalMaterial()
            if phys_mat:
                mat_info = {
                    'name': phys_mat.Name,
                    'body': body.Name if hasattr(body, 'Name') else 'Unknown',
                    'properties': {}
                }

                # Try to get common material properties
                prop_names = ['Density', 'YoungModulus', 'PoissonRatio',
                              'ThermalExpansionCoefficient', 'ThermalConductivity',
                              'SpecificHeat', 'YieldStrength', 'UltimateStrength']
                for prop_name in prop_names:
                    try:
                        val = phys_mat.GetPropertyValue(prop_name)
                        if val is not None:
                            mat_info['properties'][prop_name] = float(val)
                    except:
                        pass

                materials['assigned'].append(mat_info)
        except:
            pass

    # Get all materials in part via PhysicalMaterialManager
    try:
        pmm = part.PhysicalMaterialManager
        if pmm:
            all_mats = pmm.GetAllPhysicalMaterials()
            for mat in all_mats:
                try:
                    mat_info = {
                        'name': mat.Name,
                        'properties': {}
                    }
                    prop_names = ['Density', 'YoungModulus', 'PoissonRatio']
                    for prop_name in prop_names:
                        try:
                            val = mat.GetPropertyValue(prop_name)
                            if val is not None:
                                mat_info['properties'][prop_name] = float(val)
                        except:
                            pass
                    materials['available'].append(mat_info)
                except:
                    pass
    except Exception as e:
        materials['pmm_error'] = str(e)

    return materials


def get_body_info(part):
    """Get detailed body information."""
    body_info = {
        'solid_bodies': [],
        'sheet_bodies': [],
        'counts': {
            'solid': 0,
            'sheet': 0,
            'total': 0
        }
    }

    try:
        for body in part.Bodies:
            body_data = {
                'name': body.Name if hasattr(body, 'Name') else 'Unknown',
                'is_solid': body.IsSolidBody,
                'is_sheet': body.IsSheetBody if hasattr(body, 'IsSheetBody') else False,
                'attributes': []
            }

            # Get body attributes
            try:
                attrs = body.GetUserAttributes()
                for attr in attrs:
                    try:
                        body_data['attributes'].append({
                            'title': attr.Title,
                            'type': str(attr.Type),
                            'value': attr.StringValue if hasattr(attr, 'StringValue') else str(attr.Value)
                        })
                    except:
                        pass
            except:
                pass

            if body.IsSolidBody:
                body_info['solid_bodies'].append(body_data)
                body_info['counts']['solid'] += 1
            else:
                body_info['sheet_bodies'].append(body_data)
                body_info['counts']['sheet'] += 1

        body_info['counts']['total'] = body_info['counts']['solid'] + body_info['counts']['sheet']

    except Exception as e:
        body_info['error'] = str(e)

    return body_info


def get_part_attributes(part):
    """Get all part-level attributes."""
    attributes = []

    try:
        attrs = part.GetUserAttributes()
        for attr in attrs:
            try:
                attr_data = {
                    'title': attr.Title,
                    'type': str(attr.Type),
                }

                # Get value based on type
                try:
                    if hasattr(attr, 'StringValue'):
                        attr_data['value'] = attr.StringValue
                    elif hasattr(attr, 'Value'):
                        attr_data['value'] = attr.Value
                    elif hasattr(attr, 'IntegerValue'):
                        attr_data['value'] = attr.IntegerValue
                    elif hasattr(attr, 'RealValue'):
                        attr_data['value'] = attr.RealValue
                except:
                    attr_data['value'] = 'Unknown'

                attributes.append(attr_data)
            except:
                pass

    except Exception as e:
        pass

    return attributes


def get_groups(part):
    """Get all groups in the part."""
    groups = []

    try:
        # NX stores groups in a collection
        if hasattr(part, 'Groups'):
            for group in part.Groups:
                try:
                    group_data = {
                        'name': group.Name if hasattr(group, 'Name') else 'Unknown',
                        'member_count': 0,
                        'members': []
                    }

                    # Try to get group members
                    try:
                        members = group.GetMembers()
                        group_data['member_count'] = len(members) if members else 0
                        for member in members[:10]:  # Limit to first 10 for readability
                            try:
                                group_data['members'].append(str(type(member).__name__))
                            except:
                                pass
                    except:
                        pass

                    groups.append(group_data)
                except:
                    pass
    except Exception as e:
        pass

    return groups


def get_features(part):
    """Get summary of features in the part."""
    features = {
        'total_count': 0,
        'by_type': {},
        'first_10': []
    }

    try:
        count = 0
        for feature in part.Features:
            try:
                feat_type = str(type(feature).__name__)

                # Count by type
                if feat_type in features['by_type']:
                    features['by_type'][feat_type] += 1
                else:
                    features['by_type'][feat_type] = 1

                # Store first 10 for reference
                if count < 10:
                    features['first_10'].append({
                        'name': feature.Name if hasattr(feature, 'Name') else 'Unknown',
                        'type': feat_type
                    })

                count += 1
            except:
                pass

        features['total_count'] = count

    except Exception as e:
        features['error'] = str(e)

    return features


def get_datums(part):
    """Get datum planes and coordinate systems."""
    datums = {
        'planes': [],
        'csys': [],
        'axes': []
    }

    try:
        # Datum planes
        if hasattr(part, 'Datums'):
            for datum in part.Datums:
                try:
                    datum_type = str(type(datum).__name__)
                    datum_name = datum.Name if hasattr(datum, 'Name') else 'Unknown'

                    if 'Plane' in datum_type:
                        datums['planes'].append(datum_name)
                    elif 'Csys' in datum_type or 'Coordinate' in datum_type:
                        datums['csys'].append(datum_name)
                    elif 'Axis' in datum_type:
                        datums['axes'].append(datum_name)
                except:
                    pass
    except Exception as e:
        datums['error'] = str(e)

    return datums


def get_units_info(part):
    """Get unit system information."""
    units_info = {
        'base_units': {},
        'system': 'Unknown'
    }

    try:
        uc = part.UnitCollection

        # Get common base units
        unit_types = ['Length', 'Mass', 'Time', 'Temperature', 'Angle',
                      'Area', 'Volume', 'Force', 'Pressure', 'Density']
        for unit_type in unit_types:
            try:
                base_unit = uc.GetBase(unit_type)
                if base_unit:
                    units_info['base_units'][unit_type] = base_unit.Name
            except:
                pass

        # Determine system from length unit
        if 'Length' in units_info['base_units']:
            length_unit = units_info['base_units']['Length'].lower()
            if 'mm' in length_unit or 'millimeter' in length_unit:
                units_info['system'] = 'Metric (mm)'
            elif 'meter' in length_unit and 'milli' not in length_unit:
                units_info['system'] = 'Metric (m)'
            elif 'inch' in length_unit or 'in' in length_unit:
                units_info['system'] = 'Imperial (inch)'

    except Exception as e:
        units_info['error'] = str(e)

    return units_info


def get_linked_parts(theSession, working_dir):
    """Get information about linked/associated parts."""
    linked_parts = {
        'loaded_parts': [],
        'fem_parts': [],
        'sim_parts': [],
        'idealized_parts': []
    }

    try:
        for part in theSession.Parts:
            try:
                part_name = part.Name if hasattr(part, 'Name') else str(part)
                part_path = part.FullPath if hasattr(part, 'FullPath') else 'Unknown'

                part_info = {
                    'name': part_name,
                    'path': part_path,
                    'leaf_name': part.Leaf if hasattr(part, 'Leaf') else part_name
                }

                name_lower = part_name.lower()
                if '_sim' in name_lower or name_lower.endswith('.sim'):
                    linked_parts['sim_parts'].append(part_info)
                elif '_fem' in name_lower or name_lower.endswith('.fem'):
                    if '_i.prt' in name_lower or '_i' in name_lower:
                        linked_parts['idealized_parts'].append(part_info)
                    else:
                        linked_parts['fem_parts'].append(part_info)
                elif '_i.prt' in name_lower:
                    linked_parts['idealized_parts'].append(part_info)
                else:
                    linked_parts['loaded_parts'].append(part_info)

            except:
                pass

    except Exception as e:
        linked_parts['error'] = str(e)

    return linked_parts


def main(args):
    """Main entry point for NX journal."""

    if len(args) < 1:
        print("ERROR: No .prt file path provided")
        print("Usage: run_journal.exe introspect_part.py <prt_file> [output_dir]")
        return False

    prt_file_path = args[0]
    output_dir = args[1] if len(args) > 1 else os.path.dirname(prt_file_path)
    prt_filename = os.path.basename(prt_file_path)

    print(f"[INTROSPECT] " + "="*60)
    print(f"[INTROSPECT] NX COMPREHENSIVE PART INTROSPECTION")
    print(f"[INTROSPECT] " + "="*60)
    print(f"[INTROSPECT] Part: {prt_filename}")
    print(f"[INTROSPECT] Output: {output_dir}")

    results = {
        'part_file': prt_filename,
        'part_path': prt_file_path,
        'success': False,
        'error': None,
        'expressions': {},
        'mass_properties': {},
        'materials': {},
        'bodies': {},
        'attributes': [],
        'groups': [],
        'features': {},
        'datums': {},
        'units': {},
        'linked_parts': {}
    }

    try:
        theSession = NXOpen.Session.GetSession()

        # Set load options
        working_dir = os.path.dirname(prt_file_path)
        theSession.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory
        theSession.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])

        # Open the part file
        print(f"[INTROSPECT] Opening part file...")
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
            prt_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        partLoadStatus.Dispose()

        workPart = theSession.Parts.Work
        print(f"[INTROSPECT] Loaded: {workPart.Name}")

        # Extract all data
        print(f"[INTROSPECT] Extracting expressions...")
        results['expressions'] = get_expressions(workPart)
        print(f"[INTROSPECT]   Found {results['expressions']['user_count']} user expressions")

        print(f"[INTROSPECT] Extracting body info...")
        results['bodies'] = get_body_info(workPart)
        print(f"[INTROSPECT]   Found {results['bodies']['counts']['solid']} solid bodies")

        print(f"[INTROSPECT] Extracting mass properties...")
        bodies = get_all_solid_bodies(workPart)
        results['mass_properties'] = get_mass_properties(workPart, bodies)
        print(f"[INTROSPECT]   Mass: {results['mass_properties']['mass_kg']:.4f} kg")

        print(f"[INTROSPECT] Extracting materials...")
        results['materials'] = get_materials(workPart, bodies)
        print(f"[INTROSPECT]   Found {len(results['materials']['assigned'])} assigned materials")

        print(f"[INTROSPECT] Extracting attributes...")
        results['attributes'] = get_part_attributes(workPart)
        print(f"[INTROSPECT]   Found {len(results['attributes'])} part attributes")

        print(f"[INTROSPECT] Extracting groups...")
        results['groups'] = get_groups(workPart)
        print(f"[INTROSPECT]   Found {len(results['groups'])} groups")

        print(f"[INTROSPECT] Extracting features...")
        results['features'] = get_features(workPart)
        print(f"[INTROSPECT]   Found {results['features']['total_count']} features")

        print(f"[INTROSPECT] Extracting datums...")
        results['datums'] = get_datums(workPart)
        print(f"[INTROSPECT]   Found {len(results['datums']['planes'])} datum planes")

        print(f"[INTROSPECT] Extracting units...")
        results['units'] = get_units_info(workPart)
        print(f"[INTROSPECT]   System: {results['units']['system']}")

        print(f"[INTROSPECT] Extracting linked parts...")
        results['linked_parts'] = get_linked_parts(theSession, working_dir)
        print(f"[INTROSPECT]   Found {len(results['linked_parts']['loaded_parts'])} loaded parts")

        results['success'] = True
        print(f"[INTROSPECT] ")
        print(f"[INTROSPECT] INTROSPECTION COMPLETE!")
        print(f"[INTROSPECT] " + "="*60)

        # Summary
        print(f"[INTROSPECT] SUMMARY:")
        print(f"[INTROSPECT]   Expressions: {results['expressions']['user_count']} user, {len(results['expressions']['internal'])} internal")
        print(f"[INTROSPECT]   Mass: {results['mass_properties']['mass_kg']:.4f} kg ({results['mass_properties']['mass_g']:.2f} g)")
        print(f"[INTROSPECT]   Bodies: {results['bodies']['counts']['solid']} solid, {results['bodies']['counts']['sheet']} sheet")
        print(f"[INTROSPECT]   Features: {results['features']['total_count']}")
        print(f"[INTROSPECT]   Materials: {len(results['materials']['assigned'])} assigned")

    except Exception as e:
        results['error'] = str(e)
        results['success'] = False
        print(f"[INTROSPECT] FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()

    # Write results
    output_file = os.path.join(output_dir, "_temp_introspection.json")
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)
    print(f"[INTROSPECT] Results written to: {output_file}")

    return results['success']


if __name__ == '__main__':
    main(sys.argv[1:])
