"""
NX Journal Script: Extract Expressions from .sim File

This script:
1. Opens a .sim file
2. Extracts all expressions from the .sim and loaded .prt files
3. Saves expression data to JSON for the dashboard

Usage:
    run_journal.exe extract_expressions.py <sim_file_path> <output_json_path>
"""

import sys
import json
import NXOpen


def extract_all_expressions(sim_file_path, output_file_path):
    """
    Extract all expressions from .sim file and loaded parts.

    Args:
        sim_file_path: Path to .sim file
        output_file_path: Path to save JSON output
    """
    try:
        # Get NX session
        session = NXOpen.Session.GetSession()

        # Open the .sim file
        print(f"Opening .sim file: {sim_file_path}")
        part_load_status = None
        base_part, part_load_status = session.Parts.OpenBaseDisplay(sim_file_path)

        if part_load_status:
            part_load_status.Dispose()

        # Collect all expressions from all loaded parts
        all_expressions = {}

        # Get work parts and components
        parts_to_scan = [base_part]

        # Also scan all loaded components
        for component in base_part.ComponentAssembly.RootComponent.GetChildren():
            try:
                component_part = component.Prototype.OwningPart
                if component_part and component_part not in parts_to_scan:
                    parts_to_scan.append(component_part)
            except:
                pass

        # Extract expressions from each part
        for part in parts_to_scan:
            part_name = part.Name
            print(f"Scanning expressions from: {part_name}")

            expressions_list = []

            # Get all expressions
            for expr in part.Expressions:
                try:
                    expr_data = {
                        'name': expr.Name,
                        'value': expr.Value,
                        'formula': expr.Equation,
                        'units': expr.Units,
                        'type': 'number',  # Most expressions are numeric
                        'source_part': part_name
                    }

                    # Try to determine if it's a design variable candidate
                    # (not a formula, can be changed)
                    if '=' not in expr.Equation or expr.Equation.strip() == str(expr.Value):
                        expr_data['is_variable_candidate'] = True
                    else:
                        expr_data['is_variable_candidate'] = False

                    expressions_list.append(expr_data)

                except Exception as e:
                    print(f"Warning: Could not read expression {expr.Name}: {e}")
                    continue

            all_expressions[part_name] = expressions_list

        # Collect simulation metadata
        metadata = {
            'sim_file': sim_file_path,
            'base_part': base_part.Name,
            'num_components': len(parts_to_scan),
            'total_expressions': sum(len(exprs) for exprs in all_expressions.values()),
            'variable_candidates': sum(
                1 for exprs in all_expressions.values()
                for expr in exprs
                if expr.get('is_variable_candidate', False)
            )
        }

        # Prepare output
        output_data = {
            'metadata': metadata,
            'expressions_by_part': all_expressions
        }

        # Save to JSON
        print(f"Saving expressions to: {output_file_path}")
        with open(output_file_path, 'w') as f:
            json.dump(output_data, f, indent=2)

        print(f"Successfully extracted {metadata['total_expressions']} expressions")
        print(f"Found {metadata['variable_candidates']} potential design variables")

        # Close part
        base_part.Close(NXOpen.BasePart.CloseWholeTree.True,
                        NXOpen.BasePart.CloseModified.CloseModified, None)

        return True

    except Exception as e:
        error_data = {
            'error': str(e),
            'sim_file': sim_file_path
        }

        print(f"ERROR: {e}")

        with open(output_file_path, 'w') as f:
            json.dump(error_data, f, indent=2)

        return False


def main():
    """Main entry point for journal script."""
    if len(sys.argv) < 3:
        print("Usage: extract_expressions.py <sim_file_path> <output_json_path>")
        sys.exit(1)

    sim_file_path = sys.argv[1]
    output_file_path = sys.argv[2]

    print("="*60)
    print("NX Expression Extractor")
    print("="*60)

    success = extract_all_expressions(sim_file_path, output_file_path)

    if success:
        print("\nExpression extraction completed successfully!")
        sys.exit(0)
    else:
        print("\nExpression extraction failed!")
        sys.exit(1)


if __name__ == '__main__':
    main()