from chimerax.core.tools import ToolInstance
from chimerax.core.commands import run
from chimerax.ui import MainToolWindow
from Qt.QtWidgets import QVBoxLayout, QLabel, QComboBox, QPushButton, QWidget


class DuplicateStructureTool(ToolInstance):
    SESSION_ENDURING = False
    SESSION_SAVE = True
    help = "help:user/tools/duplicatestructure.html"  # Replace with the appropriate help URL if available

    def __init__(self, session, tool_name):
        super().__init__(session, tool_name)

        # Set name displayed on title bar
        self.display_name = "Duplicate Structure"

        # Create the main tool window
        self.tool_window = MainToolWindow(self)

        # Build the UI for the tool
        self._build_ui()

        # Show the tool on the user-preferred side
        self.tool_window.manage('side')

    def _build_ui(self):
        """Build the user interface."""
        layout = QVBoxLayout()

        # Dropdown to select model and chain
        self.combobox = QComboBox()
        layout.addWidget(QLabel("Choose model ID and chain to duplicate:"))
        layout.addWidget(self.combobox)

        # Populate the dropdown with available models and chains
        self.populate_combobox()

        # Duplicate button
        duplicate_button = QPushButton("Duplicate Structure")
        duplicate_button.clicked.connect(self.duplicate_structure)
        layout.addWidget(duplicate_button)

        # Set layout
        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)

    def populate_combobox(self):
        """Fetch available models and chains and populate the dropdown."""
        all_models, all_chains = self.get_models_and_chains(self.session)
        for chain in all_chains:
            self.combobox.addItem(chain)

    def get_models_and_chains(self, session):
        """Retrieve all available models and chains."""
        all_models = []
        all_chains = []
        for model in session.models.list():
            if hasattr(model, "residues"):
                all_models.append(model.id_string)
                for chain_id in set(res.chain_id for res in model.residues):
                    all_chains.append(f"Model {model.id_string}, Chain {chain_id}")
        return all_models, all_chains

    def duplicate_structure(self):
        """Execute the duplication process."""
        selected_item = self.combobox.currentText()
        if not selected_item:
            self.session.logger.warning("No model or chain selected.")
            return

        try:
            # Parse the selected model and chain
            parts = selected_item.split(", ")
            model_id = parts[0].split(" ")[1]
            chain_id = parts[1].split(" ")[1]

            # Construct and execute the duplication command
            command = f"combine #{model_id}/{chain_id}"
            self.session.logger.info(f"Executing duplication command: {command}")
            run(self.session, command)
            self.session.logger.info(f"Successfully duplicated Model {model_id}, Chain {chain_id}.")
        except Exception as e:
            self.session.logger.error(f"An error occurred while duplicating the structure: {e}")
