"""
NX Part Manager Hook
====================

Provides Python functions to open, close, and save NX parts.

API Reference (verified via Siemens MCP docs):
- Session.Parts() -> PartCollection
- PartCollection.OpenBase() -> Opens a part file
- Part.Close() -> Closes the part
- Part.Save() -> Saves the part
- Part.SaveAs() -> Saves the part with a new name

Usage:
    from optimization_engine.hooks.nx_cad import part_manager

    # Open a part
    part = part_manager.open_part("C:/path/to/part.prt")

    # Save the part
    part_manager.save_part(part)

    # Close the part
    part_manager.close_part(part)
"""

import os
import json
import subprocess
import tempfile
from pathlib import Path
from typing import Optional, Dict, Any, Tuple

# Import NX path from centralized config
try:
    from config import NX_BIN_DIR
    NX_BIN_PATH = str(NX_BIN_DIR)
except ImportError:
    NX_BIN_PATH = os.environ.get(
        "NX_BIN_PATH",
        r"C:\Program Files\Siemens\DesigncenterNX2512\NXBIN"
    )

# Journal template for part operations
PART_OPERATIONS_JOURNAL = '''
# NX Open Python Journal - Part Operations
# Auto-generated by Atomizer hooks

import NXOpen
import NXOpen.UF
import json
import sys
import os

def main():
    """Execute part operation based on command arguments."""
    # Get the NX session
    session = NXOpen.Session.GetSession()

    # Parse arguments: operation, part_path, [output_json]
    args = sys.argv[1:] if len(sys.argv) > 1 else []

    if len(args) < 2:
        raise ValueError("Usage: script.py <operation> <part_path> [output_json]")

    operation = args[0]
    part_path = args[1]
    output_json = args[2] if len(args) > 2 else None

    result = {"success": False, "error": None, "data": {}}

    try:
        if operation == "open":
            result = open_part(session, part_path)
        elif operation == "close":
            result = close_part(session, part_path)
        elif operation == "save":
            result = save_part(session, part_path)
        elif operation == "save_as":
            new_path = args[3] if len(args) > 3 else None
            result = save_part_as(session, part_path, new_path)
        elif operation == "info":
            result = get_part_info(session, part_path)
        else:
            result["error"] = f"Unknown operation: {operation}"
    except Exception as e:
        result["error"] = str(e)

    # Write result to output JSON if specified
    if output_json:
        with open(output_json, 'w') as f:
            json.dump(result, f, indent=2)

    return result


def open_part(session, part_path):
    """Open a part file.

    NX Open API: Session.Parts().OpenActiveDisplay()
    """
    result = {"success": False, "error": None, "data": {}}

    if not os.path.exists(part_path):
        result["error"] = f"Part file not found: {part_path}"
        return result

    try:
        # Set load options for the working directory
        working_dir = os.path.dirname(part_path)
        session.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory
        session.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])

        # Open the part using OpenActiveDisplay (more compatible with batch mode)
        part, load_status = session.Parts.OpenActiveDisplay(
            part_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        load_status.Dispose()

        if part is None:
            result["error"] = "Failed to open part - returned None"
            return result

        result["success"] = True
        result["data"] = {
            "part_name": part.Name,
            "full_path": part.FullPath,
            "leaf": part.Leaf,
            "is_modified": part.IsModified,
            "is_fully_loaded": part.IsFullyLoaded,
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def close_part(session, part_path):
    """Close a part.

    NX Open API: Part.Close()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        # Find the part in the session
        part = find_part_by_path(session, part_path)

        if part is None:
            result["error"] = f"Part not found in session: {part_path}"
            return result

        # Close the part
        # Parameters: close_whole_tree, close_modified, responses
        part.Close(
            NXOpen.BasePart.CloseWholeTree.TrueValue,
            NXOpen.BasePart.CloseModified.CloseModified,
            None
        )

        result["success"] = True
        result["data"] = {"closed": part_path}

    except Exception as e:
        result["error"] = str(e)

    return result


def save_part(session, part_path):
    """Save a part.

    NX Open API: Part.Save()
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        # Find the part in the session
        part = find_part_by_path(session, part_path)

        if part is None:
            result["error"] = f"Part not found in session: {part_path}"
            return result

        # Save the part
        # Parameters: save_component_parts, close_after_save
        save_status = part.Save(
            NXOpen.BasePart.SaveComponents.TrueValue,
            NXOpen.BasePart.CloseAfterSave.FalseValue
        )

        result["success"] = True
        result["data"] = {
            "saved": part_path,
            "is_modified": part.IsModified
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def save_part_as(session, part_path, new_path):
    """Save a part with a new name.

    NX Open API: Part.SaveAs()
    """
    result = {"success": False, "error": None, "data": {}}

    if not new_path:
        result["error"] = "New path is required for SaveAs operation"
        return result

    try:
        # Find the part in the session
        part = find_part_by_path(session, part_path)

        if part is None:
            result["error"] = f"Part not found in session: {part_path}"
            return result

        # Save as new file
        part.SaveAs(new_path)

        result["success"] = True
        result["data"] = {
            "original": part_path,
            "saved_as": new_path
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def get_part_info(session, part_path):
    """Get information about a part.

    NX Open API: Part properties
    """
    result = {"success": False, "error": None, "data": {}}

    try:
        # Find the part in the session
        part = find_part_by_path(session, part_path)

        if part is None:
            result["error"] = f"Part not found in session: {part_path}"
            return result

        # Get part info
        result["success"] = True
        result["data"] = {
            "name": part.Name,
            "full_path": part.FullPath,
            "leaf": part.Leaf,
            "is_modified": part.IsModified,
            "is_fully_loaded": part.IsFullyLoaded,
            "is_read_only": part.IsReadOnly,
            "has_write_access": part.HasWriteAccess,
            "part_units": str(part.PartUnits),
        }

    except Exception as e:
        result["error"] = str(e)

    return result


def find_part_by_path(session, part_path):
    """Find a part in the session by its file path."""
    part_path_normalized = os.path.normpath(part_path).lower()

    for part in session.Parts:
        if os.path.normpath(part.FullPath).lower() == part_path_normalized:
            return part

    return None


if __name__ == "__main__":
    main()
'''


