import maya.cmds as cmds

if cmds.about(version=True) == "2025":
    from PySide6 import QtWidgets, QtGui
    import shiboken6 as shiboken
else:
    from PySide2 import QtWidgets, QtGui
    import shiboken2 as shiboken

import maya.OpenMayaUI as omui
from enum import Enum
import random
import os
import pathlib
from datetime import datetime
import time


def get_scaling_ratio():
    screen = QtGui.QGuiApplication.primaryScreen()
    return screen.logicalDotsPerInch() / 96.0

scaling_ratio = get_scaling_ratio()

def maya_main_window():
    main_window_pointer = omui.MQtUtil.mainWindow()
    return shiboken.wrapInstance(int(main_window_pointer), QtWidgets.QWidget)

class OptionType(Enum):
    Unknown = 1
    Bool = 2
    Int = 3
    Float = 4
    Option = 5
    File = 6

class Option:
    
    option_type = OptionType.Unknown
    option_name = None
    default = None
    setting = None
    widget = None

    def __init__(self, option_name, option_type, default, setting, meta=None):
        self.option_name = option_name
        self.option_type = option_type
        self.default = default
        self.setting = setting
        self.meta = meta

    def make_widget(self, form_layout):
        if self.option_type == OptionType.Bool:
            self.widget = QtWidgets.QCheckBox()
            self.widget.setChecked(self.default)
            form_layout.addRow(self.option_name + ":", self.widget)
        elif self.option_type == OptionType.Int:
            self.widget = QtWidgets.QSpinBox()
            self.widget.setFixedWidth(70 * scaling_ratio)
            self.widget.setButtonSymbols(QtWidgets.QAbstractSpinBox.NoButtons)
            if self.meta:
                if "min" in self.meta:
                    self.widget.setMinimum(self.meta["min"])
                if "max" in self.meta:
                    self.widget.setMaximum(self.meta["max"])
            self.widget.setValue(self.default)
            form_layout.addRow(self.option_name + ":", self.widget)
        elif self.option_type == OptionType.Float:
            self.widget = QtWidgets.QDoubleSpinBox()
            self.widget.setDecimals(4)
            self.widget.setSingleStep(0.01)
            self.widget.setFixedWidth(70 * scaling_ratio)
            self.widget.setButtonSymbols(QtWidgets.QAbstractSpinBox.NoButtons)
            if self.meta:
                if "min" in self.meta:
                    self.widget.setMinimum(self.meta["min"])
                if "max" in self.meta:
                    self.widget.setMaximum(self.meta["max"])
            self.widget.setValue(self.default)
            form_layout.addRow(self.option_name + ":", self.widget)
        elif self.option_type == OptionType.Option:
            self.widget = QtWidgets.QComboBox()
            if self.meta:
                if "options" in self.meta:
                    for option in self.meta["options"]:
                        self.widget.addItem(option[1], option[0])
            form_layout.addRow(self.option_name + ":", self.widget)
            index = self.widget.findData(self.default)
            self.widget.setCurrentIndex(index)
        elif self.option_type == OptionType.File:
            self.widget = QtWidgets.QLineEdit()
            widget_layout = QtWidgets.QHBoxLayout()
            widget_layout.setContentsMargins(0, 0, 0, 0)
            browse_button = QtWidgets.QPushButton("Browse")

            def browse_action():
                path = self.widget.text()
                file_filter = ""
                if self.meta:
                    if "fileFilter" in self.meta:
                        file_filter = self.meta["fileFilter"]
                result = cmds.fileDialog2(dir=path, fileFilter=file_filter, dialogStyle=2)
                if result:
                    self.widget.setText(result[0])

            browse_button.clicked.connect(browse_action)

            widget_layout.addWidget(self.widget)
            widget_layout.addWidget(browse_button)
            form_layout.addRow(self.option_name + ":", widget_layout)
            self.widget.setText(self.default)

    def load(self):
        if not self.widget:
            return
        if self.option_type == OptionType.Bool:
            if cmds.optionVar(exists=self.setting):
                value = bool(cmds.optionVar(q=self.setting))
                self.widget.setChecked(value)
        elif self.option_type == OptionType.Int:
            if cmds.optionVar(exists=self.setting):
                value = cmds.optionVar(q=self.setting)
                self.widget.setValue(value)
        elif self.option_type == OptionType.Float:
            if cmds.optionVar(exists=self.setting):
                value = cmds.optionVar(q=self.setting)
                self.widget.setValue(value)
        elif self.option_type == OptionType.Option:
            if cmds.optionVar(exists=self.setting):
                value = cmds.optionVar(q=self.setting)
                index = self.widget.findData(value)
                self.widget.setCurrentIndex(index)
        elif self.option_type == OptionType.File:
            if cmds.optionVar(exists=self.setting):
                value = cmds.optionVar(q=self.setting)
                self.widget.setText(value)

    def save(self):
        if not self.widget:
            return
        if self.option_type == OptionType.Bool:
            if self.widget.isChecked():
                cmds.optionVar(iv=(self.setting, 1))
            else:
                cmds.optionVar(iv=(self.setting, 0))
        elif self.option_type == OptionType.Int:
            cmds.optionVar(iv=(self.setting, self.widget.value()))
        elif self.option_type == OptionType.Float:
            cmds.optionVar(fv=(self.setting, self.widget.value()))
        elif self.option_type == OptionType.Option:
            cmds.optionVar(iv=(self.setting, self.widget.currentData()))
        elif self.option_type == OptionType.File:
            cmds.optionVar(sv=(self.setting, self.widget.text()))

    def reset(self):
        if not self.widget:
            return
        if self.option_type == OptionType.Bool:
            self.widget.setChecked(self.default)
            if self.default:
                cmds.optionVar(iv=(self.setting, 1))
            else:
                cmds.optionVar(iv=(self.setting, 0))
        elif self.option_type == OptionType.Int:
            self.widget.setValue(self.default)
            cmds.optionVar(iv=(self.setting, self.default))
        elif self.option_type == OptionType.Float:
            self.widget.setValue(self.default)
            cmds.optionVar(fv=(self.setting, self.default))
        elif self.option_type == OptionType.Option:
            index = self.widget.findData(self.default)
            self.widget.setCurrentIndex(index)
            cmds.optionVar(iv=(self.setting, self.default))
        elif self.option_type == OptionType.File:
            self.widget.setText(self.default)
            cmds.optionVar(sv=(self.setting, self.default))

