"""
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]: tip_thickness value (optional)
              args[2]: support_angle value (optional)
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print("Usage: run_journal.exe solve_simulation.py <sim_file_path> [tip_thickness] [support_angle]")
        return False

    sim_file_path = args[0]

    # 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 values if provided
    tip_thickness = float(args[1]) if len(args) > 1 else None
    support_angle = float(args[2]) if len(args) > 2 else None

    print(f"[JOURNAL] Opening simulation: {sim_file_path}")
    print(f"[JOURNAL] Detected part base name: {part_base_name}")
    if tip_thickness is not None:
        print(f"[JOURNAL] Will update tip_thickness = {tip_thickness}")
    if support_angle is not None:
        print(f"[JOURNAL] Will update support_angle = {support_angle}")

    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 = []

                if tip_thickness is not None:
                    print(f"[JOURNAL]   Applying tip_thickness = {tip_thickness}")
                    expr_tip = workPart.Expressions.FindObject("tip_thickness")
                    if expr_tip:
                        unit_mm = workPart.UnitCollection.FindObject("MilliMeter")
                        workPart.Expressions.EditExpressionWithUnits(expr_tip, unit_mm, str(tip_thickness))
                        expressions_updated.append(expr_tip)
                        print(f"[JOURNAL]     tip_thickness updated")
                    else:
                        print(f"[JOURNAL]     WARNING: tip_thickness expression not found!")

                if support_angle is not None:
                    print(f"[JOURNAL]   Applying support_angle = {support_angle}")
                    expr_angle = workPart.Expressions.FindObject("support_angle")
                    if expr_angle:
                        unit_deg = workPart.UnitCollection.FindObject("Degrees")
                        workPart.Expressions.EditExpressionWithUnits(expr_angle, unit_deg, str(support_angle))
                        expressions_updated.append(expr_angle)
                        print(f"[JOURNAL]     support_angle updated")
                    else:
                        print(f"[JOURNAL]     WARNING: support_angle expression not found!")

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