from __future__ import annotations

import html
import sys
from pathlib import Path

from PySide6.QtCore import QProcess, Qt
from PySide6.QtGui import QIcon
from PySide6.QtWidgets import QFileDialog, QMessageBox, QWidget

from benchmarks.GP_regressor.ui.gp_regressor import Ui_Form

DATA_SENTINEL = ".morehub_data.json"


class GPRegressorWidget(QWidget):
    """GUI controller for GP benchmark execution."""

    def __init__(self, workspace_provider, package_root: Path, parent=None) -> None:
        super().__init__(parent)
        self.ui = Ui_Form()
        self.ui.setupUi(self)

        self._workspace_provider = workspace_provider
        self._package_root = package_root

        self._single_embeddings_path: Path | None = None
        self._targets_path: Path | None = None
        self._batch_root_path: Path | None = None
        self._process: QProcess | None = None

        self._compute_rmse_toggle = None
        self._split_data_toggle = None
        self._current_run_succeeded = False

        self._configure_ui()
        self._bind_controls()
        self.refresh_workspace_state()

    def _resolve_toggle_widgets(self) -> None:
        self._compute_rmse_toggle = getattr(self.ui, "checkBox_8", None)
        if self._compute_rmse_toggle is None:
            self._compute_rmse_toggle = getattr(self.ui, "radioButton_8", None)

        self._split_data_toggle = getattr(self.ui, "checkBox_9", None)
        if self._split_data_toggle is None:
            self._split_data_toggle = getattr(self.ui, "radioButton_9", None)

    def _configure_ui(self) -> None:
        self.setWindowTitle("GP Regression")
        self._resolve_toggle_widgets()

        if self._compute_rmse_toggle is not None and not self._compute_rmse_toggle.isChecked():
            self._compute_rmse_toggle.setChecked(True)

        self.ui.spinBox_2.setMinimum(1)
        if self.ui.spinBox_2.value() < 1:
            self.ui.spinBox_2.setValue(5)

        self.ui.model_installed_4.setText("Environment ready")
        self.ui.checkpoints_downloaded_4.setText("Run succeeded")
        for checkbox in (self.ui.model_installed_4, self.ui.checkpoints_downloaded_4):
            checkbox.setEnabled(True)
            checkbox.setCheckable(True)
            checkbox.setFocusPolicy(Qt.NoFocus)
            checkbox.setAttribute(Qt.WA_TransparentForMouseEvents, True)

        icons_dir = self._package_root / "benchmarks" / "GP_regressor" / "ui" / "icons"
        if not icons_dir.exists():
            icons_dir = self._package_root / "ui" / "icons"

        folder_icon = QIcon(str(icons_dir / "folder.png"))
        run_icon = QIcon(str(icons_dir / "run.png"))

        self.ui.pushButton_5.setIcon(folder_icon)
        self.ui.pushButton_7.setIcon(folder_icon)
        self.ui.pushButton_9.setIcon(folder_icon)
        self.ui.pushButton_6.setIcon(run_icon)
        self.ui.pushButton_8.setIcon(run_icon)

        self._update_split_controls_enabled()

    def _bind_controls(self) -> None:
        self.ui.pushButton_5.clicked.connect(self._on_select_single_embeddings)
        self.ui.pushButton_7.clicked.connect(self._on_select_targets)
        self.ui.pushButton_9.clicked.connect(self._on_select_batch_root)

        self.ui.pushButton_6.clicked.connect(self._on_run_single)
        self.ui.pushButton_8.clicked.connect(self._on_run_batch)

        if self._split_data_toggle is not None:
            self._split_data_toggle.toggled.connect(self._on_split_data_toggled)

    def _on_split_data_toggled(self, _checked: bool) -> None:
        self._update_split_controls_enabled()

    def _update_split_controls_enabled(self) -> None:
        enabled = self._is_split_data_enabled()
        self.ui.spinBox_2.setEnabled(enabled)
        self.ui.doubleSpinBox_2.setEnabled(enabled)

    def _is_compute_rmse_enabled(self) -> bool:
        if self._compute_rmse_toggle is None:
            return True
        return bool(self._compute_rmse_toggle.isChecked())

    def _is_split_data_enabled(self) -> bool:
        if self._split_data_toggle is None:
            return False
        return bool(self._split_data_toggle.isChecked())

    def _workspace_root(self) -> Path | None:
        workspace = self._workspace_provider()
        if workspace is None:
            return None
        return Path(workspace)

    def _default_root_dir(self) -> Path:
        workspace = self._workspace_root()
        if workspace is None:
            return Path.home()
        return workspace

    def _find_data_root(self, start_path: Path) -> Path | None:
        path = start_path.resolve()
        for candidate in [path, *path.parents]:
            if (candidate / DATA_SENTINEL).exists():
                return candidate
        return None

    def _is_data_instance(self, data_root: Path) -> bool:
        return (data_root / DATA_SENTINEL).exists()

    def _resolve_data_root_from_path(self, selected_path: Path | None) -> Path | None:
        if selected_path is None:
            return None
        data_root = self._find_data_root(selected_path)
        if data_root is None or not self._is_data_instance(data_root):
            return None
        return data_root

    def _active_data_root(self) -> Path | None:
        for path in (self._single_embeddings_path, self._batch_root_path, self._targets_path):
            data_root = self._resolve_data_root_from_path(path)
            if data_root is not None:
                return data_root
        return None

    def _workspace_data_roots(self) -> list[Path]:
        workspace = self._workspace_root()
        if workspace is None or not workspace.exists():
            return []

        roots = []
        for candidate in workspace.iterdir():
            if not candidate.is_dir() or not candidate.name.startswith("data_"):
                continue
            if self._is_data_instance(candidate):
                roots.append(candidate)

        roots.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return roots

    def refresh_workspace_state(self) -> None:
        data_root = self._active_data_root()
        if data_root is None:
            roots = self._workspace_data_roots()
            if roots:
                data_root = roots[0]
            else:
                self.ui.model_installed_4.setChecked(False)
                self.ui.checkpoints_downloaded_4.setChecked(False)
                return
        benchmark_root = data_root / "data" / "benchmarks" / "GP_regressor"
        venv_python = benchmark_root / ".venv" / ("Scripts" if sys.platform.startswith("win") else "bin") / (
            "python.exe" if sys.platform.startswith("win") else "python"
        )

        self.ui.model_installed_4.setChecked(venv_python.exists())

        self.ui.checkpoints_downloaded_4.setChecked(self._current_run_succeeded)

    def _set_inputs_enabled(self, enabled: bool) -> None:
        for widget in (
            self.ui.pushButton_5,
            self.ui.pushButton_7,
            self.ui.pushButton_9,
            self.ui.pushButton_6,
            self.ui.pushButton_8,
        ):
            widget.setEnabled(enabled)

    def _log_color_for_line(self, line: str) -> str | None:
        normalized = line.strip().lower()
        if "completed successfully" in normalized or "run completed successfully" in normalized:
            return "#1b8f3a"
        if "failed" in normalized or "error" in normalized or "traceback" in normalized:
            return "#b3261e"
        return None

    def _append_log(self, text: str) -> None:
        stripped = text.rstrip()
        for line in stripped.splitlines() or [stripped]:
            safe_line = html.escape(line)
            color = self._log_color_for_line(line) or "#000000"
            self.ui.textEdit_3.append(f"<span style=\"color: {color};\">{safe_line}</span>")

    def _on_select_single_embeddings(self) -> None:
        selected, _ = QFileDialog.getOpenFileName(
            self,
            "Select embeddings .npy file",
            str(self._default_root_dir()),
            "NumPy files (*.npy);;All files (*.*)",
        )
        if not selected:
            return

        path = Path(selected).resolve()
        self._single_embeddings_path = path
        self.ui.label_23.setText(path.name)
        self.ui.label_23.setToolTip(str(path))
        self._append_log(f"[GP] Embeddings selected: {path}")

        data_root = self._resolve_data_root_from_path(path)
        if data_root is not None:
            self._append_log(f"[GP] Data root: {data_root}")
        else:
            self._append_log("[GP] Warning: selected embeddings file is outside a data_<name> instance.")

        self.refresh_workspace_state()

    def _on_select_targets(self) -> None:
        selected, _ = QFileDialog.getOpenFileName(
            self,
            "Select targets file",
            str(self._default_root_dir()),
            "Tabular files (*.csv *.txt);;All files (*.*)",
        )
        if not selected:
            return

        path = Path(selected).resolve()
        self._targets_path = path
        self.ui.label_27.setText(path.name)
        self.ui.label_27.setToolTip(str(path))
        self._append_log(f"[GP] Targets selected: {path}")
        self.refresh_workspace_state()

    def _on_select_batch_root(self) -> None:
        selected = QFileDialog.getExistingDirectory(
            self,
            "Select batch root folder",
            str(self._default_root_dir()),
        )
        if not selected:
            return

        path = Path(selected).resolve()
        self._batch_root_path = path
        self.ui.label_31.setText(path.name)
        self.ui.label_31.setToolTip(str(path))
        self._append_log(f"[GP] Batch root selected: {path}")

        data_root = self._resolve_data_root_from_path(path)
        if data_root is not None:
            self._append_log(f"[GP] Data root: {data_root}")
        else:
            self._append_log("[GP] Warning: selected batch root is outside a data_<name> instance.")

        self.refresh_workspace_state()

    def _selected_kernels(self) -> list[str]:
        kernels = []
        if self.ui.gaussian_checkbox_2.isChecked():
            kernels.append("gaussian")
        if self.ui.laplacian_checkbox_2.isChecked():
            kernels.append("laplacian")
        if self.ui.mattern52_checkbox_2.isChecked():
            kernels.append("matern52")
        if not kernels:
            kernels = ["gaussian"]
            self._append_log("[GP] No kernel selected. Defaulting to gaussian.")
        return kernels

    def _build_base_launcher_args(self, mode: str, data_root: Path) -> list[str]:
        n_splits = max(1, self.ui.spinBox_2.value())
        train_ratio = self.ui.doubleSpinBox_2.value()
        compute_rmse = self._is_compute_rmse_enabled()
        regenerate_splits = self._is_split_data_enabled()

        launcher = self._package_root / "benchmarks" / "GP_regressor" / "launcher.py"
        if not launcher.exists():
            raise FileNotFoundError(f"GP launcher script not found at {launcher}")

        args = [
            str(launcher),
            "--data_root",
            str(data_root),
            "--mode",
            mode,
            "--n_splits",
            str(n_splits),
            "--train_ratio",
            f"{train_ratio}",
            "--compute_rmse",
            "1" if compute_rmse else "0",
            "--regenerate_splits",
            "1" if regenerate_splits else "0",
        ]

        for kernel in self._selected_kernels():
            args.extend(["--kernel", kernel])

        return args

    def _start_process(self, args: list[str]) -> None:
        if self._process is not None:
            self._append_log("[GP] A benchmark run is already in progress.")
            return

        # New run requested: reset success state until the process exits successfully.
        self._current_run_succeeded = False
        self.ui.checkpoints_downloaded_4.setChecked(False)

        process = QProcess(self)
        process.setProgram(sys.executable)
        process.setArguments(args)
        process.setWorkingDirectory(str(self._package_root / "benchmarks" / "GP_regressor"))
        process.setProcessChannelMode(QProcess.SeparateChannels)

        process.readyReadStandardOutput.connect(
            lambda p=process: self._append_log(bytes(p.readAllStandardOutput()).decode(errors="replace"))
        )
        process.readyReadStandardError.connect(
            lambda p=process: self._append_log(bytes(p.readAllStandardError()).decode(errors="replace"))
        )
        process.finished.connect(self._on_process_finished)

        self._process = process
        self._set_inputs_enabled(False)
        self._append_log("[GP] Starting benchmark run...")
        self._append_log(f"[GP] Running: {sys.executable} {' '.join(args)}")
        process.start()

    def _on_run_single(self) -> None:
        try:
            if self._single_embeddings_path is None:
                raise RuntimeError("Select an embeddings .npy file for single mode.")
            if self._targets_path is None:
                raise RuntimeError("Select a targets file before running.")

            data_root = self._resolve_data_root_from_path(self._single_embeddings_path)
            if data_root is None:
                raise RuntimeError("Selected embeddings file must be inside a data_<name> instance.")

            args = self._build_base_launcher_args(mode="single", data_root=data_root)
            args.extend(["--embeddings_path", str(self._single_embeddings_path)])
            args.extend(["--targets_path", str(self._targets_path)])
            self._start_process(args)
        except Exception as exc:
            QMessageBox.warning(self, "GP Regression", str(exc))
            self._append_log(f"[GP] {exc}")

    def _on_run_batch(self) -> None:
        try:
            if self._batch_root_path is None:
                raise RuntimeError("Select an embeddings root before running batch mode.")
            if self._targets_path is None:
                raise RuntimeError("Select a targets file before running.")

            data_root = self._resolve_data_root_from_path(self._batch_root_path)
            if data_root is None:
                raise RuntimeError("Selected batch root must be inside a data_<name> instance.")

            args = self._build_base_launcher_args(mode="batch", data_root=data_root)
            args.extend(["--root_path", str(self._batch_root_path)])
            args.extend(["--targets_path", str(self._targets_path)])
            self._start_process(args)
        except Exception as exc:
            QMessageBox.warning(self, "GP Regression", str(exc))
            self._append_log(f"[GP] {exc}")

    def _on_process_finished(self, exit_code: int, _status) -> None:
        self._set_inputs_enabled(True)
        self._process = None

        if exit_code == 0:
            self._current_run_succeeded = True
            self._append_log("[GP] Benchmark run completed successfully.")
        else:
            self._current_run_succeeded = False
            self._append_log(f"[GP] Benchmark run failed with exit code {exit_code}.")

        self.ui.checkpoints_downloaded_4.setChecked(self._current_run_succeeded)
        self.refresh_workspace_state()