class OptionsDialogType(Enum):
    Create = 1
    Export = 2

class OptionsDialog(QtWidgets.QDialog):
    
    def __init__(self, options_dialog_type, title, command, options, parent=None):
        super(OptionsDialog, self).__init__(parent)
        
        self.setWindowTitle(title)

        self.options_dialog_type = options_dialog_type
        self.command = command
        self.options = options

        layout = QtWidgets.QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(layout)

        menu_bar = QtWidgets.QMenuBar()
        edit_menu = menu_bar.addMenu("Edit")
        edit_menu.addAction("Save Settings", self.save)
        edit_menu.addAction("Reset Settings", self.reset)

        if self.options_dialog_type == OptionsDialogType.Export:
            edit_menu.addSeparator()
            edit_menu.addAction("Clear Log", self.clear_log)

        layout.addWidget(menu_bar)

        form_layout = QtWidgets.QFormLayout()
        form_layout_margins = 8 * scaling_ratio
        form_layout.setContentsMargins(
            form_layout_margins, form_layout_margins, form_layout_margins, form_layout_margins)

        for option in self.options:
            option.make_widget(form_layout)

        layout.addLayout(form_layout)

        if options_dialog_type == OptionsDialogType.Create:
            layout.addStretch()

        buttons_layout = QtWidgets.QHBoxLayout()       
        layout.addLayout(buttons_layout)
        buttons_layout_margins = 4 * scaling_ratio
        buttons_layout.setContentsMargins(
            buttons_layout_margins, buttons_layout_margins, buttons_layout_margins, buttons_layout_margins)
        buttons_layout.setSpacing(4)

        if options_dialog_type == OptionsDialogType.Create:
            create_button = QtWidgets.QPushButton("Create")
            apply_button = QtWidgets.QPushButton("Apply")
            close_button = QtWidgets.QPushButton("Close")

            create_button.clicked.connect(self.create)
            apply_button.clicked.connect(self.apply)
            close_button.clicked.connect(self.close)

            buttons_layout.addWidget(create_button)
            buttons_layout.addWidget(apply_button)
            buttons_layout.addWidget(close_button)
        elif options_dialog_type == OptionsDialogType.Export:
            self.log = QtWidgets.QTextEdit()
            #self.log.setFixedHeight(80)
            self.log.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.MinimumExpanding)

            self.log.setReadOnly(True)
            form_layout.addRow("Log:", self.log)

            export_button = QtWidgets.QPushButton("Export")
            close_button = QtWidgets.QPushButton("Close")

            export_button.clicked.connect(self.apply)
            close_button.clicked.connect(self.close)

            buttons_layout.addStretch()
            buttons_layout.addWidget(export_button)
            buttons_layout.addWidget(close_button)

        self.load()
        self.show()
    
    def load(self):
        for option in self.options:
            option.load()

    def save(self):
        for option in self.options:
            option.save()
    
    def reset(self):
        for option in self.options:
            option.reset()

    def create(self):
        self.save()
        self.command()
        self.close()

    def apply(self):
        self.save()
        if self.options_dialog_type == OptionsDialogType.Export:
            start_time = time.time()
            start_date = datetime.now()
            self.log.moveCursor(QtGui.QTextCursor.End)
            self.log.append("<font color=grey>Export started</font> " + str(start_date))
            try:
                self.command()
                end_time = time.time()
                end_date = datetime.now()
                self.log.moveCursor(QtGui.QTextCursor.End)
                self.log.append("<font color=grey>Export ended</font> " + str(end_date))
                self.log.append("<font color=grey>Export took</font> %.4f <font color=grey>seconds</font>" % (end_time - start_time))
                self.log.append("<font color=green>Done</font>")
            except Exception as e:
                self.log.moveCursor(QtGui.QTextCursor.End)
                self.log.append("<font color=orange>%s</font>" % e)
                self.log.append("<font color=red>Failure</font>")
        else:
            self.command()
    
    def activate(self):
        self.load()
        self.show()
        self.activateWindow()
    
    def clear_log(self):
        self.log.clear()

