Atomizer/optimization_engine/utils/nx_session_manager.py

"""
NX Session Manager - Tracks NX sessions to prevent accidental closure of user sessions.

This module ensures that Atomizer only closes NX sessions that it explicitly started,
never touching sessions that were already running or started by the user.

Key Features:
- Automatically starts NX if not running
- Tracks PID of the NX instance we started
- Only closes NX instances we started (never user's sessions)
- Waits for NX to fully load before proceeding
"""

import os
import json
import time
import subprocess
from pathlib import Path
from typing import Optional, Tuple, List


# Session tracking file location
SESSION_LOCK_DIR = Path(os.environ.get('TEMP', '/tmp')) / 'atomizer_nx_sessions'

# Default NX installation paths (in order of preference)
DEFAULT_NX_PATHS = [
    Path(r"C:\Program Files\Siemens\DesigncenterNX2512\NXBIN\ugraf.exe"),  # DesignCenter (preferred)
    Path(r"C:\Program Files\Siemens\NX2506\NXBIN\ugraf.exe"),
    Path(r"C:\Program Files\Siemens\NX2412\NXBIN\ugraf.exe"),
    Path(r"C:\Program Files\Siemens\Simcenter3D_2506\NXBIN\ugraf.exe"),
    Path(r"C:\Program Files\Siemens\Simcenter3D_2412\NXBIN\ugraf.exe"),
]