def _get_run_journal_exe() -> str:
    """Get the path to run_journal.exe."""
    return os.path.join(NX_BIN_PATH, "run_journal.exe")


def _run_journal(journal_path: str, args: list) -> Tuple[bool, str]:
    """Run an NX journal with arguments.

    Returns:
        Tuple of (success, output_or_error)
    """
    run_journal = _get_run_journal_exe()

    if not os.path.exists(run_journal):
        return False, f"run_journal.exe not found at {run_journal}"

    cmd = [run_journal, journal_path, "-args"] + args

    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=120  # 2 minute timeout
        )

        if result.returncode != 0:
            return False, f"Journal execution failed: {result.stderr}"

        return True, result.stdout

    except subprocess.TimeoutExpired:
        return False, "Journal execution timed out"
    except Exception as e:
        return False, str(e)


def _execute_part_operation(
    operation: str,
    part_path: str,
    extra_args: list = None
) -> Dict[str, Any]:
    """Execute a part operation via NX journal.

    Args:
        operation: The operation to perform (open, close, save, save_as, info)
        part_path: Path to the part file
        extra_args: Additional arguments for the operation

    Returns:
        Dict with operation result
    """
    # Create temporary journal file
    with tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.py',
        delete=False
    ) as journal_file:
        journal_file.write(PART_OPERATIONS_JOURNAL)
        journal_path = journal_file.name

    # Create temporary output file
    output_file = tempfile.NamedTemporaryFile(
        mode='w',
        suffix='.json',
        delete=False
    ).name

    try:
        # Build arguments
        args = [operation, part_path, output_file]
        if extra_args:
            args.extend(extra_args)

        # Run the journal
        success, output = _run_journal(journal_path, args)

        if not success:
            return {"success": False, "error": output, "data": {}}

        # Read the result
        if os.path.exists(output_file):
            with open(output_file, 'r') as f:
                return json.load(f)
        else:
            return {"success": False, "error": "Output file not created", "data": {}}

    finally:
        # Cleanup temporary files
        if os.path.exists(journal_path):
            os.unlink(journal_path)
        if os.path.exists(output_file):
            os.unlink(output_file)


# =============================================================================
# Public API
# =============================================================================

def open_part(part_path: str) -> Dict[str, Any]:
    """Open an NX part file.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with part_name, full_path, leaf, is_modified, is_fully_loaded

    Example:
        >>> result = open_part("C:/models/bracket.prt")
        >>> if result["success"]:
        ...     print(f"Opened: {result['data']['part_name']}")
    """
    part_path = os.path.abspath(part_path)

    if not os.path.exists(part_path):
        return {
            "success": False,
            "error": f"Part file not found: {part_path}",
            "data": {}
        }

    return _execute_part_operation("open", part_path)


def close_part(part_path: str) -> Dict[str, Any]:
    """Close an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with closed path
    """
    part_path = os.path.abspath(part_path)
    return _execute_part_operation("close", part_path)


def save_part(part_path: str) -> Dict[str, Any]:
    """Save an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with saved path and is_modified flag
    """
    part_path = os.path.abspath(part_path)
    return _execute_part_operation("save", part_path)


def save_part_as(part_path: str, new_path: str) -> Dict[str, Any]:
    """Save an NX part with a new name.

    Args:
        part_path: Full path to the original .prt file
        new_path: Full path for the new file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with original and saved_as paths
    """
    part_path = os.path.abspath(part_path)
    new_path = os.path.abspath(new_path)
    return _execute_part_operation("save_as", part_path, [new_path])


def get_part_info(part_path: str) -> Dict[str, Any]:
    """Get information about an NX part.

    Args:
        part_path: Full path to the .prt file

    Returns:
        Dict with keys:
            - success: bool
            - error: Optional error message
            - data: Dict with name, full_path, leaf, is_modified,
                    is_fully_loaded, is_read_only, has_write_access, part_units
    """
    part_path = os.path.abspath(part_path)
    return _execute_part_operation("info", part_path)