def get_value(setting):
    if cmds.optionVar(exists=setting[1]):
        return cmds.optionVar(q=setting[1])
    return setting[0]

create_ahoge_window = None

create_ahoge_randomize_seed = [True, "ahoge_randomize_seed"]
create_ahoge_swap_uv = [True, "ahoge_randomize_swap_uv"]
create_ahoge_segments = [16, "ahoge_segments"]
create_ahoge_number_of_curves = [10, "ahoge_number_of_curves"]
create_ahoge_width = [0.01, "ahoge_width"]

def create_ahoge():
    cmds.undoInfo(chunkName="create_ahoge", openChunk=True)
    try:
        randomize_seed = get_value(create_ahoge_randomize_seed)
        swap_uv = get_value(create_ahoge_swap_uv)
        segments = get_value(create_ahoge_segments)
        number_of_curves = get_value(create_ahoge_number_of_curves)
        width = get_value(create_ahoge_width)

        surfaces = cmds.ls(sl=True, dag=True, type="nurbsSurface")
        ahoge_shape = cmds.createNode("ahogeShape")
        
        if randomize_seed:
            cmds.setAttr(ahoge_shape + ".seed", random.randint(0, 10000))

        cmds.setAttr(ahoge_shape + ".numSegments", segments)
        cmds.setAttr(ahoge_shape + ".numCurves", number_of_curves)
        cmds.setAttr(ahoge_shape + ".width", width)
        
        for index, surface in enumerate(surfaces):
            if swap_uv and cmds.getAttr(surface + ".fu") == 0 and cmds.getAttr(surface + ".fv") == 2:
                cmds.warning(surface + " swapping UV")
                cmds.reverseSurface(surface, d=3, ch=1, rpo=1)
                cmds.reverseSurface(surface, d=0, ch=1, rpo=1)
            cmds.connectAttr(surface + ".worldSpace[0]", ahoge_shape + ".inputSurfaces[" + str(index) + "]")
    finally:
        cmds.undoInfo(closeChunk=True)