class NXSessionManager:
    """
    Manages NX sessions to ensure we only close sessions we started.

    Usage:
        manager = NXSessionManager()

        # Check if NX is already running before starting optimization
        was_running = manager.is_nx_running()

        # Start NX if needed (records that we started it)
        if not was_running:
            manager.start_nx()

        # ... run optimization ...

        # Only close NX if we started it
        if manager.can_close_nx():
            manager.close_nx()
    """

    def __init__(self, session_id: Optional[str] = None, nx_exe_path: Optional[Path] = None):
        """
        Initialize the session manager.

        Args:
            session_id: Unique identifier for this optimization session.
                       If None, uses process ID.
            nx_exe_path: Path to ugraf.exe. If None, auto-detects from common locations.
        """
        self.session_id = session_id or f"atomizer_{os.getpid()}"
        self.lock_file = SESSION_LOCK_DIR / f"{self.session_id}.lock"
        self._nx_was_running_before_start = None
        self._started_nx_pid = None

        # Find NX executable
        self.nx_exe_path = nx_exe_path or self._find_nx_executable()

        # Ensure lock directory exists
        SESSION_LOCK_DIR.mkdir(parents=True, exist_ok=True)

    def _find_nx_executable(self) -> Optional[Path]:
        """Find NX executable from common installation paths."""
        for path in DEFAULT_NX_PATHS:
            if path.exists():
                return path
        return None

    def is_nx_running(self) -> bool:
        """Check if NX (ugraf.exe) is currently running."""
        try:
            result = subprocess.run(
                ['tasklist', '/FI', 'IMAGENAME eq ugraf.exe'],
                capture_output=True,
                text=True,
                timeout=10
            )
            return 'ugraf.exe' in result.stdout.lower()
        except Exception:
            return False

    def get_nx_pids(self) -> list:
        """Get list of all NX process IDs currently running."""
        pids = []
        try:
            result = subprocess.run(
                ['tasklist', '/FI', 'IMAGENAME eq ugraf.exe', '/FO', 'CSV'],
                capture_output=True,
                text=True,
                timeout=10
            )
            for line in result.stdout.strip().split('\n')[1:]:  # Skip header
                if 'ugraf.exe' in line.lower():
                    parts = line.strip('"').split('","')
                    if len(parts) >= 2:
                        try:
                            pids.append(int(parts[1]))
                        except ValueError:
                            pass
        except Exception:
            pass
        return pids

    def start_nx(self, wait_for_ready: bool = True, timeout_seconds: int = 60) -> Optional[int]:
        """
        Start a new NX instance and track its PID.

        Args:
            wait_for_ready: If True, wait for NX to fully load before returning
            timeout_seconds: Maximum time to wait for NX to start (default: 60s)

        Returns:
            PID of the started NX process, or None if failed

        Raises:
            RuntimeError: If NX executable not found or failed to start
        """
        if not self.nx_exe_path or not self.nx_exe_path.exists():
            raise RuntimeError(
                f"NX executable not found. Checked paths:\n"
                f"{chr(10).join(str(p) for p in DEFAULT_NX_PATHS)}\n"
                f"Please install NX or specify the path manually."
            )

        # Get PIDs before starting
        pids_before = set(self.get_nx_pids())

        print(f"[NX SESSION] Starting NX from: {self.nx_exe_path}")
        print(f"[NX SESSION]   Existing NX PIDs: {pids_before or 'none'}")

        try:
            # Start NX as a detached process (won't block)
            # Using subprocess.Popen with CREATE_NEW_PROCESS_GROUP
            process = subprocess.Popen(
                [str(self.nx_exe_path)],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS
            )

            # The Popen PID might be a launcher, not ugraf.exe itself
            # We need to wait and detect the actual ugraf.exe PID
            if wait_for_ready:
                print(f"[NX SESSION]   Waiting for NX to start (up to {timeout_seconds}s)...")
                start_time = time.time()

                while (time.time() - start_time) < timeout_seconds:
                    current_pids = set(self.get_nx_pids())
                    new_pids = current_pids - pids_before

                    if new_pids:
                        # Found a new NX process!
                        self._started_nx_pid = list(new_pids)[0]
                        elapsed = time.time() - start_time
                        print(f"[NX SESSION]   NX started! PID: {self._started_nx_pid} (took {elapsed:.1f}s)")
                        return self._started_nx_pid

                    # Show progress
                    elapsed = time.time() - start_time
                    if int(elapsed) % 5 == 0 and elapsed > 0:
                        print(f"[NX SESSION]   Still waiting... ({elapsed:.0f}s)")

                    time.sleep(1)

                # Timeout reached
                print(f"[NX SESSION]   WARNING: Timeout waiting for NX to start")
                return None

            else:
                # Don't wait - just return the Popen PID (may not be ugraf.exe)
                self._started_nx_pid = process.pid
                return process.pid

        except Exception as e:
            print(f"[NX SESSION]   ERROR starting NX: {e}")
            raise RuntimeError(f"Failed to start NX: {e}")

    def record_session_start(self, nx_was_running: bool, started_pid: Optional[int] = None):
        """
        Record that we're starting an optimization session.

        Args:
            nx_was_running: Whether NX was already running before we started
            started_pid: PID of NX process we started (if any)
        """
        self._nx_was_running_before_start = nx_was_running

        session_data = {
            'session_id': self.session_id,
            'start_time': time.time(),
            'nx_was_running_before': nx_was_running,
            'started_nx_pid': started_pid,
            'original_nx_pids': self.get_nx_pids() if nx_was_running else []
        }

        with open(self.lock_file, 'w') as f:
            json.dump(session_data, f, indent=2)

        print(f"[NX SESSION] Recorded session start:")
        print(f"  Session ID: {self.session_id}")
        print(f"  NX was already running: {nx_was_running}")
        if started_pid:
            print(f"  Started NX PID: {started_pid}")

    def load_session_data(self) -> Optional[dict]:
        """Load session data from lock file."""
        if self.lock_file.exists():
            try:
                with open(self.lock_file, 'r') as f:
                    return json.load(f)
            except Exception:
                pass
        return None

    def can_close_nx(self) -> bool:
        """
        Determine if we're allowed to close NX.

        Returns True only if:
        1. We have a valid session record
        2. NX was NOT running before we started
        3. We explicitly started NX for this session
        """
        session_data = self.load_session_data()

        if not session_data:
            print("[NX SESSION] WARNING: No session data found - refusing to close NX")
            return False

        if session_data.get('nx_was_running_before', True):
            print("[NX SESSION] NX was running before optimization - NOT closing")
            return False

        if session_data.get('started_nx_pid') is None:
            print("[NX SESSION] We didn't start NX - NOT closing")
            return False

        print("[NX SESSION] We started NX for this session - safe to close")
        return True

    def get_pids_we_can_close(self) -> list:
        """
        Get list of NX PIDs that we're allowed to close.

        Only returns PIDs that:
        1. Are currently running
        2. Were started by us (not in original_nx_pids)
        """
        session_data = self.load_session_data()
        if not session_data:
            return []

        original_pids = set(session_data.get('original_nx_pids', []))
        current_pids = set(self.get_nx_pids())

        # PIDs we can close = current PIDs minus original PIDs
        closeable = list(current_pids - original_pids)

        if closeable:
            print(f"[NX SESSION] PIDs we can close: {closeable}")
        else:
            print(f"[NX SESSION] No PIDs to close (original: {original_pids}, current: {current_pids})")

        return closeable

    def close_nx_if_allowed(self) -> bool:
        """
        Close NX only if we're allowed to.

        Returns:
            True if NX was closed, False if not allowed or already closed
        """
        if not self.can_close_nx():
            return False

        pids_to_close = self.get_pids_we_can_close()

        if not pids_to_close:
            print("[NX SESSION] No NX processes to close")
            return False

        for pid in pids_to_close:
            try:
                print(f"[NX SESSION] Closing NX PID {pid}...")
                subprocess.run(
                    ['taskkill', '/F', '/PID', str(pid)],
                    capture_output=True,
                    timeout=30
                )
            except Exception as e:
                print(f"[NX SESSION] Failed to close PID {pid}: {e}")

        return True

    def cleanup(self):
        """Remove session lock file."""
        if self.lock_file.exists():
            try:
                self.lock_file.unlink()
                print(f"[NX SESSION] Cleaned up session lock: {self.lock_file}")
            except Exception:
                pass


