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

class FoldseekAnalysis(ToolInstance):

    SESSION_ENDURING = False
    SESSION_SAVE = True
    help = "https://www.cgl.ucsf.edu/chimerax/docs/user/tools/foldseek.html"

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

        # Set name displayed on title bar
        self.display_name = "Foldseek Analysis Tool"
        self.tool_window = MainToolWindow(self)
        self._build_ui()
        self.tool_window.manage('side')

    def _build_ui(self):
        layout = QVBoxLayout()

        # Title
        title = QLabel("Foldseek: Structural Alignment and Analysis")
        title_font = QFont()
        title_font.setBold(True)
        title_font.setPointSize(14)
        title.setFont(title_font)
        layout.addWidget(title)

        # Description
        description = QLabel(
            "Foldseek provides tools for structural alignment and analysis."
        )
        description.setWordWrap(True)
        layout.addWidget(description)

        # Model Selection
        model_label = QLabel("Select Model:")
        layout.addWidget(model_label)

        self.model_selector = QComboBox()
        self.chain_selector = QComboBox()
        self._populate_models_and_chains()
        layout.addWidget(self.model_selector)

        # Chain Selection
        chain_label = QLabel("Select Chain:")
        layout.addWidget(chain_label)
        layout.addWidget(self.chain_selector)

        # Run Button
        run_button = QPushButton("ChopChop Foldseek")
        run_button.clicked.connect(self._run_foldseek)
        layout.addWidget(run_button)

        # Set layout to the tool window
        container = QWidget()
        container.setLayout(layout)
        self.tool_window.ui_area.setLayout(layout)

    def _populate_models_and_chains(self):
        self.model_selector.clear()
        self.chain_selector.clear()
        models, chains = self.select_models_and_chains(self.session)

        for model in models:
            self.model_selector.addItem(model)
        for chain in chains:
            self.chain_selector.addItem(chain)

    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

    def _run_foldseek(self):
        model = self.model_selector.currentText()
        chain = self.chain_selector.currentText()

        if not model or not chain:
            self.session.logger.error("Model and chain must be selected.")
            return

        command = f"foldseek #{model} /{chain.split(':')[1]}"
        self.session.logger.info(f"Running command: {command}")
        run(self.session, command)

# Register the tool in ChimeraX
from chimerax.core.toolshed import BundleAPI

class _FoldseekAnalysisBundleAPI(BundleAPI):

    @staticmethod
    def start_tool(session, tool_name):
        return FoldseekAnalysis(session, tool_name)

bundle_api = _FoldseekAnalysisBundleAPI()