def create_ahoge_options():
    global create_ahoge_window
    if not create_ahoge_window:
        title = "Create Ahoge Options"
        options = []
        options.append(Option("Randomize Seed", OptionType.Bool, *create_ahoge_randomize_seed))
        options.append(Option("Swap UV", OptionType.Bool, *create_ahoge_swap_uv))
        options.append(Option("Segments", OptionType.Int, *create_ahoge_segments, {"min": 1, "max": 1000000}))
        options.append(Option("Number of Curves", OptionType.Int, *create_ahoge_number_of_curves, {"min": 1, "max": 1000000}))
        options.append(Option("Width", OptionType.Float, *create_ahoge_width, {"min": 0.0001, "max": 1000000.0}))

        create_ahoge_window = OptionsDialog(OptionsDialogType.Create, title, create_ahoge, options, maya_main_window())
    else:
        create_ahoge_window.activate()

create_nurbopus_window = None

create_nurbopus_method = [1, "ahoge_nurbopus_method"]
create_nurbopus_shape = [0, "ahoge_nurbopus_shape"]
create_nurbopus_useg = [4, "ahoge_nurbopus_useg"]
create_nurbopus_vseg = [8, "ahoge_nurbopus_vseg"]

def create_nurbopus():
    cmds.undoInfo(chunkName="create_nurbopus", openChunk=True)
    try:
        method = get_value(create_nurbopus_method)
        shape = get_value(create_nurbopus_shape)
        useg = get_value(create_nurbopus_useg)
        vseg = get_value(create_nurbopus_vseg)

        curves = cmds.ls(sl=True, dag=True, type="nurbsCurve")
        if not curves:
            cmds.error("No curves selected")
        meshes = cmds.ls(sl=True, dag=True, type="mesh")

        nurbopus = cmds.createNode("nurbopusNode")

        cmds.setAttr(nurbopus + ".method", method)
        cmds.setAttr(nurbopus + ".shape", shape)
        cmds.setAttr(nurbopus + ".uSegments", useg)
        cmds.setAttr(nurbopus + ".vSegments", vseg)

        for index, curve in enumerate(curves):
            cmds.connectAttr(curve + ".worldSpace[0]", nurbopus + ".inputCurves[" + str(index) + "]")
            surface = cmds.createNode("nurbsSurface")

            cmds.connectAttr(nurbopus + ".outputSurfaces[" + str(index) + "]", surface + ".create")
            cmds.hyperShade(surface, assign="initialShadingGroup")

        if meshes:
            cmds.connectAttr(meshes[0] + ".worldMesh[0]", nurbopus + ".inputMesh")
        
        cmds.select(nurbopus)

    finally:
        cmds.undoInfo(closeChunk=True)

