from __future__ import annotations

from typing import Any

from qtpy import QtCore, QtWidgets

import chisurf
from chisurf.plots.plotbase import Plot

# Moview / ProtView is an optional, heavy dependency. Import it lazily so
# that chisurf startup does not pull in the entire viewer stack. This
# module is imported at GUI initialization time, but the actual ProtView
# widget is only needed when a ProtViewPlot is instantiated.
ProtViewType = None  # type: ignore[assignment]


def _lazy_import_protview():
    """Return the ProtView widget class, importing it on first use.

    If the import fails (e.g. missing OpenGL), ``None`` is returned and
    callers should fall back to a placeholder widget.
    """

    global ProtViewType
    if ProtViewType is not None:
        return ProtViewType

    try:
        from chisurf.plugins.moview.viewer import ProtView as _ProtView
    except Exception:  # pragma: no cover - ProtView may not be available
        ProtViewType = None
    else:
        ProtViewType = _ProtView
    return ProtViewType


class ProtViewPlot(Plot):

    """Plot wrapper embedding the ProtView protein viewer.

    This class is intended as a lightweight replacement for the
    PyMOL-based MolView plot in structure-based models such as
    ProteinMC. It reuses the existing Plot interface but renders the
    protein using :class:`ProtView` (pyqtgraph/OpenGL) instead of
    PyMOL.
    """

    name = "ProtView"

    def __init__(self, fit, mode: str = "cartoon", **kwargs: Any) -> None:
        super().__init__(fit=fit, **kwargs)

        self.mode: str = mode
        self.structure = None

        # Replace base layout with our own
        self.layout = QtWidgets.QVBoxLayout(self)
        self.layout.setContentsMargins(0, 0, 0, 0)
        self.layout.setSpacing(0)

        pv_cls = _lazy_import_protview()
        if pv_cls is None:
            placeholder = QtWidgets.QLabel(
                "ProtViewPlot: ProtView widget is not available.\n"
                "Ensure chisurf.gui.widgets.protview can be imported."
            )
            placeholder.setAlignment(QtCore.Qt.AlignCenter)
            self.layout.addWidget(placeholder)
            self.viewer = None
            return

        self.viewer = pv_cls(self)
        self.layout.addWidget(self.viewer)

        # Try to obtain an initial structure from the fit/model
        structure = None
        try:
            model = getattr(fit, "model", None)
            structure = getattr(model, "structure", None)
            if structure is None:
                structure = getattr(fit, "data", None)
        except Exception:
            structure = None

        self.structure = structure
        if structure is not None:
            try:
                self.open_structure(structure, mode=self.mode)
            except Exception as e:  # pragma: no cover - best-effort logging
                try:
                    chisurf.logging.warning(
                        "ProtViewPlot.__init__: open_structure failed: %s", e
                    )
                except Exception:
                    pass

    # ------------------------------------------------------------------
    # API similar to MolView
    # ------------------------------------------------------------------
    def open_structure(self, structure, bfact=None, mode: str | None = None) -> None:
        """Display a new structure in the embedded ProtView widget.

        Parameters
        ----------
        structure:
            ChiSurf :class:`Structure`-like object exposing an ``atoms``
            array compatible with :mod:`chisurf.structure`.
        bfact:
            Optional B-factor array (currently ignored, reserved for
            future use such as coloring by B-factor).
        mode:
            Optional view mode string; if omitted the previously stored
            :attr:`mode` is used. Supported values are "cartoon",
            "coarse", and "surface". Other values are currently
            ignored.
        """

        if self.viewer is None:
            return

        self.structure = structure
        if mode is not None:
            self.mode = mode

        try:
            self.viewer.set_structure(structure)
        except Exception as e:  # pragma: no cover - best-effort logging
            try:
                chisurf.logging.warning(
                    "ProtViewPlot.open_structure: set_structure failed: %s", e
                )
            except Exception:
                pass
            return

        # Map legacy MolView "mode" strings to ProtView visual settings.
        m = (self.mode or "").lower()

        try:
            # Prefer the new global representation API if available.
            set_repr = getattr(self.viewer, "set_representation", None)

            if callable(set_repr):
                if m == "coarse":
                    # Coarse: emphasize atoms/balls.
                    set_repr("atoms")
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(False)
                elif m == "surface":
                    # Surface: keep cartoon backbone and enable surface point cloud.
                    set_repr("cartoon")
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(True)
                else:  # "cartoon" and any other fall back here
                    set_repr("cartoon")
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(False)
            else:
                # Fallback for older ProtView versions: use per-residue masks.
                n_res = 0
                if hasattr(self.viewer, "residue_oneletter"):
                    seq = self.viewer.residue_oneletter
                    if seq is not None:
                        n_res = len(seq)
                if n_res > 0:
                    indices = list(range(n_res))
                else:
                    indices = []

                if m == "coarse":
                    if indices:
                        self.viewer.set_residue_representation(
                            indices, cartoon=True, ball=True
                        )
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(False)
                elif m == "surface":
                    if indices:
                        self.viewer.set_residue_representation(
                            indices, cartoon=True, ball=False
                        )
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(True)
                else:
                    if indices:
                        self.viewer.set_residue_representation(
                            indices, cartoon=True, ball=False
                        )
                    if hasattr(self.viewer, "set_surface_visible"):
                        self.viewer.set_surface_visible(False)
        except Exception:  # pragma: no cover - visualization is best-effort
            pass

        try:
            self.viewer.reset_view()
        except Exception:
            pass

    def reset(self) -> None:
        """Reset viewer camera (for compatibility with MolView.reset)."""

        if getattr(self, "viewer", None) is None:
            return
        try:
            self.viewer.reset_view()
        except Exception:
            pass
