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


class CropStructureTool(ToolInstance):
    SESSION_ENDURING = False
    SESSION_SAVE = True
    help = "help:user/tools/cropstructure.html"

    def __init__(self, session, tool_name):
        super().__init__(session, tool_name)
        self.display_name = "Crop Structure"
        self.tool_window = MainToolWindow(self)
        self._build_ui()

    def _build_ui(self):
        all_models, all_chains = self.select_models_and_chains(self.session)

        layout = QVBoxLayout()
        self.combobox1 = QComboBox()
        for chain in all_chains:
            model, chain_id = chain.split(":")
            self.combobox1.addItem(f"Model {model}, Chain {chain_id}")
        layout.addWidget(QLabel("Choose model ID and chain ID for cropping"))
        layout.addWidget(self.combobox1)

        self.residue_input = QLineEdit()
        layout.addWidget(QLabel("Enter Residue Range you want to keep -> the rest of the Chain will be deleted (e.g., 1-50,60):"))
        layout.addWidget(self.residue_input)

        crop_button = QPushButton("ChopChop Crop")
        crop_button.clicked.connect(self.crop_structure)
        layout.addWidget(crop_button)

        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)
        self.tool_window.manage('side')

    def crop_structure(self):
        # Retrieve inputs
        selection = self.combobox1.currentText().split(", ")
        model_id = selection[0].split(" ")[1]
        chain_id = selection[1].split(" ")[1]
        residue_range = self.residue_input.text().strip()

        if not model_id or not chain_id or not residue_range:
            self.session.logger.warning("Please fill in all fields: Model ID, Chain ID, and Residue Range.")
            return

        try:
            # Parse residue range
            self.session.logger.info(f"Parsing residue range: '{residue_range}'")
            residues_to_keep = self.parse_residue_range(residue_range)
            if not residues_to_keep:
                self.session.logger.warning("Residue range is empty or invalid. Please provide a valid range (e.g., 1-50,60).")
                return
            self.session.logger.info(f"Residues to keep: {residues_to_keep}")

            # Calculate total residues in the chain
            total_residues = self.get_total_residues(model_id, chain_id)
            self.session.logger.info(f"Total residues in chain {chain_id}: {total_residues}")

            # Determine residues to remove
            residues_to_remove = [res for res in range(1, total_residues + 1) if res not in residues_to_keep]
            self.session.logger.info(f"Residues to remove: {residues_to_remove}")

            # Select residues to remove
            if residues_to_remove:
                select_command = f"select #{model_id}/{chain_id}:{','.join(map(str, residues_to_remove))}"
                self.session.logger.info(f"Executing select command: {select_command}")
                run(self.session, select_command)

                # Delete selected residues
                delete_command = "delete sel"
                self.session.logger.info(f"Executing delete command: {delete_command}")
                run(self.session, delete_command)

                # Clear selection
                clear_command = "select clear"
                self.session.logger.info(f"Executing clear command: {clear_command}")
                run(self.session, clear_command)

            self.session.logger.info(f"Successfully cropped the structure for Model {model_id}, Chain {chain_id}.")
        except ValueError as ve:
            self.session.logger.error(f"Residue range parsing error: {ve}")
        except Exception as e:
            self.session.logger.error(f"An error occurred while cropping the structure: {e}")

    def get_total_residues(self, model_id, chain_id):
        """Calculate the total number of residues in a specific chain."""
        model = next((m for m in self.session.models if m.id_string == model_id), None)
        if not model:
            raise ValueError(f"Model {model_id} not found.")
        chain_residues = [res for res in model.residues if res.chain_id == chain_id]
        return len(chain_residues)

    @staticmethod
    def parse_residue_range(residue_range):
        """Parse a residue range string into a list of integers."""
        residues = []
        for part in residue_range.split(','):
            part = part.strip()
            if not part:  # Skip empty parts
                continue
            if '-' in part:
                try:
                    start, end = map(int, part.split('-'))
                    residues.extend(range(start, end + 1))
                except ValueError:
                    raise ValueError(f"Invalid range '{part}' in residue range '{residue_range}'.")
            else:
                try:
                    residues.append(int(part))
                except ValueError:
                    raise ValueError(f"Invalid residue '{part}' in residue range '{residue_range}'.")
        return residues

    def select_models_and_chains(self, session):
        """Fetch all available models and chains."""
        all_models = []
        all_chains = []
        self.session.logger.info("Fetching available models and 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.id_string}:{chain_id}")
        self.session.logger.info(f"Found models: {all_models}")
        self.session.logger.info(f"Found chains: {all_chains}")
        return all_models, all_chains