def create_nurbopus_options():
    global create_nurbopus_window
    if not create_nurbopus_window:
        title = "Create Nurbopus Options"
        options = []
        options.append(Option("Method", OptionType.Option, *create_nurbopus_method, {"options": [[0, "Up Vector"], [1, "Untwist"]]}))
        options.append(Option("Shape", OptionType.Option, *create_nurbopus_shape, {"options": [[0, "circle"], [1, "line"]]}))
        options.append(Option("U Segments", OptionType.Int, *create_nurbopus_useg, {"min": 1, "max": 1000000}))
        options.append(Option("V Segments", OptionType.Int, *create_nurbopus_vseg, {"min": 1, "max": 1000000}))

        create_nurbopus_window = OptionsDialog(OptionsDialogType.Create, title, create_nurbopus, options, maya_main_window())
    else:
        create_nurbopus_window.activate()

export_alembic_window = None

abc_path = pathlib.Path(os.environ["MAYA_APP_DIR"])
abc_path = abc_path / "projects" / "default" / "data" / "test.abc"
export_alembic_path = [str(abc_path), "ahoge_export_alembic_path"]

export_alembic_mode = [0, "ahoge_export_alembic_mode"]
export_alembic_time_range = [0, "ahoge_export_alembic_time_range"]
export_alembic_start = [1, "ahoge_export_alembic_start"]
export_alembic_end = [24, "ahoge_export_alembic_end"]
export_alembic_relative_sample = [False, "ahoge_export_alembic_relative_sample"]
export_alembic_low = [-0.2, "ahoge_export_alembic_low"]
export_alembic_high = [0.2, "ahoge_export_alembic_high"]

def export_alembic():
    alembic_mode = get_value(export_alembic_mode)
    alembic_time_range = get_value(export_alembic_time_range)
    alembic_path = get_value(export_alembic_path)

    start = get_value(export_alembic_start)
    end = get_value(export_alembic_end)

    relative_sample = get_value(export_alembic_relative_sample)
    low = get_value(export_alembic_low)
    high = get_value(export_alembic_high)

    export_mode="default"
    if alembic_mode == 1:
        export_mode = "unreal"
    elif alembic_mode == 2:
        export_mode = "marmoset"

    extra = ""

    time_range = "current_frame"
    if alembic_time_range == 1:
        time_range = "time_slider"
    elif alembic_time_range == 2:
        time_range = "start_end"
        extra += " -startEnd %d %d" % (start, end)

    if relative_sample:
        extra += " -relativeSample true -lowHigh %f %f" % (low, high)

    print("ahogeCmd -e \"%s\" -exportMode \"%s\" -timeRange \"%s\"%s;" % 
          (alembic_path.replace("\\", "/"), export_mode, time_range, extra))
    cmds.ahogeCmd(e=alembic_path, 
                  exportMode=export_mode, 
                  timeRange=time_range, 
                  startEnd=(start, end), 
                  relativeSample=relative_sample, 
                  lowHigh=(low, high))

def export_alembic_options():
    global export_alembic_window
    if not export_alembic_window:
        title = "Ahoge Export: Alembic"
        options = []
        options.append(Option("Path", OptionType.File, *export_alembic_path, {"fileFilter": "*.abc"}))
        options.append(Option("Mode", OptionType.Option, *export_alembic_mode, {"options": [[0, "Default"], [1, "Unreal"], [2, "Marmoset"]]}))
        options.append(Option("Time Range", OptionType.Option, *export_alembic_time_range, {"options": [[0, "Current Frame"], [1, "Time Slider"], [2, "Start/End"]]}))
        options.append(Option("Start Frame", OptionType.Int, *export_alembic_start, {"min": -1000000, "max": 1000000}))
        options.append(Option("End Frame", OptionType.Int, *export_alembic_end, {"min": -1000000, "max": 1000000}))
        options.append(Option("Relative Sample", OptionType.Bool, *export_alembic_relative_sample))
        options.append(Option("Low", OptionType.Float, *export_alembic_low, {"min": -1.0, "max": 0.0}))
        options.append(Option("High", OptionType.Float, *export_alembic_high, {"min": 0.0, "max": 1.0}))

        export_alembic_window = OptionsDialog(OptionsDialogType.Export, title, export_alembic, options, maya_main_window())
    else:
        export_alembic_window.activate()