def ensure_nx_running(
    session_id: Optional[str] = None,
    auto_start: bool = True,
    start_timeout: int = 60
) -> Tuple[NXSessionManager, bool]:
    """
    Ensure NX is running for optimization. Starts NX automatically if needed.

    This is the main entry point for optimization scripts. It:
    1. Checks if NX is already running
    2. If not, starts a fresh NX instance
    3. Records the session so we know which instance to close later

    Args:
        session_id: Unique session identifier (defaults to process ID)
        auto_start: If True, automatically start NX if not running
        start_timeout: Seconds to wait for NX to start (default: 60)

    Returns:
        Tuple of (session_manager, nx_was_started)
        - session_manager: Use this to close NX when done
        - nx_was_started: True if we started NX, False if it was already running

    Example:
        manager, we_started_nx = ensure_nx_running()

        # ... run optimization ...

        # This will only close NX if we started it
        manager.close_nx_if_allowed()
        manager.cleanup()
    """
    manager = NXSessionManager(session_id)

    was_running = manager.is_nx_running()
    started_pid = None

    if was_running:
        print("[NX SESSION] NX is already running - will NOT close after optimization")
        # Record existing PIDs so we don't close them
        original_pids = manager.get_nx_pids()
        print(f"[NX SESSION]   Existing PIDs: {original_pids}")
    else:
        if auto_start:
            print("[NX SESSION] NX is not running - starting fresh instance...")
            started_pid = manager.start_nx(wait_for_ready=True, timeout_seconds=start_timeout)

            if started_pid:
                print(f"[NX SESSION] Successfully started NX (PID: {started_pid})")
                print(f"[NX SESSION] Will close this instance when optimization completes")
            else:
                raise RuntimeError(
                    "Failed to start NX. Please ensure NX is installed and try again.\n"
                    f"Checked paths: {', '.join(str(p) for p in DEFAULT_NX_PATHS)}"
                )
        else:
            raise RuntimeError(
                "NX is not running and auto_start=False.\n"
                "Please start NX manually before running optimization."
            )

    manager.record_session_start(
        nx_was_running=was_running,
        started_pid=started_pid
    )

    return manager, not was_running


# Global session manager for module-level usage
_global_session_manager: Optional[NXSessionManager] = None


def get_session_manager() -> Optional[NXSessionManager]:
    """Get the global session manager if one exists."""
    return _global_session_manager


def set_session_manager(manager: NXSessionManager):
    """Set the global session manager."""
    global _global_session_manager
    _global_session_manager = manager


