import html
import json
import re
import sys
import hashlib
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from PySide6.QtCore import QProcess, Qt
from PySide6.QtGui import QFont, QIcon
from PySide6.QtWidgets import (
    QApplication,
    QCheckBox,
    QDialog,
    QFileDialog,
    QLabel,
    QMainWindow,
    QMessageBox,
    QPushButton,
    QRadioButton,
    QTextEdit,
)

try:
    from .runner import build_embeddings_command
    from .ui.mainui import Ui_MainWindow
    from .ui.new_data_form import Ui_Form as Ui_NewDataForm
except ImportError:
    from runner import build_embeddings_command
    from ui.mainui import Ui_MainWindow
    from ui.new_data_form import Ui_Form as Ui_NewDataForm

try:
    from .data_manager import (
        create_data_instance,
        find_data_root,
        is_data_instance,
        model_embeddings_dir,
        resolve_data_source,
    )
except ImportError:
    from data_manager import (
        create_data_instance,
        find_data_root,
        is_data_instance,
        model_embeddings_dir,
        resolve_data_source,
    )

try:
    from benchmarks.benchmark_gp import BenchmarkGP
except Exception:
    BenchmarkGP = None


WORKSPACE_MARKER_NAME = ".morehub_workspace.json"
APP_STATE_DIR = Path.home() / ".morehub"
APP_STATE_FILE = APP_STATE_DIR / "state.json"


