"""
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.

    Args:
        args: Command line arguments, args[0] should be the .sim file path
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print("Usage: run_journal.exe solve_simulation.py <sim_file_path>")
        return False

    sim_file_path = args[0]
    print(f"[JOURNAL] Opening simulation: {sim_file_path}")

    try:
        theSession = NXOpen.Session.GetSession()

        # 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
        partLoadStatus1.Dispose()

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

        simPart1 = workSimPart
        theSession.Post.UpdateUserGroupsFromSimPart(simPart1)

        # Reload all components to pick up parameter changes from .prt files
        print("[JOURNAL] Reloading components to pick up .prt parameter changes...")
        try:
            workSimPart.ComponentAssembly.ReloadComponents(
                NXOpen.Assemblies.ComponentAssembly.ReloadOption.AllLoaded
            )
            print("[JOURNAL] Components reloaded")
        except Exception as e:
            print(f"[JOURNAL] Warning: Could not reload components: {e}")

        # Make FEM work component (to ensure it's updated)
        # Find the FEM component - pattern: "COMPONENT <name>_fem1 1"
        markId1 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Make Work Part")

        # Get all components and find the FEM one
        rootComponent = workSimPart.ComponentAssembly.RootComponent
        femComponent = None

        for component in rootComponent.GetChildren():
            if "_fem" in component.DisplayName.lower():
                femComponent = component
                break

        if femComponent:
            print(f"[JOURNAL] Switching to FEM component: {femComponent.DisplayName}")

            # Make FEM the work component (this is what your recorded journal does)
            partLoadStatus2 = theSession.Parts.SetWorkComponent(
                femComponent,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible
            )
            partLoadStatus2.Dispose()

            # Get the FEM part and try to update it
            workFemPart = theSession.Parts.BaseWork
            try:
                # Try to update the FEM to pick up geometry/parameter changes
                print("[JOURNAL] Updating FEM to recognize parameter changes...")
                if hasattr(workFemPart, 'FemPart'):
                    workFemPart.FemPart.UpdateFeModel()
                    print("[JOURNAL] FEM updated")
            except Exception as e:
                print(f"[JOURNAL] Note: Could not update FEM (may not be necessary): {e}")

            # Switch back to sim part (this forces NX to recognize updates)
            markId2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Make Work Part")
            partLoadStatus3 = theSession.Parts.SetWorkComponent(
                NXOpen.Assemblies.Component.Null,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible
            )
            workSimPart = theSession.Parts.BaseWork
            partLoadStatus3.Dispose()
            print("[JOURNAL] Switched back to sim part")
        else:
            print("[JOURNAL] WARNING: No FEM component found, proceeding with solve anyway")

        # 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 first solution from the simulation
        simSimulation1 = workSimPart.FindObject("Simulation")
        simSolution1 = simSimulation1.FindObject("Solution[Solution 1]")

        psolutions1 = [simSolution1]

        # Solve in background mode
        numsolutionssolved1, numsolutionsfailed1, numsolutionsskipped1 = theCAESimSolveManager.SolveChainOfSolutions(
            psolutions1,
            NXOpen.CAE.SimSolution.SolveOption.Solve,
            NXOpen.CAE.SimSolution.SetupCheckOption.CompleteDeepCheckAndOutputErrors,
            NXOpen.CAE.SimSolution.SolveMode.Background
        )

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

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

        # NOTE: In Background mode, these values may not be accurate since the solve
        # runs asynchronously. The solve will continue after this journal finishes.
        # We rely on the Save operation and file existence checks to verify success.

        # 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)