def validate_nx_running(
    require_nx: bool = True,
    timeout_seconds: int = 0,
    auto_prompt: bool = True
) -> bool:
    """
    Validate that NX is running before starting optimization.

    This function checks if NX (ugraf.exe) is running and provides
    clear feedback to the user if it's not.

    Args:
        require_nx: If True, raise an error if NX is not running.
                   If False, just return the status.
        timeout_seconds: If > 0, wait this many seconds for user to open NX.
                        Check every 5 seconds with a countdown.
        auto_prompt: If True, print helpful message when NX is not running.

    Returns:
        True if NX is running, False otherwise.

    Raises:
        RuntimeError: If require_nx=True and NX is not running (after timeout).

    Example:
        # At the start of optimization:
        validate_nx_running(require_nx=True, timeout_seconds=60)
        # Will wait up to 60 seconds for user to open NX
    """
    manager = NXSessionManager("validation_check")

    # Initial check
    if manager.is_nx_running():
        print("[NX VALIDATION] ✓ NX is running")
        return True

    # NX is not running
    if auto_prompt:
        print("\n" + "="*70)
        print("  NX IS NOT RUNNING")
        print("="*70)
        print("\n  Atomizer requires NX to be open to run FEA simulations.")
        print("  Please open NX (Simcenter 3D) and try again.")
        print("\n  Note: You can leave NX open - Atomizer will not close it")
        print("        since it was already running before optimization started.")
        print("="*70 + "\n")

    # If timeout specified, wait for user to open NX
    if timeout_seconds > 0:
        print(f"[NX VALIDATION] Waiting up to {timeout_seconds}s for NX to start...")

        start_time = time.time()
        check_interval = 5  # seconds

        while (time.time() - start_time) < timeout_seconds:
            remaining = int(timeout_seconds - (time.time() - start_time))

            if manager.is_nx_running():
                print(f"\n[NX VALIDATION] ✓ NX detected! Proceeding with optimization...")
                return True

            print(f"  Waiting for NX... ({remaining}s remaining)")
            time.sleep(check_interval)

        print(f"\n[NX VALIDATION] Timeout reached. NX was not started.")

    # NX is still not running
    if require_nx:
        raise RuntimeError(
            "NX is not running. Please open NX (Simcenter 3D) before starting optimization.\n"
            "Atomizer requires NX to execute FEA simulations via journal scripts."
        )

    return False


def require_nx_or_exit(
    auto_start: bool = True,
    start_timeout: int = 60
) -> NXSessionManager:
    """
    Convenience function: Ensure NX is running, starting it if needed.

    This is the recommended function to call at the start of any optimization.
    It will:
    1. Check if NX is already running
    2. If not, automatically start NX (takes ~10-15 seconds)
    3. Track which NX instance we started
    4. Only close that specific instance when done (never user's sessions)

    Args:
        auto_start: If True, start NX automatically if not running (default: True)
        start_timeout: Seconds to wait for NX to start (default: 60)

    Returns:
        NXSessionManager configured with session tracking

    Example:
        # At start of run_optimization.py:
        from optimization_engine.utils import require_nx_or_exit

        session_manager = require_nx_or_exit()
        # ... run optimization ...

        # At end - only closes NX if we started it!
        session_manager.close_nx_if_allowed()
        session_manager.cleanup()
    """
    # Use ensure_nx_running which handles auto-start
    manager, nx_was_started = ensure_nx_running(
        auto_start=auto_start,
        start_timeout=start_timeout
    )

    return manager


if __name__ == '__main__':
    # Test the session manager
    print("Testing NX Session Manager...")

    manager = NXSessionManager("test_session")

    print(f"NX running: {manager.is_nx_running()}")
    print(f"NX PIDs: {manager.get_nx_pids()}")

    # Simulate session start
    manager.record_session_start(
        nx_was_running=manager.is_nx_running(),
        started_pid=None
    )

    print(f"Can close NX: {manager.can_close_nx()}")
    print(f"PIDs we can close: {manager.get_pids_we_can_close()}")

    # Cleanup
    manager.cleanup()
    print("Test complete.")