class App(QMainWindow):
    """Main GUI application that owns UI and button behaviors."""

    def __init__(self) -> None:
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self._apply_visual_consistency()

        self.package_root = Path(__file__).resolve().parent
        self.metadata_root = self.package_root / "metadata"
        self.workspace_root: Path | None = None

        # State tracking for processes, consoles, and input selections
        self._process_by_model = {}
        self._console_by_model = {}
        self._model_by_button = {}
        self._status_widgets_by_model = {}
        self._input_button_to_model = {}
        self._input_label_by_model = {}
        self._selected_input_by_model = {}
        self._global_data_root: Path | None = None
        self._run_embeddings_after_setup = {}
        self._format_widgets_by_model = {}
        self._embeddings_extra_args_builders = {}
        self._grover_option_widgets = {}
        self._chemberta_variant_widgets = {}
        self._novomolgen_variant_widgets = {}
        self._molgen_variant_widgets = {}
        self._molclr_backbone_widgets = {}
        self._overwrite_widget_by_model = {}
        self._active_embeddings_runs = {}
        self._benchmark_gp = BenchmarkGP(package_root=self.package_root, workspace_provider=lambda: self.workspace_root, owner=self) if BenchmarkGP is not None else None

        self._bind_run_buttons()
        self._bind_input_buttons()
        self._bind_format_controls()
        self._bind_embeddings_option_controls()
        self._bind_menu_actions()
        self._bind_status_widgets()
        self._clear_all_data_selections()
        self.ui.tabWidget.currentChanged.connect(self._on_model_tab_changed)
        self._load_persisted_workspace()
        self._update_window_title()
        self._refresh_all_statuses()
        if self.workspace_root is None:
            self._broadcast_log("[WORKSPACE] Select File -> New.. or File -> Load to choose a working directory.")

    def _apply_visual_consistency(self) -> None:
        """Normalize font and indicator sizing across machines/styles."""
        radio_button_names = [
            "radioButton",
            "radioButton_2",
            "radioButton_3",
            "radioButton_4",
            "radioButton_5",
            "radioButton_6",
            "radioButton_7",
            "radioButton_8",
            "radioButton_9",
            "radioButton_10",
            "radioButton_11",
            "radioButton_12",
            "radioButton_13",
            "radioButton_14",
            "radioButton_15",
            "radioButton_16",
            "radioButton_17",
            "radioButton_18",
            "radioButton_19",
            "radioButton_20",
            "radioButton_21",
            "radioButton_22",
            "radioButton_23",
            "radioButton_25",
            "radioButton_26",
            "radioButton_27",
            "radioButton_28",
            "radioButton_24",
            "radioButton_29",
            "radioButton_30",
            "radioButton_31",
            "radioButton_32",
            "radioButton_33",
            "radioButton_34",
            "radioButton_35",
            "radioButton_36",
            "radioButton_37",
            "radioButton_38",
        ]

        radio_font = QFont()
        radio_font.setFamilies(["Oswald", "Nirmala UI", "Segoe UI", "Arial"])
        radio_font.setPointSize(10)

        for name in radio_button_names:
            widget = getattr(self.ui, name, None)
            if not isinstance(widget, QRadioButton):
                continue
            widget.setFont(radio_font)
            widget.setStyleSheet(
                "QRadioButton { spacing: 8px; } "
                "QRadioButton::indicator { width: 14px; height: 14px; }"
            )
    def _bind_menu_actions(self) -> None:
        if hasattr(self.ui, "actionNew"):
            self.ui.actionNew.triggered.connect(self._on_new_workspace_clicked)
        if hasattr(self.ui, "actionLoad"):
            self.ui.actionLoad.triggered.connect(self._on_load_workspace_clicked)
        if hasattr(self.ui, "actionNew_Data_Setup"):
            self.ui.actionNew_Data_Setup.triggered.connect(self._on_new_data_setup_clicked)

        benchmark_action = getattr(self.ui, "actionGP_Regression", None)
        if benchmark_action is not None:
            benchmark_action.triggered.connect(self._open_gp_regression_widget)
        else:
            benchmarks_menu = self.menuBar().addMenu("Benchmarks")
            self._action_gp_regression = benchmarks_menu.addAction("GP Regression")
            self._action_gp_regression.triggered.connect(self._open_gp_regression_widget)

    def _open_gp_regression_widget(self) -> None:
        if self._benchmark_gp is None:
            QMessageBox.warning(
                self,
                "Benchmarks Unavailable",
                "GP benchmark interface could not be imported.",
            )
            return

        self._benchmark_gp.show()

    def _on_new_data_setup_clicked(self) -> None:
        if self.workspace_root is None:
            QMessageBox.information(
                self,
                "Workspace Required",
                "Select File -> New.. or File -> Load before creating a data setup.",
            )
            return

        dialog = NewDataSetupDialog(workspace_root=self.workspace_root, parent=self)
        if dialog.exec() != QDialog.Accepted:
            return

        if dialog.created_data_root is not None:
            self._broadcast_log(f"[DATA] Created data instance: {dialog.created_data_root}")
            self._set_data_selection_for_all_models(dialog.created_data_root)
            self._broadcast_log(f"[DATA] Active data instance for all models: {dialog.created_data_root}")

    def _on_new_workspace_clicked(self) -> None:
        if self._process_by_model:
            QMessageBox.information(
                self,
                "Workspace Busy",
                "Stop running setup/embedding jobs before switching workspace.",
            )
            return

        selected_dir = QFileDialog.getExistingDirectory(
            self,
            "Select new workspace directory",
            str(self.workspace_root or Path.home()),
        )
        if not selected_dir:
            return

        try:
            workspace = Path(selected_dir).resolve()
            workspace.mkdir(parents=True, exist_ok=True)
            (workspace / "models").mkdir(parents=True, exist_ok=True)
            marker = self._workspace_marker_path(workspace)
            if not marker.exists():
                marker.write_text(
                    json.dumps({"format": "morehub-workspace", "version": 1}, indent=2) + "\n",
                    encoding="utf-8",
                )
            self._set_workspace(workspace, persist=True)
        except Exception as exc:
            QMessageBox.critical(self, "Workspace Error", f"Failed to initialize workspace:\n{exc}")

    def _on_load_workspace_clicked(self) -> None:
        if self._process_by_model:
            QMessageBox.information(
                self,
                "Workspace Busy",
                "Stop running setup/embedding jobs before switching workspace.",
            )
            return

        selected_dir = QFileDialog.getExistingDirectory(
            self,
            "Load existing workspace directory",
            str(self.workspace_root or Path.home()),
        )
        if not selected_dir:
            return

        workspace = Path(selected_dir).resolve()
        if not self._is_workspace_valid(workspace):
            QMessageBox.warning(
                self,
                "Invalid Workspace",
                (
                    "Selected folder is not a valid workspace.\n\n"
                    "Expected either:\n"
                    "- a ./models directory, or\n"
                    f"- a {WORKSPACE_MARKER_NAME} marker file."
                ),
            )
            return

        self._set_workspace(workspace, persist=True)

    def _workspace_marker_path(self, workspace_root: Path) -> Path:
        return workspace_root / WORKSPACE_MARKER_NAME

    def _workspace_layout_paths(self, workspace_root: Path) -> dict[str, Path]:
        return {
            "models": workspace_root / "models",
        }

    def _ensure_workspace_layout(self, workspace_root: Path) -> None:
        for path in self._workspace_layout_paths(workspace_root).values():
            path.mkdir(parents=True, exist_ok=True)

    def _default_input_directory(self) -> Path:
        if self.workspace_root is None:
            return Path.home()
        return self.workspace_root

    def _is_workspace_valid(self, workspace_root: Path) -> bool:
        if not workspace_root.exists() or not workspace_root.is_dir():
            return False
        return (workspace_root / "models").is_dir() or self._workspace_marker_path(workspace_root).exists()

    def _set_workspace(self, workspace_root: Path, persist: bool) -> None:
        previous_workspace = self.workspace_root
        self.workspace_root = workspace_root.resolve()
        if previous_workspace != self.workspace_root:
            self._clear_all_data_selections()
        self._ensure_workspace_layout(self.workspace_root)
        marker = self._workspace_marker_path(self.workspace_root)
        if not marker.exists():
            marker.write_text(
                json.dumps({"format": "morehub-workspace", "version": 1}, indent=2) + "\n",
                encoding="utf-8",
            )
        if persist:
            self._persist_workspace(self.workspace_root)
        self._update_window_title()
        self._refresh_all_statuses()
        if self._benchmark_gp is not None:
            self._benchmark_gp.refresh_workspace_state()
        self._broadcast_log(f"[WORKSPACE] Active workspace: {self.workspace_root}")

    def _persist_workspace(self, workspace_root: Path) -> None:
        APP_STATE_DIR.mkdir(parents=True, exist_ok=True)
        payload = {"workspace_root": str(workspace_root)}
        APP_STATE_FILE.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")

    def _load_persisted_workspace(self) -> None:
        if not APP_STATE_FILE.exists():
            return
        try:
            payload = json.loads(APP_STATE_FILE.read_text(encoding="utf-8"))
            raw_workspace = payload.get("workspace_root")
            if not raw_workspace:
                return
            workspace = Path(raw_workspace).resolve()
            if self._is_workspace_valid(workspace):
                self._set_workspace(workspace, persist=False)
        except Exception:
            # Best-effort restore only.
            return

    def _update_window_title(self) -> None:
        if self.workspace_root is None:
            self.setWindowTitle("MoreHub - No workspace selected")
            return
        self.setWindowTitle(f"MoreHub - {self.workspace_root}")

    def _broadcast_log(self, text: str) -> None:
        for console in self._console_by_model.values():
            console.append(text.rstrip())

    def _set_model_data_label(self, model_name: str, data_root: Path | None) -> None:
        label = self._input_label_by_model.get(model_name)
        if label is None:
            return

        if data_root is None:
            label.setText("No data selected")
            label.setToolTip("")
            return

        label.setText(data_root.name)
        label.setToolTip(str(data_root))

    def _clear_all_data_selections(self) -> None:
        self._global_data_root = None
        self._selected_input_by_model.clear()
        for model_name in self._input_label_by_model:
            self._set_model_data_label(model_name, None)

    def _set_data_selection_for_all_models(self, data_root: Path) -> None:
        resolved = Path(data_root).resolve()
        self._global_data_root = resolved
        for model_name in self._input_label_by_model:
            self._selected_input_by_model[model_name] = resolved
            self._set_model_data_label(model_name, resolved)

    def _on_model_tab_changed(self, index: int) -> None:
        if index < 0:
            return
        model_name = self.ui.tabWidget.tabText(index).strip().upper()
        data_root = self._selected_input_by_model.get(model_name) or self._global_data_root
        if data_root is None:
            self._set_model_data_label(model_name, None)
            return
        self._set_model_data_label(model_name, Path(data_root).resolve())

    def _bind_run_buttons(self) -> None:
        run_button_names = [
            "pushButton_2",
            "pushButton_4",
            "pushButton_6",
            "pushButton_8",
            "pushButton_10",
            "pushButton_12",
            "pushButton_14",
            "pushButton_16",
            "pushButton_18",
            "pushButton_20",
            "pushButton_22",
            "pushButton_24",
        ]
        console_names = [
            "textEdit",
            "textEdit_2",
            "textEdit_3",
            "textEdit_4",
            "textEdit_5",
            "textEdit_6",
            "textEdit_7",
            "textEdit_8",
            "textEdit_9",
            "textEdit_10",
            "textEdit_11",
            "textEdit_12",
        ]

        for index, button_name in enumerate(run_button_names):
            button = getattr(self.ui, button_name, None)
            console = getattr(self.ui, console_names[index], None)
            if button is None or not isinstance(console, QTextEdit):
                continue

            tab_label = self.ui.tabWidget.tabText(index).strip()
            model_name = tab_label.upper()

            self._model_by_button[button] = model_name
            self._console_by_model[model_name] = console
            button.clicked.connect(self._on_run_clicked)
    def _bind_input_buttons(self) -> None:
        input_button_names = [
            "pushButton",
            "pushButton_3",
            "pushButton_5",
            "pushButton_7",
            "pushButton_9",
            "pushButton_11",
            "pushButton_13",
            "pushButton_15",
            "pushButton_17",
            "pushButton_19",
            "pushButton_21",
            "pushButton_23",
        ]
        input_label_names = [
            "label_2",
            "label_11",
            "label_23",
            "label_33",
            "label_43",
            "label_49",
            "label_29",
            "label_55",
            "label_59",
            "label_65",
            "label_71",
            "label_77",
        ]

        for index, button_name in enumerate(input_button_names):
            button = getattr(self.ui, button_name, None)
            label = getattr(self.ui, input_label_names[index], None)
            if not isinstance(button, QPushButton) or not isinstance(label, QLabel):
                continue

            model_name = self.ui.tabWidget.tabText(index).strip().upper()
            self._input_button_to_model[button] = model_name
            self._input_label_by_model[model_name] = label
            button.clicked.connect(self._on_input_select_clicked)

    def _on_input_select_clicked(self) -> None:
        button = self.sender()
        model_name = self._input_button_to_model.get(button)
        if not model_name:
            return

        selected_dir = QFileDialog.getExistingDirectory(
            self,
            f"Select data folder for {model_name}",
            str(self._default_input_directory()),
        )
        if not selected_dir:
            return

        data_root = Path(selected_dir).resolve()
        if not is_data_instance(data_root):
            inferred = find_data_root(data_root)
            if inferred is None or not is_data_instance(inferred):
                QMessageBox.warning(
                    self,
                    "Invalid Data Folder",
                    "Select a valid data instance folder (data_<name>) that contains .morehub_data.json.",
                )
                self._append_log(model_name, f"[{model_name}] Invalid data folder: {data_root}")
                return
            data_root = inferred

        self._set_data_selection_for_all_models(data_root)
        self._append_log(model_name, f"[{model_name}] Data selected for all models: {data_root}")
        try:
            molecules_path, _target_path, detected_format = resolve_data_source(data_root)
            self._append_log(model_name, f"[{model_name}] Data selected: {data_root}")
            self._append_log(model_name, f"[{model_name}] Source molecules: {molecules_path}")
            self._append_log(model_name, f"[{model_name}] Input format from source metadata: {detected_format}")
        except Exception:
            self._append_log(model_name, f"[{model_name}] Data selected: {data_root}")

    def _bind_format_controls(self) -> None:
        smiles_names = [
            "radioButton",
            "radioButton_3",
            "radioButton_5",
            "radioButton_7",
            "radioButton_9",
            "radioButton_11",
            "radioButton_13",
            "radioButton_18",
            "radioButton_20",
            "radioButton_24",
            "radioButton_30",
            "radioButton_35",
        ]
        selfies_names = [
            "radioButton_2",
            "radioButton_4",
            "radioButton_6",
            "radioButton_8",
            "radioButton_10",
            "radioButton_12",
            "radioButton_14",
            "radioButton_19",
            "radioButton_21",
            "radioButton_29",
            "radioButton_31",
            "radioButton_36",
        ]

        for index, smiles_name in enumerate(smiles_names):
            smiles_radio = getattr(self.ui, smiles_name, None)
            selfies_radio = getattr(self.ui, selfies_names[index], None)
            if not isinstance(smiles_radio, QRadioButton) or not isinstance(selfies_radio, QRadioButton):
                continue
            model_name = self.ui.tabWidget.tabText(index).strip().upper()
            self._format_widgets_by_model[model_name] = {
                "smiles": smiles_radio,
                "selfies": selfies_radio,
            }
            # Safe default for now while format toggle behavior is evolving.
            if not smiles_radio.isChecked() and not selfies_radio.isChecked():
                selfies_radio.setChecked(True)
    def _get_selected_input_format(self, model_name: str) -> str:
        radios = self._format_widgets_by_model.get(model_name)
        if not radios:
            return "selfies"
        if radios["smiles"].isChecked():
            return "smiles"
        return "selfies"

    def _bind_embeddings_option_controls(self) -> None:
        # Grover-specific controls from the third tab.
        use_large = getattr(self.ui, "checkBox_10", None)
        use_bonds = getattr(self.ui, "checkBox_9", None)
        use_atoms = getattr(self.ui, "checkBox_8", None)
        if (
            isinstance(use_large, QCheckBox)
            and isinstance(use_bonds, QCheckBox)
            and isinstance(use_atoms, QCheckBox)
        ):
            self._grover_option_widgets = {
                "use_large": use_large,
                "use_bonds": use_bonds,
                "use_atoms": use_atoms,
            }
            if not use_atoms.isChecked() and not use_bonds.isChecked():
                use_atoms.setChecked(True)
            self._embeddings_extra_args_builders["GROVER"] = self._build_grover_extra_args

        # ChemBERTa-specific checkpoint variant selection from the seventh tab.
        chemberta_zinc = getattr(self.ui, "radioButton_15", None)
        chemberta_mtr = getattr(self.ui, "radioButton_16", None)
        chemberta_mlm = getattr(self.ui, "radioButton_17", None)
        if (
            isinstance(chemberta_zinc, QRadioButton)
            and isinstance(chemberta_mtr, QRadioButton)
            and isinstance(chemberta_mlm, QRadioButton)
        ):
            self._chemberta_variant_widgets = {
                "zinc-base-v1": chemberta_zinc,
                "77m-mtr": chemberta_mtr,
                "77m-mlm": chemberta_mlm,
            }
            if not any(widget.isChecked() for widget in self._chemberta_variant_widgets.values()):
                chemberta_zinc.setChecked(True)
            self._embeddings_extra_args_builders["CHEMBERTA"] = self._build_chemberta_extra_args

        # NovoMolGen-specific checkpoint variant selection from the ninth tab.
        novomolgen_32m_atomwise = getattr(self.ui, "radioButton_22", None)
        novomolgen_157m_atomwise = getattr(self.ui, "radioButton_23", None)
        novomolgen_300m_atomwise = getattr(self.ui, "radioButton_25", None)
        novomolgen_32m_bpe = getattr(self.ui, "radioButton_26", None)
        novomolgen_157m_bpe = getattr(self.ui, "radioButton_27", None)
        novomolgen_300m_bpe = getattr(self.ui, "radioButton_28", None)
        if (
            isinstance(novomolgen_32m_atomwise, QRadioButton)
            and isinstance(novomolgen_157m_atomwise, QRadioButton)
            and isinstance(novomolgen_300m_atomwise, QRadioButton)
            and isinstance(novomolgen_32m_bpe, QRadioButton)
            and isinstance(novomolgen_157m_bpe, QRadioButton)
            and isinstance(novomolgen_300m_bpe, QRadioButton)
        ):
            self._novomolgen_variant_widgets = {
                "32m-atomwise": novomolgen_32m_atomwise,
                "157m-atomwise": novomolgen_157m_atomwise,
                "300m-atomwise": novomolgen_300m_atomwise,
                "32m-bpe": novomolgen_32m_bpe,
                "157m-bpe": novomolgen_157m_bpe,
                "300m-bpe": novomolgen_300m_bpe,
            }
            if not any(widget.isChecked() for widget in self._novomolgen_variant_widgets.values()):
                novomolgen_32m_atomwise.setChecked(True)
            self._embeddings_extra_args_builders["NOVOMOLGEN"] = self._build_novomolgen_extra_args

        # MolGen-specific checkpoint selection from the eleventh tab.
        molgen_large = getattr(self.ui, "radioButton_32", None)
        molgen_large_opt = getattr(self.ui, "radioButton_33", None)
        molgen_7b = getattr(self.ui, "radioButton_34", None)
        if (
            isinstance(molgen_large, QRadioButton)
            and isinstance(molgen_large_opt, QRadioButton)
            and isinstance(molgen_7b, QRadioButton)
        ):
            self._molgen_variant_widgets = {
                "large": molgen_large,
                "large-opt": molgen_large_opt,
                "7b": molgen_7b,
            }
            if not any(widget.isChecked() for widget in self._molgen_variant_widgets.values()):
                molgen_7b.setChecked(True)
            self._embeddings_extra_args_builders["MOLGEN"] = self._build_molgen_extra_args

        # MolCLR-specific backbone selection from the twelfth tab.
        molclr_gin = getattr(self.ui, "radioButton_37", None)
        molclr_gcn = getattr(self.ui, "radioButton_38", None)
        if molclr_gin is None or molclr_gcn is None:
            molclr_gin = getattr(self.ui, "checkBox_11", None)
            molclr_gcn = getattr(self.ui, "checkBox_12", None)

        if isinstance(molclr_gin, (QCheckBox, QRadioButton)) and isinstance(molclr_gcn, (QCheckBox, QRadioButton)):
            self._molclr_backbone_widgets = {
                "gin": molclr_gin,
                "gcn": molclr_gcn,
            }
            if not molclr_gin.isChecked() and not molclr_gcn.isChecked():
                molclr_gcn.setChecked(True)
            self._embeddings_extra_args_builders["MOLCLR"] = self._build_molclr_extra_args

        # Optional per-model overwrite toggles. If present in a tab and the label contains
        # "overwrite", that checkbox controls replacement of previous embeddings for that model.
        for index in range(self.ui.tabWidget.count()):
            tab = self.ui.tabWidget.widget(index)
            if tab is None:
                continue
            model_name = self.ui.tabWidget.tabText(index).strip().upper()
            for checkbox in tab.findChildren(QCheckBox):
                text = checkbox.text().strip().lower()
                if "overwrite" in text:
                    self._overwrite_widget_by_model[model_name] = checkbox
                    break
    def _build_model_extra_args(self, model_name: str) -> list[str]:
        builder = self._embeddings_extra_args_builders.get(model_name)
        if builder is None:
            return []
        return builder()

    def _extract_variant_arg(self, extra_args: list[str]) -> str | None:
        for index, token in enumerate(extra_args):
            if token == "--variant" and (index + 1) < len(extra_args):
                value = str(extra_args[index + 1]).strip()
                return value or None
        return None

    def _normalize_output_token(self, value: str) -> str:
        token = re.sub(r"[^A-Za-z0-9]+", "_", str(value)).strip("_")
        return token or "default"

    def _build_embeddings_version_tag(self, model_name: str, extra_args: list[str]) -> str | None:
        upper = model_name.upper()

        if upper == "GROVER":
            variant = self._extract_variant_arg(extra_args) or "base"
            use_atoms = "--use-atoms" in extra_args
            use_bonds = "--use-bonds" in extra_args
            if use_atoms and use_bonds:
                mode = "atomsBonds"
            elif use_atoms:
                mode = "atoms"
            elif use_bonds:
                mode = "bonds"
            else:
                mode = "atoms"
            return f"{variant}_{mode}"

        if upper == "NOVOMOLGEN":
            variant = (self._extract_variant_arg(extra_args) or "").lower()
            variant_map = {
                "32m-atomwise": "NovoMolGen_32M_SMILES_AtomWise",
                "157m-atomwise": "NovoMolGen_157M_SMILES_AtomWise",
                "300m-atomwise": "NovoMolGen_300M_SMILES_AtomWise",
                "32m-bpe": "NovoMolGen_32M_SMILES_BPE",
                "157m-bpe": "NovoMolGen_157M_SMILES_BPE",
                "300m-bpe": "NovoMolGen_300M_SMILES_BPE",
            }
            if variant in variant_map:
                return variant_map[variant]
            return variant or None

        variant = self._extract_variant_arg(extra_args)
        return variant or None

    def _build_embeddings_output_stem(self, model_name: str, extra_args: list[str]) -> tuple[str, str | None]:
        version_tag = self._build_embeddings_version_tag(model_name, extra_args)
        parts = ["embeddings", model_name.upper()]
        if version_tag:
            parts.append(version_tag)
        stem = "_".join(self._normalize_output_token(part) for part in parts if part)
        return stem, version_tag

    def _build_grover_extra_args(self) -> list[str]:
        widgets = self._grover_option_widgets
        if not widgets:
            return []

        use_large = widgets["use_large"].isChecked()
        use_bonds = widgets["use_bonds"].isChecked()
        use_atoms = widgets["use_atoms"].isChecked()

        args = ["--variant", "large" if use_large else "base"]
        if use_atoms:
            args.append("--use-atoms")
        if use_bonds:
            args.append("--use-bonds")
        if not use_atoms and not use_bonds:
            args.append("--use-atoms")
            self._append_log("GROVER", "[GROVER] Neither atoms nor bonds selected; defaulting to atoms.")
        return args

    def _build_chemberta_extra_args(self) -> list[str]:
        widgets = self._chemberta_variant_widgets
        if not widgets:
            return []

        for variant, radio in widgets.items():
            if radio.isChecked():
                return ["--variant", variant]

        # Defensive fallback for unexpected UI state.
        self._append_log("CHEMBERTA", "[CHEMBERTA] No variant selected; defaulting to zinc-base-v1.")
        return ["--variant", "zinc-base-v1"]

    def _build_novomolgen_extra_args(self) -> list[str]:
        widgets = self._novomolgen_variant_widgets
        if not widgets:
            return []

        for variant, radio in widgets.items():
            if radio.isChecked():
                return ["--variant", variant]

        # Defensive fallback for unexpected UI state.
        self._append_log("NOVOMOLGEN", "[NOVOMOLGEN] No variant selected; defaulting to 32m-atomwise.")
        return ["--variant", "32m-atomwise"]

    def _build_molgen_extra_args(self) -> list[str]:
        widgets = self._molgen_variant_widgets
        if not widgets:
            return []

        for variant, radio in widgets.items():
            if radio.isChecked():
                return ["--variant", variant]

        # Defensive fallback for unexpected UI state.
        self._append_log("MOLGEN", "[MOLGEN] No variant selected; defaulting to 7b.")
        return ["--variant", "7b"]

    def _build_molclr_extra_args(self) -> list[str]:
        widgets = self._molclr_backbone_widgets
        if not widgets:
            return []

        use_gin = widgets["gin"].isChecked()
        use_gcn = widgets["gcn"].isChecked()

        if use_gin and use_gcn:
            self._append_log("MOLCLR", "[MOLCLR] Both GIN and GCN selected; defaulting to GIN.")
            return ["--variant", "gin"]
        if use_gin:
            return ["--variant", "gin"]
        if use_gcn:
            return ["--variant", "gcn"]

        self._append_log("MOLCLR", "[MOLCLR] No backbone selected; defaulting to GCN.")
        return ["--variant", "gcn"]

    def _should_overwrite_embeddings(self, model_name: str) -> bool:
        widget = self._overwrite_widget_by_model.get(model_name)
        if widget is None:
            return False
        return bool(widget.isChecked())

    def _bind_status_widgets(self) -> None:
        model_installed_names = [
            "model_installed",
            "model_installed_2",
            "model_installed_3",
            "model_installed_4",
            "model_installed_5",
            "model_installed_6",
            "model_installed_7",
            "model_installed_8",
            "model_installed_9",
            "model_installed_10",
            "model_installed_11",
            "model_installed_12",
        ]
        checkpoints_names = [
            "checkpoints_downloaded",
            "checkpoints_downloaded_2",
            "checkpoints_downloaded_3",
            "checkpoints_downloaded_4",
            "checkpoints_downloaded_5",
            "checkpoints_downloaded_6",
            "checkpoints_downloaded_7",
            "checkpoints_downloaded_8",
            "checkpoints_downloaded_9",
            "checkpoints_downloaded_10",
            "checkpoints_downloaded_11",
            "checkpoints_downloaded_12",
        ]

        for index in range(len(model_installed_names)):
            model_checkbox = getattr(self.ui, model_installed_names[index], None)
            checkpoints_checkbox = getattr(self.ui, checkpoints_names[index], None)
            if not isinstance(model_checkbox, QCheckBox) or not isinstance(checkpoints_checkbox, QCheckBox):
                continue

            model_name = self.ui.tabWidget.tabText(index).strip().upper()
            self._status_widgets_by_model[model_name] = {
                "model_installed": model_checkbox,
                "checkpoints_downloaded": checkpoints_checkbox,
            }

            for checkbox in (model_checkbox, checkpoints_checkbox):
                checkbox.setEnabled(True)
                checkbox.setCheckable(True)
                checkbox.setFocusPolicy(Qt.NoFocus)
                checkbox.setAttribute(Qt.WA_TransparentForMouseEvents, True)
    def _refresh_all_statuses(self) -> None:
        for model_name in self._status_widgets_by_model:
            self._refresh_model_status(model_name)

    def _refresh_model_status(self, model_name: str) -> None:
        widgets = self._status_widgets_by_model.get(model_name)
        if not widgets:
            return
        model_installed, checkpoints_downloaded = self._get_model_state(model_name)
        widgets["model_installed"].setChecked(model_installed)
        widgets["checkpoints_downloaded"].setChecked(checkpoints_downloaded)

    def _load_model_metadata(self, model_name: str) -> dict[str, Any] | None:
        metadata_dir = self.metadata_root / model_name
        candidates = [
            metadata_dir / "model.json",
            metadata_dir / f"{metadata_dir.name}.json",
            metadata_dir / f"{model_name}.json",
        ]
        for path in candidates:
            if not path.exists():
                continue
            try:
                with path.open("r", encoding="utf-8") as handle:
                    return json.load(handle)
            except Exception:
                return None
        return None

    def _repository_artifacts_present(self, model_dir: Path, metadata: dict[str, Any]) -> bool | None:
        runtime_cfg = metadata.get("runtime") or {}
        runtime_source = str(runtime_cfg.get("source", "")).strip().lower()
        if not runtime_source:
            runtime_source = "repository" if metadata.get("repository") else "standalone_script"

        if runtime_source != "repository":
            return None

        repo_cfg = metadata.get("repository")
        if not isinstance(repo_cfg, dict) or not bool(repo_cfg.get("artifactsInRepository", False)):
            return None

        raw_artifact_paths = repo_cfg.get("artifactPaths") or []
        if isinstance(raw_artifact_paths, str):
            raw_artifact_paths = [raw_artifact_paths]

        model_dir_abs = model_dir.resolve()
        expected_paths = []
        for raw in raw_artifact_paths:
            if not isinstance(raw, str) or not raw.strip():
                continue
            candidate = (model_dir_abs / raw).resolve()
            try:
                candidate.relative_to(model_dir_abs)
            except ValueError:
                continue
            expected_paths.append(candidate)

        if expected_paths:
            return all(path.exists() for path in expected_paths)

        repo_subdir = str(repo_cfg.get("dir", "")).strip()
        if repo_subdir:
            return (model_dir_abs / repo_subdir).exists()
        return (model_dir_abs / ".repo").exists()

    def _get_model_state(self, model_name: str) -> tuple[bool, bool]:
        if self.workspace_root is None:
            return False, False

        model_dir = self.workspace_root / "models" / model_name
        sentinel = model_dir / ".setup_complete"
        venv_dir = model_dir / ".venv"
        checkpoints_dir_data = model_dir / "data"
        checkpoints_dir_cache = model_dir / ".cache"

        model_installed = venv_dir.exists() and sentinel.exists()

        checkpoints_downloaded: bool
        metadata = self._load_model_metadata(model_name)
        if metadata is None:
            checkpoints_downloaded = checkpoints_dir_data.exists() or checkpoints_dir_cache.exists()
        else:
            repo_artifacts_present = self._repository_artifacts_present(model_dir, metadata)
            if repo_artifacts_present is None:
                checkpoints_downloaded = checkpoints_dir_data.exists() or checkpoints_dir_cache.exists()
            else:
                checkpoints_downloaded = repo_artifacts_present

        return model_installed, checkpoints_downloaded

    def _on_run_clicked(self) -> None:
        button = self.sender()  # sender() returns the button that was clicked
        model_name = self._model_by_button.get(button)
        if not model_name:
            return

        if self.workspace_root is None:
            self._append_log(model_name, "[WORKSPACE] Select File -> New.. or File -> Load before running.")
            QMessageBox.information(
                self,
                "Workspace Required",
                "Select File -> New.. or File -> Load before running a model.",
            )
            return

        model_installed, checkpoints_downloaded = self._get_model_state(model_name)

        if model_installed and checkpoints_downloaded:
            self.run_model_embeddings(model_name)
            return

        if model_installed and not checkpoints_downloaded:
            self._append_log(
                model_name,
                f"[{model_name}] Model is installed but checkpoints are missing. Downloading checkpoints only...",
            )
            self.run_model_setup(model_name, download_only=True, run_embeddings_after=True)
            return

        if not model_installed and checkpoints_downloaded:
            self._append_log(
                model_name,
                f"[{model_name}] Checkpoints exist but model is not installed. Installing environment...",
            )
            self.run_model_setup(model_name, download_only=False, run_embeddings_after=True)
            return

        self._append_log(model_name, f"[{model_name}] Model and checkpoints missing. Running full setup...")
        self.run_model_setup(model_name, download_only=False, run_embeddings_after=True)

    def _log_color_for_line(self, line: str) -> str | None:
        normalized = line.strip().lower()
        if "setup finished successfully" in normalized or "embeddings finished:" in normalized:
            return "#1b8f3a"
        return None

    def _append_log(self, model_name: str, text: str) -> None:
        console = self._console_by_model.get(model_name)
        if console is None:
            return

        stripped = text.rstrip()
        for line in stripped.splitlines() or [stripped]:
            safe_line = html.escape(line)
            color = self._log_color_for_line(line) or "#000000"
            console.append(f"<span style=\"color: {color};\">{safe_line}</span>")

    def _sha256_file(self, path: Path) -> str:
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        return digest.hexdigest()

    def _capture_embeddings_output(self, model_name: str, process: QProcess, stream_name: str) -> None:
        if stream_name == "stdout":
            raw = bytes(process.readAllStandardOutput())
        else:
            raw = bytes(process.readAllStandardError())

        text = raw.decode(errors="replace")
        if not text:
            return

        self._append_log(model_name, text)
        run_context = self._active_embeddings_runs.get(model_name)
        if run_context is not None:
            run_context["log_chunks"].append(f"[{stream_name}] {text}")

    def _finalize_embeddings_run(self, model_name: str, exit_code: int, output_path: Path) -> None:
        run_context = self._active_embeddings_runs.pop(model_name, None)
        if run_context is None:
            return

        run_dir = run_context["run_dir"]
        log_path = run_context["log_path"]
        manifest_path = run_context["manifest_path"]

        try:
            log_path.write_text("".join(run_context["log_chunks"]), encoding="utf-8")
        except Exception as exc:
            self._append_log(model_name, f"[{model_name}] Warning: could not write run log ({exc})")

        output_files = []
        if run_dir.exists():
            for item in sorted(run_dir.iterdir()):
                if item.is_file():
                    output_files.append(item.name)

        manifest = {
            "tool": "embeddings",
            "model": model_name,
            "model_version": run_context.get("model_version"),
            "run_id": run_context["run_id"],
            "started_at": run_context["started_at"],
            "finished_at": datetime.now(timezone.utc).isoformat(),
            "exit_code": exit_code,
            "status": "success" if exit_code == 0 else "failed",
            "workspace_root": run_context.get("workspace_root") or "",
            "data_root": run_context.get("data_root") or "",
            "input": {
                "path": run_context["input_path"],
                "sha256": run_context["input_sha256"],
                "format": run_context["input_format"],
                "targets_path": run_context.get("source_targets_path"),
            },
            "extra_args": run_context["extra_args"],
            "command": run_context["command"],
            "output": {
                "stem": run_context.get("output_stem"),
                "path": str(output_path),
                "exists": output_path.exists(),
                "files": output_files,
            },
            "log_file": str(log_path),
        }

        try:
            manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
            self._append_log(model_name, f"[{model_name}] Run manifest: {manifest_path}")
        except Exception as exc:
            self._append_log(model_name, f"[{model_name}] Warning: could not write run manifest ({exc})")


    def run_model_setup(
        self,
        model_name: str,
        download_only: bool = False,
        run_embeddings_after: bool = False,
    ) -> None:
        if model_name in self._process_by_model:
            self._append_log(model_name, f"[{model_name}] Setup is already running.")
            return
        if self.workspace_root is None:
            self._append_log(model_name, "[WORKSPACE] No workspace selected.")
            return

        metadata_dir = self.metadata_root / model_name
        setup_script = metadata_dir / "setup.py"

        if not setup_script.exists():
            self._append_log(model_name, f"[{model_name}] setup.py not found at {setup_script}")
            return

        process = QProcess(self)
        process.setWorkingDirectory(str(metadata_dir))
        process.setProgram(sys.executable)
        args = [str(setup_script), "--model", model_name, "--workspace", str(self.workspace_root)]
        if download_only:
            args.append("--download-only")
        process.setArguments(args)
        process.setProcessChannelMode(QProcess.SeparateChannels)

        process.readyReadStandardOutput.connect(
            lambda m=model_name, p=process: self._append_log(
                m, bytes(p.readAllStandardOutput()).decode(errors="replace")
            )
        )
        process.readyReadStandardError.connect(
            lambda m=model_name, p=process: self._append_log(
                m, bytes(p.readAllStandardError()).decode(errors="replace")
            )
        )
        process.finished.connect(
            lambda exit_code, _status, m=model_name: self._on_setup_finished(m, exit_code)
        )

        self._process_by_model[model_name] = process
        self._run_embeddings_after_setup[model_name] = run_embeddings_after
        phase = "download-only setup" if download_only else "setup"
        self._append_log(model_name, f"[{model_name}] Starting {phase}...")
        process.start()

    def run_model_embeddings(self, model_name: str) -> None:
        if model_name in self._process_by_model:
            self._append_log(model_name, f"[{model_name}] Another process is already running.")
            return
        if self.workspace_root is None:
            self._append_log(model_name, "[WORKSPACE] No workspace selected.")
            return

        data_root = self._selected_input_by_model.get(model_name) or self._global_data_root
        if data_root is None:
            self._append_log(model_name, f"[{model_name}] No data folder selected.")
            self._append_log(model_name, "[DATA] Select once using any folder button; selection is shared across all models.")
            return
        data_root = Path(data_root).resolve()
        if not is_data_instance(data_root):
            inferred = find_data_root(data_root)
            if inferred is None or not is_data_instance(inferred):
                self._append_log(model_name, f"[{model_name}] Invalid data folder: {data_root}")
                return
            data_root = inferred
            self._set_data_selection_for_all_models(data_root)

        try:
            molecules_path, target_path, source_input_format = resolve_data_source(data_root)
        except Exception as exc:
            self._append_log(model_name, f"[{model_name}] Failed to resolve data source: {exc}")
            return

        radio_input_format = self._get_selected_input_format(model_name)
        input_format = source_input_format
        if radio_input_format != source_input_format:
            self._append_log(
                model_name,
                (
                    f"[{model_name}] UI format '{radio_input_format}' differs from source metadata "
                    f"'{source_input_format}'. Using source metadata format."
                ),
            )

        extra_args = self._build_model_extra_args(model_name)
        output_stem, model_version = self._build_embeddings_output_stem(model_name, extra_args)
        model_output_base = model_embeddings_dir(data_root, model_name)
        if model_version:
            variant_dir = self._normalize_output_token(model_version).lower()
            model_output_root = model_output_base / variant_dir
        else:
            model_output_root = model_output_base

        if model_output_root.exists():
            shutil.rmtree(model_output_root, ignore_errors=True)
            self._append_log(model_name, f"[{model_name}] Overwriting previous embeddings in: {model_output_root}")

        model_output_root.mkdir(parents=True, exist_ok=True)
        output_path = model_output_root / f"{output_stem}.npy"

        try:
            program, args, cwd, output_path, run_dir, _run_id = build_embeddings_command(
                workspace_root=self.workspace_root,
                model_name=model_name,
                input_file=molecules_path,
                input_format=input_format,
                output_file=output_path,
                extra_args=extra_args,
                metadata_root=self.metadata_root,
            )
        except FileNotFoundError as exc:
            self._append_log(model_name, f"[{model_name}] {exc}")
            return
        except Exception as exc:
            self._append_log(model_name, f"[{model_name}] Failed to build embeddings command: {exc}")
            return

        run_context = {
            "run_id": "latest",
            "model_version": model_version,
            "output_stem": output_stem,
            "run_dir": run_dir,
            "manifest_path": run_dir / "manifest.json",
            "log_path": run_dir / "stdout.log",
            "started_at": datetime.now(timezone.utc).isoformat(),
            "workspace_root": str(self.workspace_root),
            "data_root": str(data_root),
            "input_path": str(molecules_path.resolve()),
            "input_sha256": self._sha256_file(molecules_path),
            "input_format": input_format,
            "source_targets_path": str(target_path.resolve()) if target_path is not None else None,
            "extra_args": list(extra_args),
            "command": {
                "program": program,
                "args": args,
                "cwd": cwd,
            },
            "log_chunks": [],
        }
        self._active_embeddings_runs[model_name] = run_context

        process = QProcess(self)
        process.setWorkingDirectory(cwd)
        process.setProgram(program)
        process.setArguments(args)
        process.setProcessChannelMode(QProcess.SeparateChannels)

        process.readyReadStandardOutput.connect(
            lambda m=model_name, p=process: self._capture_embeddings_output(m, p, "stdout")
        )
        process.readyReadStandardError.connect(
            lambda m=model_name, p=process: self._capture_embeddings_output(m, p, "stderr")
        )
        process.finished.connect(
            lambda exit_code, _status, m=model_name, out=output_path: self._on_embeddings_finished(
                m, exit_code, out
            )
        )

        self._process_by_model[model_name] = process
        self._append_log(model_name, f"[{model_name}] Starting embeddings...")
        self._append_log(model_name, f"[{model_name}] Data root: {data_root}")
        self._append_log(model_name, f"[{model_name}] Input format: {input_format}")
        if extra_args:
            self._append_log(model_name, f"[{model_name}] Extra args: {' '.join(extra_args)}")
        self._append_log(model_name, f"[{model_name}] Output: {output_path}")
        process.start()

    def _on_setup_finished(self, model_name: str, exit_code: int) -> None:
        # Clean up the process reference and check if we need to run embeddings next
        self._process_by_model.pop(model_name, None)
        run_embeddings_after = self._run_embeddings_after_setup.pop(model_name, False)
        if exit_code == 0:
            self._append_log(model_name, f"[{model_name}] Setup finished successfully.")
        else:
            self._append_log(model_name, f"[{model_name}] Setup failed with exit code {exit_code}.")
        self._refresh_model_status(model_name)
        if exit_code == 0 and run_embeddings_after:
            self.run_model_embeddings(model_name)

    def _on_embeddings_finished(self, model_name: str, exit_code: int, output_path: Path) -> None:
        self._process_by_model.pop(model_name, None)
        if exit_code == 0:
            self._append_log(model_name, f"[{model_name}] Embeddings finished: {output_path}")
        else:
            self._append_log(model_name, f"[{model_name}] Embeddings failed with exit code {exit_code}.")
        self._finalize_embeddings_run(model_name, exit_code, output_path)



