"""
NX Journal Script to Solve Simulation in Batch Mode

This script opens a .sim file, updates the FEM, and solves it through the NX API.
Usage: run_journal.exe solve_simulation.py <sim_file_path>

Based on recorded NX journal pattern for solving simulations.
"""

import sys
import NXOpen
import NXOpen.Assemblies
import NXOpen.CAE


def main(args):
    """
    Open and solve a simulation file with updated expression values.

    Args:
        args: Command line arguments
              args[0]: .sim file path
              args[1]: solution_name (optional, e.g., "Solution_Normal_Modes" or None for default)
              args[2+]: expression updates as "name=value" pairs
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print("Usage: run_journal.exe solve_simulation.py <sim_file_path> [solution_name] [expr1=val1] [expr2=val2] ...")
        return False

    sim_file_path = args[0]

    # Parse solution name if provided (args[1])
    solution_name = args[1] if len(args) > 1 and args[1] != 'None' else None

    # Extract base name from sim file (e.g., "Beam_sim1.sim" -> "Beam")
    import os
    sim_filename = os.path.basename(sim_file_path)
    part_base_name = sim_filename.split('_sim')[0] if '_sim' in sim_filename else sim_filename.split('.sim')[0]

    # Parse expression updates from args[2+] as "name=value" pairs
    expression_updates = {}
    for arg in args[2:]:
        if '=' in arg:
            name, value = arg.split('=', 1)
            expression_updates[name] = float(value)

    print(f"[JOURNAL] Opening simulation: {sim_file_path}")
    print(f"[JOURNAL] Detected part base name: {part_base_name}")
    if solution_name:
        print(f"[JOURNAL] Will solve specific solution: {solution_name}")
    else:
        print(f"[JOURNAL] Will solve default solution (Solution 1)")
    if expression_updates:
        print(f"[JOURNAL] Will update expressions:")
        for name, value in expression_updates.items():
            print(f"[JOURNAL]   {name} = {value}")

    try:
        theSession = NXOpen.Session.GetSession()

        # Set load options to load linked parts from directory
        print("[JOURNAL] Setting load options for linked parts...")
        import os
        working_dir = os.path.dirname(os.path.abspath(sim_file_path))

        # Complete load options setup (from recorded journal)
        theSession.Parts.LoadOptions.LoadLatest = False
        theSession.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory

        searchDirectories = [working_dir]
        searchSubDirs = [True]
        theSession.Parts.LoadOptions.SetSearchDirectories(searchDirectories, searchSubDirs)

        theSession.Parts.LoadOptions.ComponentsToLoad = NXOpen.LoadOptions.LoadComponents.All
        theSession.Parts.LoadOptions.PartLoadOption = NXOpen.LoadOptions.LoadOption.FullyLoad
        theSession.Parts.LoadOptions.SetInterpartData(True, NXOpen.LoadOptions.Parent.All)
        theSession.Parts.LoadOptions.AllowSubstitution = False
        theSession.Parts.LoadOptions.GenerateMissingPartFamilyMembers = True
        theSession.Parts.LoadOptions.AbortOnFailure = False

        referenceSets = ["As Saved", "Use Simplified", "Use Model", "Entire Part", "Empty"]
        theSession.Parts.LoadOptions.SetDefaultReferenceSets(referenceSets)
        theSession.Parts.LoadOptions.ReferenceSetOverride = False

        print(f"[JOURNAL] Load directory set to: {working_dir}")

        # Close any currently open sim file to force reload from disk
        print("[JOURNAL] Checking for open parts...")
        try:
            current_work = theSession.Parts.BaseWork
            if current_work and hasattr(current_work, 'FullPath'):
                current_path = current_work.FullPath
                print(f"[JOURNAL] Closing currently open part: {current_path}")
                # Close without saving (we want to reload from disk)
                partCloseResponses1 = [NXOpen.BasePart.CloseWholeTree]
                theSession.Parts.CloseAll(partCloseResponses1)
                print("[JOURNAL] Parts closed")
        except Exception as e:
            print(f"[JOURNAL] No parts to close or error closing: {e}")

        # Open the .sim file (now will load fresh from disk with updated .prt files)
        print(f"[JOURNAL] Opening simulation fresh from disk...")
        basePart1, partLoadStatus1 = theSession.Parts.OpenActiveDisplay(
            sim_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )

        workSimPart = theSession.Parts.BaseWork
        displaySimPart = theSession.Parts.BaseDisplay

        print(f"[JOURNAL] Simulation opened successfully")
        partLoadStatus1.Dispose()

        # Switch to simulation application
        theSession.ApplicationSwitchImmediate("UG_APP_SFEM")

        simPart1 = workSimPart
        theSession.Post.UpdateUserGroupsFromSimPart(simPart1)

        # STEP 1: Try to switch to part and update expressions (optional for some models)
        print(f"[JOURNAL] STEP 1: Checking for {part_base_name}.prt geometry...")
        geometry_updated = False
        try:
            # Find the main part (may not exist for embedded geometry models)
            bracketPart = None
            try:
                bracketPart = theSession.Parts.FindObject(part_base_name)
            except:
                pass

            if bracketPart:
                print(f"[JOURNAL]   Found {part_base_name} part, updating geometry...")
                # Make Bracket the active display part
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    bracketPart,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.UseLast
                )
                partLoadStatus.Dispose()

                workPart = theSession.Parts.Work

                # CRITICAL: Apply expression changes BEFORE updating geometry
                expressions_updated = []

                # Apply all expression updates dynamically
                for expr_name, expr_value in expression_updates.items():
                    print(f"[JOURNAL]   Applying {expr_name} = {expr_value}")
                    try:
                        expr_obj = workPart.Expressions.FindObject(expr_name)
                        if expr_obj:
                            # Use millimeters as default unit for geometric parameters
                            unit_mm = workPart.UnitCollection.FindObject("MilliMeter")
                            workPart.Expressions.EditExpressionWithUnits(expr_obj, unit_mm, str(expr_value))
                            expressions_updated.append(expr_obj)
                            print(f"[JOURNAL]     {expr_name} updated successfully")
                        else:
                            print(f"[JOURNAL]     WARNING: {expr_name} expression not found!")
                    except Exception as e:
                        print(f"[JOURNAL]     ERROR updating {expr_name}: {e}")

                # Make expressions up to date
                if expressions_updated:
                    print(f"[JOURNAL]   Making {len(expressions_updated)} expression(s) up to date...")
                    for expr in expressions_updated:
                        markId_expr = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Make Up to Date")
                        objects1 = [expr]
                        theSession.UpdateManager.MakeUpToDate(objects1, markId_expr)
                        theSession.DeleteUndoMark(markId_expr, None)

                # CRITICAL: Update the geometry model - rebuilds features with new expressions
                print(f"[JOURNAL]   Rebuilding geometry with new expression values...")
                markId_update = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "NX update")
                nErrs = theSession.UpdateManager.DoUpdate(markId_update)
                theSession.DeleteUndoMark(markId_update, "NX update")
                print(f"[JOURNAL]   {part_base_name} geometry updated ({nErrs} errors)")

                # Extract mass from expression p173 if it exists and write to temp file
                try:
                    mass_expr = workPart.Expressions.FindObject("p173")
                    if mass_expr:
                        mass_kg = mass_expr.Value
                        mass_output_file = os.path.join(working_dir, "_temp_mass.txt")
                        with open(mass_output_file, 'w') as f:
                            f.write(str(mass_kg))
                        print(f"[JOURNAL]   Mass from p173: {mass_kg:.6f} kg ({mass_kg * 1000:.2f} g)")
                        print(f"[JOURNAL]   Mass written to: {mass_output_file}")
                except:
                    pass  # Expression p173 might not exist in all models

                geometry_updated = True
            else:
                print(f"[JOURNAL]   {part_base_name} part not found - may be embedded in sim file")
        except Exception as e:
            print(f"[JOURNAL]   Could not update {part_base_name}.prt: {e}")
            print(f"[JOURNAL]   Continuing with sim-only solve...")

        # STEP 2: Try to switch to FEM part and update (optional for some models)
        fem_part_name = f"{part_base_name}_fem1"
        print(f"[JOURNAL] STEP 2: Checking for {fem_part_name}.fem...")
        fem_updated = False
        try:
            # Find the FEM part (may not exist or may have different name)
            femPart1 = None
            try:
                femPart1 = theSession.Parts.FindObject(fem_part_name)
            except:
                # Try with _i suffix for idealized FEM
                try:
                    femPart1 = theSession.Parts.FindObject(f"{fem_part_name}_i")
                except:
                    pass

            if femPart1:
                print(f"[JOURNAL]   Found FEM part, updating...")
                # Make FEM the active display part
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    femPart1,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.SameAsDisplay
                )
                partLoadStatus.Dispose()

                workFemPart = theSession.Parts.BaseWork

                # CRITICAL: Update FE Model - regenerates FEM with new geometry
                print("[JOURNAL]   Updating FE Model...")
                fEModel1 = workFemPart.FindObject("FEModel")
                if fEModel1:
                    fEModel1.UpdateFemodel()
                    print("[JOURNAL]   FE Model updated with new geometry!")
                    fem_updated = True
                else:
                    print("[JOURNAL]   WARNING: Could not find FEModel object")
            else:
                print(f"[JOURNAL]   FEM part not found - may be embedded in sim file")
        except Exception as e:
            print(f"[JOURNAL]   Could not update FEM: {e}")
            print(f"[JOURNAL]   Continuing with sim-only solve...")

        # STEP 3: Switch back to sim part
        print("[JOURNAL] STEP 3: Switching back to sim part...")
        try:
            status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                simPart1,
                NXOpen.DisplayPartOption.AllowAdditional,
                NXOpen.PartDisplayPartWorkPartOption.UseLast
            )
            partLoadStatus.Dispose()
            workSimPart = theSession.Parts.BaseWork
            print("[JOURNAL]   Switched back to sim part")
        except Exception as e:
            print(f"[JOURNAL]   WARNING: Error switching to sim part: {e}")

        # Note: Old output files are deleted by nx_solver.py before calling this journal
        # This ensures NX performs a fresh solve

        # Solve the simulation
        print("[JOURNAL] Starting solve...")
        markId3 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId3, "Solve Dialog")

        markId5 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Solve")

        theCAESimSolveManager = NXOpen.CAE.SimSolveManager.GetSimSolveManager(theSession)

        # Get the simulation object
        simSimulation1 = workSimPart.FindObject("Simulation")

        # Get the solution(s) to solve - either specific or all
        if solution_name:
            # Solve specific solution in background mode
            solution_obj_name = f"Solution[{solution_name}]"
            print(f"[JOURNAL] Looking for solution: {solution_obj_name}")
            simSolution1 = simSimulation1.FindObject(solution_obj_name)
            psolutions1 = [simSolution1]

            numsolutionssolved1, numsolutionsfailed1, numsolutionsskipped1 = theCAESimSolveManager.SolveChainOfSolutions(
                psolutions1,
                NXOpen.CAE.SimSolution.SolveOption.Solve,
                NXOpen.CAE.SimSolution.SetupCheckOption.CompleteDeepCheckAndOutputErrors,
                NXOpen.CAE.SimSolution.SolveMode.Background
            )
        else:
            # Solve ALL solutions using SolveAllSolutions API (Foreground mode)
            # This ensures all solutions (static + modal, etc.) complete before returning
            print(f"[JOURNAL] Solving all solutions using SolveAllSolutions API (Foreground mode)...")

            numsolutionssolved1, numsolutionsfailed1, numsolutionsskipped1 = theCAESimSolveManager.SolveAllSolutions(
                NXOpen.CAE.SimSolution.SolveOption.Solve,
                NXOpen.CAE.SimSolution.SetupCheckOption.CompleteCheckAndOutputErrors,
                NXOpen.CAE.SimSolution.SolveMode.Foreground,
                False
            )

        theSession.DeleteUndoMark(markId5, None)
        theSession.SetUndoMarkName(markId3, "Solve")

        print(f"[JOURNAL] Solve completed!")
        print(f"[JOURNAL]   Solutions solved: {numsolutionssolved1}")
        print(f"[JOURNAL]   Solutions failed: {numsolutionsfailed1}")
        print(f"[JOURNAL]   Solutions skipped: {numsolutionsskipped1}")

        # NOTE: When solution_name=None, we use Foreground mode to ensure all solutions
        # complete before returning. When solution_name is specified, Background mode is used.

        # Save the simulation to write all output files
        print("[JOURNAL] Saving simulation to ensure output files are written...")
        simPart2 = workSimPart
        partSaveStatus1 = simPart2.Save(
            NXOpen.BasePart.SaveComponents.TrueValue,
            NXOpen.BasePart.CloseAfterSave.FalseValue
        )
        partSaveStatus1.Dispose()
        print("[JOURNAL] Save complete!")

        return True

    except Exception as e:
        print(f"[JOURNAL] ERROR: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == '__main__':
    success = main(sys.argv[1:])
    sys.exit(0 if success else 1)