"""
NX Journal Script to Extract All Expressions from a Part

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

Output:
    _temp_expressions.json with all expressions from the part
"""

import sys
import os
import json
import NXOpen


def main(args):
    if len(args) < 1:
        print("ERROR: No .prt file path provided")
        return False

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

    print(f"[JOURNAL] Extracting expressions from: {os.path.basename(prt_file_path)}")

    results = {
        'part_file': os.path.basename(prt_file_path),
        'part_path': prt_file_path,
        'expressions': [],
        'expression_count': 0,
        'user_expression_count': 0,
        'success': False,
        'error': None
    }

    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"[JOURNAL] Opening part file...")
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
            prt_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        partLoadStatus.Dispose()

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

        # Extract all expressions
        print(f"[JOURNAL] Extracting expressions...")
        for expr in workPart.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',
                }

                # Check if it's a user expression (not internal p0, p1, etc.)
                is_internal = expr.Name.startswith('p') and len(expr.Name) > 1 and expr.Name[1:].replace('.', '').replace('_', '').isdigit()
                expr_data['is_internal'] = is_internal

                results['expressions'].append(expr_data)

                if not is_internal:
                    results['user_expression_count'] += 1

            except Exception as e:
                print(f"[JOURNAL] Warning: Could not read expression: {e}")

        results['expression_count'] = len(results['expressions'])
        results['success'] = True

        print(f"[JOURNAL] Found {results['expression_count']} total expressions")
        print(f"[JOURNAL] Found {results['user_expression_count']} user expressions")

        # Print user expressions
        print(f"\n[JOURNAL] USER EXPRESSIONS:")
        print(f"[JOURNAL] " + "=" * 50)
        for expr in results['expressions']:
            if not expr['is_internal']:
                units_str = f" [{expr['units']}]" if expr['units'] else ""
                print(f"[JOURNAL]   {expr['name']}: {expr['value']}{units_str}")

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

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

    return results['success']


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