# NX Journal: Capture Study Images for Atomizer Documentation
#
# Purpose: Capture top view and isometric view images of a part for study documentation
# Usage: run_journal.exe capture_study_images.py -args "part_file_path" "output_directory" ["prefix"]
#
# Arguments:
#   part_file_path: Full path to the .prt file to capture
#   output_directory: Directory where images will be saved
#   prefix (optional): Prefix for image filenames (default: part name)
#
# Output:
#   {prefix}_Top.png - Top view image
#   {prefix}_iso.png - Isometric view image
#
# Author: Atomizer
# Created: 2025-12-18

import sys
import os
import math
import NXOpen
import NXOpen.Gateway


def capture_images(part_path: str, output_dir: str, prefix: str = None):
    """
    Capture top view and isometric view images of a part.

    Args:
        part_path: Full path to the .prt file
        output_dir: Directory to save images
        prefix: Optional prefix for image filenames
    """
    theSession = NXOpen.Session.GetSession()

    # Open the part if not already open
    try:
        workPart, loadStatus = theSession.Parts.OpenDisplay(part_path, NXOpen.Part.LoadStatically)
        loadStatus.Dispose()
    except:
        workPart = theSession.Parts.Work

    if workPart is None:
        print(f"ERROR: Could not open part: {part_path}")
        return False

    # Determine prefix from part name if not provided
    if prefix is None:
        prefix = os.path.splitext(os.path.basename(part_path))[0]

    # Ensure output directory exists
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Hide construction geometry for cleaner images
    _hide_construction_geometry(theSession, workPart)

    # Capture top view
    top_image_path = os.path.join(output_dir, f"{prefix}_Top.png")
    _capture_top_view(theSession, workPart, top_image_path)
    print(f"Saved: {top_image_path}")

    # Capture isometric view
    iso_image_path = os.path.join(output_dir, f"{prefix}_iso.png")
    _capture_isometric_view(theSession, workPart, iso_image_path)
    print(f"Saved: {iso_image_path}")

    return True


def _hide_construction_geometry(theSession, workPart):
    """Hide datums, curves, and sketches for cleaner visualization."""
    markId = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Hide Construction")

    # Hide datums
    theSession.DisplayManager.HideByType("SHOW_HIDE_TYPE_DATUMS",
                                          NXOpen.DisplayManager.ShowHideScope.AnyInAssembly)

    # Hide curves
    theSession.DisplayManager.HideByType("SHOW_HIDE_TYPE_CURVES",
                                          NXOpen.DisplayManager.ShowHideScope.AnyInAssembly)

    # Hide sketches
    theSession.DisplayManager.HideByType("SHOW_HIDE_TYPE_SKETCHES",
                                          NXOpen.DisplayManager.ShowHideScope.AnyInAssembly)

    theSession.UpdateManager.DoUpdate(markId)
    workPart.ModelingViews.WorkView.FitAfterShowOrHide(NXOpen.View.ShowOrHideType.HideOnly)
    theSession.DeleteUndoMark(markId, None)


def _capture_top_view(theSession, workPart, output_path):
    """Capture top view (looking down Z-axis)."""
    # Set top view orientation (looking down -Z)
    matrix = NXOpen.Matrix3x3()
    matrix.Xx = 0.0
    matrix.Xy = -1.0
    matrix.Xz = 0.0
    matrix.Yx = -1.0
    matrix.Yy = 0.0
    matrix.Yz = 0.0
    matrix.Zx = 0.0
    matrix.Zy = 0.0
    matrix.Zz = -1.0
    workPart.ModelingViews.WorkView.Orient(matrix)

    # Fit view
    workPart.ModelingViews.WorkView.Fit()

    # Export image
    _export_image(workPart, output_path)


def _capture_isometric_view(theSession, workPart, output_path):
    """Capture isometric view (standard ISO angle showing backface)."""
    # Set isometric orientation showing backface structure
    rotMatrix = NXOpen.Matrix3x3()
    rotMatrix.Xx = -0.32736574141345925
    rotMatrix.Xy = -0.94489752125198745
    rotMatrix.Xz = -0.00058794613984273266
    rotMatrix.Yx = -0.71924452681462514
    rotMatrix.Yy = 0.24959027079525001
    rotMatrix.Yz = -0.64837643955618585
    rotMatrix.Zx = 0.61279603621108569
    rotMatrix.Zy = -0.21183335680718612
    rotMatrix.Zz = -0.76131967460967154

    # Get current scale and set orientation
    translation = NXOpen.Point3d(0, 0, 0)
    workPart.ModelingViews.WorkView.SetRotationTranslationScale(rotMatrix, translation, 0.25)

    # Fit view
    workPart.ModelingViews.WorkView.Fit()

    # Export image
    _export_image(workPart, output_path)


def _export_image(workPart, output_path, width=1200, height=1000):
    """Export current view as PNG image."""
    imageExportBuilder = workPart.Views.CreateImageExportBuilder()

    try:
        # Configure export settings
        imageExportBuilder.RegionMode = False  # Use entire view
        imageExportBuilder.DeviceWidth = width
        imageExportBuilder.DeviceHeight = height
        imageExportBuilder.FileFormat = NXOpen.Gateway.ImageExportBuilder.FileFormats.Png
        imageExportBuilder.FileName = output_path
        imageExportBuilder.BackgroundOption = NXOpen.Gateway.ImageExportBuilder.BackgroundOptions.Original
        imageExportBuilder.EnhanceEdges = False

        # Commit export
        imageExportBuilder.Commit()
    finally:
        imageExportBuilder.Destroy()


def main(args):
    """Main entry point for journal."""
    if len(args) < 2:
        print("Usage: capture_study_images.py -args \"part_path\" \"output_dir\" [\"prefix\"]")
        print("  part_path: Full path to .prt file")
        print("  output_dir: Directory for output images")
        print("  prefix: Optional filename prefix (default: part name)")
        return

    part_path = args[0]
    output_dir = args[1]
    prefix = args[2] if len(args) > 2 else None

    print(f"Capturing images for: {part_path}")
    print(f"Output directory: {output_dir}")

    success = capture_images(part_path, output_dir, prefix)

    if success:
        print("Image capture complete!")
    else:
        print("Image capture failed!")


if __name__ == '__main__':
    main(sys.argv[1:])