from PySide6.QtWidgets import (
    QRadioButton,
    QButtonGroup,
    QWidget,
    QHBoxLayout,
    QVBoxLayout,
    QFrame,
    QLabel,
)
from nomad_camels.ui_widgets.add_remove_table import AddRemoveTable
from nomad_camels.main_classes.loop_step import Loop_Step, Loop_Step_Config
from nomad_camels.ui_widgets.path_button_edit import Path_Button_Edit
import os
from nomad_camels.utility import variables_handling


class Execute_Python_File(Loop_Step):
    """
    A loop step to execute a python file. The python.exe path as well as the python file path must be given.
    """

    def __init__(self, name="", parent_step=None, step_info=None, **kwargs):
        super().__init__(name, parent_step, step_info, **kwargs)
        self.step_type = "Execute Python File"
        if step_info is None:
            step_info = {}
        # Get the python file path and the python exe path from the step info
        self.file_path = step_info["file_path"] if "file_path" in step_info else ""
        self.python_exe_path = (
            step_info["python_exe_path"] if "python_exe_path" in step_info else ""
        )
        # Get the python packages and versions from the step info
        self.python_packages_versions = (
            step_info["python_packages_versions"]
            if "python_packages_versions" in step_info
            else []
        )
        # Get the radio button selected from the step info
        self.radio_button_selected = (
            step_info["radio_button_selected"]
            if "radio_button_selected" in step_info
            else None
        )

        # Get the python package and version variables and their value from the step info
        self.variables_passing = (
            step_info["variables_passing"] if "variables_passing" in step_info else {}
        )
        # Get the variables returned by the file from the step info
        self.returned_values_variables = (
            step_info["returned_values_variables"]
            if "returned_values_variables" in step_info
            else {}
        )

        self.use_existing_env = (
            step_info["use_existing_env"] if "use_existing_env" in step_info else False
        )
        self.use_specific_packages = (
            step_info["use_specific_packages"]
            if "use_specific_packages" in step_info
            else False
        )
        self.use_camels_python = (
            step_info["use_camels_python"]
            if "use_camels_python" in step_info
            else False
        )

    def get_protocol_short_string(self, n_tabs=0):
        short_string = super().get_protocol_short_string(n_tabs)
        short_string = f"{short_string[:-1]}: {os.path.basename(self.file_path)}"
        return short_string

    def get_protocol_string(self, n_tabs=1):
        """
        This function runs the python file with the python exe using the subprocess module.
        The the cwd of the subprocess is changed to the python files location.
        """
        if not self.file_path:
            raise Exception(f"You need to specify the path of the file for {self.name}")
        tabs = "\t" * n_tabs
        protocol_string = super().get_protocol_string(n_tabs)
        # use subprocess to run the python file
        if self.variables_passing:
            formatted_strings = []
            variable_names = self.variables_passing["Variable Name"]
            values = self.variables_passing["Value"]
            for i in range(len(variable_names)):
                # Format each "name = value" pair
                if values[i] == "":
                    formatted_string = (
                        f"f'{variable_names[i]}={{eva.eval(\"{variable_names[i]}\")}}',"
                    )
                else:
                    formatted_string = f"f'{variable_names[i]}={{{values[i]}}}',"
                # Append to the list
                formatted_strings.append(formatted_string)
            # Step 4: Join all formatted strings with a space
            variables_string = " ".join(formatted_strings)
            variables_string = variables_string[:-1]
        else:
            variables_string = ""
        if self.use_existing_env:
            protocol_string += (
                f"{tabs}result_python_file = subprocess.run(['"
                f'r"{os.path.abspath(self.python_exe_path)}", '
                f'r"{os.path.abspath(self.file_path)}", '
                f"{variables_string}], "
                f'cwd=r"{os.path.abspath(os.path.dirname(self.file_path))}", '
                f"capture_output=True, text=True)\n"
                f"{tabs}helper_functions.evaluate_python_file_output(result_python_file.stdout, namespace)\n"
            )

        if self.use_specific_packages:
            protocol_string += (
                f"{tabs}result_python_file = helper_functions.create_venv_run_file_delete_venv({self.python_packages_versions}, "
                f'r"{os.path.abspath(self.file_path)}", '
                f"{variables_string})\n"
                f"{tabs}helper_functions.evaluate_python_file_output(result_python_file.stdout, namespace)\n"
            )

        if self.use_camels_python:
            import sys

            protocol_string += (
                f"{tabs}result_python_file = subprocess.run(["
                f'r"{os.path.abspath(sys.executable)}", '
                f'r"{os.path.abspath(self.file_path)}", '
                f"{variables_string}], "
                f'cwd=r"{os.path.abspath(os.path.dirname(self.file_path))}", '
                f"capture_output=True, text=True)\n"
                f"{tabs}helper_functions.evaluate_python_file_output(result_python_file.stdout, namespace)\n"
            )
        return protocol_string

    def update_variables(self):
        """ """
        # Update the global variables to have access to the variables defined to pass to the python file
        variable_names = self.variables_passing.get("Variable Name", [])
        values = self.variables_passing.get("Value", [])
        # Create a dictionary from the lists
        new_variables = dict(zip(variable_names, values))
        # remove keys and value if the key is already in variables_handling.channels.keys()
        # remove keys and value if the value is empty so that the value is taken from the running script
        for key in list(new_variables.keys()):
            if (key in variables_handling.channels.keys()) or (
                new_variables[key] == ""
            ):
                new_variables.pop(key)
        # Update the global variables to have access to the variables returned by the python file
        if "Variable Name" in self.returned_values_variables:
            for returned_values_variables in self.returned_values_variables[
                "Variable Name"
            ]:
                if returned_values_variables:
                    if (
                        returned_values_variables
                        not in variables_handling.loop_step_variables
                    ) and (
                        returned_values_variables
                        not in variables_handling.channels.keys()
                    ):
                        variables_handling.loop_step_variables.update(
                            {returned_values_variables: 0}
                        )