class NewDataSetupDialog(QDialog):
    def __init__(self, workspace_root: Path, parent=None) -> None:
        super().__init__(parent)
        self.ui = Ui_NewDataForm()
        self.ui.setupUi(self)

        self.workspace_root = workspace_root
        self.created_data_root: Path | None = None
        self._molecules_path: Path | None = None
        self._targets_path: Path | None = None

        self.setWindowTitle("New Data Setup")
        self._configure_icons(parent)

        self.ui.pushButton.clicked.connect(self._on_pick_molecules)
        self.ui.pushButton_3.clicked.connect(self._on_pick_targets)
        self.ui.pushButton_2.clicked.connect(self._on_process)

    def _configure_icons(self, parent) -> None:
        package_root = getattr(parent, "package_root", Path(__file__).resolve().parent)
        icons_dir = Path(package_root) / "ui" / "icons"
        folder_icon = QIcon(str(icons_dir / "folder.png"))
        run_icon = QIcon(str(icons_dir / "run.png"))

        self.ui.pushButton.setIcon(folder_icon)
        self.ui.pushButton_3.setIcon(folder_icon)
        self.ui.pushButton_2.setIcon(run_icon)

    def _on_pick_molecules(self) -> None:
        selected, _ = QFileDialog.getOpenFileName(
            self,
            "Select molecules file (selfies/smiles)",
            str(self.workspace_root),
            "Text files (*.txt);;All files (*.*)",
        )
        if not selected:
            return
        self._molecules_path = Path(selected).resolve()
        self.ui.label_2.setText(self._molecules_path.name)
        self.ui.label_2.setToolTip(str(self._molecules_path))

    def _on_pick_targets(self) -> None:
        selected, _ = QFileDialog.getOpenFileName(
            self,
            "Select targets file (optional)",
            str(self.workspace_root),
            "Text/CSV files (*.txt *.csv);;All files (*.*)",
        )
        if not selected:
            return
        self._targets_path = Path(selected).resolve()
        self.ui.label_6.setText(self._targets_path.name)
        self.ui.label_6.setToolTip(str(self._targets_path))

    def _on_process(self) -> None:
        if self._molecules_path is None:
            QMessageBox.warning(self, "Data Setup", "Select a molecules file first.")
            return

        data_name = self.ui.lineEdit.text().strip()
        if not data_name:
            QMessageBox.warning(self, "Data Setup", "Provide a name for your data.")
            return

        copy_only = bool(self.ui.checkBox.isChecked())

        try:
            self.created_data_root = create_data_instance(
                workspace_root=self.workspace_root,
                data_name=data_name,
                molecules_path=self._molecules_path,
                target_path=self._targets_path,
                copy_only=copy_only,
            )
        except Exception as exc:
            QMessageBox.critical(self, "Data Setup", f"Failed to create data setup:\n{exc}")
            return

        QMessageBox.information(self, "Data Setup", f"Created data instance:\n{self.created_data_root}")
        self.accept()


def main() -> None:
    qt_app = QApplication(sys.argv)
    qt_app.setStyle("Fusion")
    window = App()
    window.show()
    sys.exit(qt_app.exec())


if __name__ == "__main__":
    main()