class Execute_Python_File_Config(Loop_Step_Config):
    """The configuration settings for the Execute_Python_File loop step."""

    def __init__(
        self,
        loop_step: Execute_Python_File,
        parent=None,
    ):
        super().__init__(parent, loop_step)
        self.loop_step = loop_step

        # Labels setup
        self.label_python_file_path = QLabel("Python File Path")
        self.label_python_exe_path = QLabel("Python Executable Path")
        self.python_packages_versions_label = QLabel("Python Packages and Versions")
        self.passing_variables_label = QLabel("Variables to pass to the Python file")
        self.read_variables_label = QLabel("Values returned by the Python file.")

        # Path_Button_Edit widgets setup
        self.path_button_edit_python_file = Path_Button_Edit()
        self.path_button_edit_python_exe = Path_Button_Edit()

        # Radio buttons setup
        self.radio_existing_env = QRadioButton("Use existing python environment")
        self.radio_specific_packages = QRadioButton("Use specific packages")
        self.radio_camels_python = QRadioButton("Use CAMELS python")
        # Button group to manage exclusive selection
        self.radio_button_group = QButtonGroup()
        self.radio_button_group.addButton(self.radio_existing_env)
        self.radio_button_group.addButton(self.radio_specific_packages)
        self.radio_button_group.addButton(self.radio_camels_python)
        self.radio_button_group.buttonToggled.connect(self.handle_radio_button_clicked)
        # Layout for radio buttons
        radio_layout = QHBoxLayout()
        radio_layout.addWidget(self.radio_camels_python)
        radio_layout.addWidget(self.radio_existing_env)
        radio_layout.addWidget(self.radio_specific_packages)

        # Connect radio button signals
        self.radio_existing_env.toggled.connect(self.toggle_add_remove_table)
        self.radio_specific_packages.toggled.connect(self.toggle_python_exe_path)

        # Horizontal line setup
        h_line = QFrame()
        h_line.setFrameShape(QFrame.HLine)
        h_line.setFrameShadow(QFrame.Sunken)

        # AddRemoveTable widget setup for python packages and their version when selecting the python packages explicitly
        self.add_remove_table_packages = AddRemoveTable(
            headerLabels=["Python Package", "Version"]
        )

        # AddRemoveTable for passing variables to the Python file
        # If no variable is given in the table, it will take the variable values from the running script
        self.add_remove_table_variables_passing = AddRemoveTable(
            headerLabels=["Variable Name", "Value"]
        )

        # AddRemoveTable for reading from the Python file
        # Value is either the index for the returned value if the file returns multiple things or the dictionary key if the file returns a dictionary
        self.add_remove_table_returned_values_variables = (
            AddRemoveTable_Returned_Variables(headerLabels=["Variable Name"])
        )

        # Tooltips
        self.label_python_file_path.setToolTip(
            "Specify the path to the Python file to be executed."
        )
        self.label_python_exe_path.setToolTip(
            "Specify the path to the Python executable (python.exe)."
        )
        self.path_button_edit_python_file.setToolTip(
            "Specify the path to the Python file to be executed."
        )
        self.path_button_edit_python_exe.setToolTip(
            "Specify the path to the Python executable (python.exe)."
        )
        self.radio_existing_env.setToolTip(
            "Run the Python file using the existing Python environment."
        )
        self.radio_specific_packages.setToolTip(
            "Create a virtual environment with specified packages to run the Python file."
        )
        self.radio_camels_python.setToolTip(
            "Run the Python file using the Python environment CAMELS is using."
        )
        self.add_remove_table_packages.setToolTip(
            "Add or remove Python packages and their versions for the virtual environment."
        )
        self.add_remove_table_variables_passing.setToolTip(
            "Specify variables and their values to be passed to the Python file during execution."
        )
        self.add_remove_table_returned_values_variables.setToolTip(
            "The Python file must return a dictionary.\nSpecify the keys of the dictionary that are returned from the Python file after execution."
        )

        # Initial paths setup
        self.path_button_edit_python_file.set_path(self.loop_step.file_path)
        self.path_button_edit_python_exe.set_path(self.loop_step.python_exe_path)
        self.path_button_edit_python_file.path_changed.connect(self.update_file_path)
        self.path_button_edit_python_exe.path_changed.connect(self.update_exe_path)

        # Save changes of the python packages table to the loop step
        self.add_remove_table_packages.table_model.itemChanged.connect(
            self.update_python_packages
        )
        self.add_remove_table_packages.change_table_data(
            self.loop_step.python_packages_versions
        )
        self.add_remove_table_packages.table_model.rowsRemoved.connect(
            self.on_rows_removed_packages
        )
        # Save changes of the variables passing table to the loop step
        self.add_remove_table_variables_passing.table_model.itemChanged.connect(
            self.update_variables_passing
        )
        self.add_remove_table_variables_passing.change_table_data(
            self.loop_step.variables_passing
        )
        self.add_remove_table_variables_passing.table_model.rowsRemoved.connect(
            self.on_rows_removed_passing
        )
        # Save changes of the returned variables table to the loop step
        self.add_remove_table_returned_values_variables.table_model.itemChanged.connect(
            self.update_returned_value_variables
        )
        self.add_remove_table_returned_values_variables.change_table_data(
            self.loop_step.returned_values_variables
        )
        self.add_remove_table_returned_values_variables.table_model.rowsRemoved.connect(
            self.on_rows_removed_returned
        )

        # Layout setup
        layout = self.layout()

        layout.addWidget(self.label_python_file_path, 1, 0)
        layout.addWidget(self.path_button_edit_python_file, 1, 1)

        # Create a widget to hold radio buttons and add it to the main layout
        radio_widget = QWidget()
        radio_widget.setLayout(radio_layout)
        layout.addWidget(radio_widget, 2, 0, 1, 2)

        layout.addWidget(self.label_python_exe_path, 3, 0)
        layout.addWidget(self.path_button_edit_python_exe, 3, 1)
        layout.addWidget(self.python_packages_versions_label, 4, 0, 1, 2)
        layout.addWidget(self.add_remove_table_packages, 5, 0, 1, 2)
        layout.addWidget(h_line, 6, 0, 1, 2)
        layout.addWidget(self.passing_variables_label, 7, 0, 1, 2)
        layout.addWidget(self.add_remove_table_variables_passing, 8, 0, 1, 2)
        layout.addWidget(self.read_variables_label, 9, 0, 1, 2)
        layout.addWidget(self.add_remove_table_returned_values_variables, 10, 0, 1, 2)

        # Set the selected radio button from the loop step
        if self.loop_step.radio_button_selected is not None:
            if (
                self.loop_step.radio_button_selected
                == "Use existing python environment"
            ):
                self.radio_existing_env.setChecked(True)
            elif self.loop_step.radio_button_selected == "Use specific packages":
                self.radio_specific_packages.setChecked(True)
            elif self.loop_step.radio_button_selected == "Use CAMELS python":
                self.radio_camels_python.setChecked(True)
        else:
            self.radio_camels_python.setChecked(True)

    def update_exe_path(self, path):
        self.loop_step.python_exe_path = path

    def update_file_path(self, path):
        self.loop_step.file_path = path

    def update_python_packages(self):
        self.loop_step.python_packages_versions = (
            self.add_remove_table_packages.update_table_data()
        )

    def update_variables_passing(self):
        self.loop_step.variables_passing = (
            self.add_remove_table_variables_passing.update_table_data()
        )

    def update_returned_value_variables(self):

        self.loop_step.returned_values_variables = (
            self.add_remove_table_returned_values_variables.update_table_data()
        )

    def handle_radio_button_clicked(self, button):
        self.loop_step.radio_button_selected = button.text()
        if button == self.radio_existing_env:
            # Handle selection for existing environment
            self.add_remove_table_packages.setVisible(False)
            self.path_button_edit_python_exe.setVisible(True)
            self.label_python_exe_path.setVisible(True)
            self.python_packages_versions_label.setVisible(False)
            self.loop_step.use_existing_env = True
            self.loop_step.use_specific_packages = False
            self.loop_step.use_camels_python = False
        elif button == self.radio_specific_packages:
            # Handle selection for specific packages
            self.add_remove_table_packages.setVisible(True)
            self.path_button_edit_python_exe.setVisible(False)
            self.label_python_exe_path.setVisible(False)
            self.python_packages_versions_label.setVisible(True)
            self.loop_step.use_specific_packages = True
            self.loop_step.use_existing_env = False
            self.loop_step.use_camels_python = False
        elif button == self.radio_camels_python:
            # Handle selection for CAMELS python
            # Hide AddRemoveTable and Python exe path
            self.add_remove_table_packages.setVisible(False)
            self.path_button_edit_python_exe.setVisible(False)
            self.label_python_exe_path.setVisible(False)
            self.python_packages_versions_label.setVisible(False)
            self.loop_step.use_camels_python = True
            self.loop_step.use_specific_packages = False
            self.loop_step.use_existing_env = False

    def toggle_add_remove_table(self, checked):
        if checked:
            self.add_remove_table_packages.setVisible(False)

    def toggle_python_exe_path(self, checked):
        if checked:
            self.path_button_edit_python_exe.setVisible(not checked)
            self.label_python_exe_path.setVisible(not checked)

    def on_rows_removed_returned(self, parent, first, last):
        self.loop_step.returned_values_variables = (
            self.add_remove_table_returned_values_variables.update_table_data()
        )

    def on_rows_removed_passing(self, parent, first, last):
        self.loop_step.variables_passing = (
            self.add_remove_table_variables_passing.update_table_data()
        )

    def on_rows_removed_packages(self, parent, first, last):
        self.loop_step.python_packages_versions = (
            self.add_remove_table_packages.update_table_data()
        )


class AddRemoveTable_Returned_Variables(AddRemoveTable):
    """Subclasses AddRemoveTable and changes the remove method to also remove the variable from the loop step variable dict."""

    def remove(self):
        try:
            index = self.table.selectedIndexes()[0]
        except IndexError:
            raise Exception("You need to select a row first!")
        row = index.row()
        variable_name_to_be_removed = self.tableData["Variable Name"][row]
        variables_handling.loop_step_variables.pop(variable_name_to_be_removed)
        super().remove()
