from __future__ import annotations

from typing import TYPE_CHECKING, Literal, overload

from bec_lib.endpoints import MessageEndpoints
from bec_lib.logger import bec_logger
from bec_lib.messages import ScanStatusMessage

from bec_widgets.utils.error_popups import SafeSlot
from bec_widgets.widgets.containers.dock_area.dock_area import BECDockArea
from bec_widgets.widgets.containers.main_window.main_window import BECMainWindow
from bec_widgets.widgets.containers.qt_ads import CDockWidget

if TYPE_CHECKING:  # pragma: no cover
    from bec_widgets.utils.bec_widget import BECWidget
    from bec_widgets.widgets.plots.image.image import Image
    from bec_widgets.widgets.plots.motor_map.motor_map import MotorMap
    from bec_widgets.widgets.plots.multi_waveform.multi_waveform import MultiWaveform
    from bec_widgets.widgets.plots.scatter_waveform.scatter_waveform import ScatterWaveform
    from bec_widgets.widgets.plots.waveform.waveform import Waveform


logger = bec_logger.logger


class AutoUpdates(BECMainWindow):
    _default_dock: CDockWidget | None
    USER_ACCESS = ["enabled", "enabled.setter", "selected_device", "selected_device.setter"]
    RPC = True
    PLUGIN = False

    # enforce that subclasses have the same rpc widget class
    rpc_widget_class = "AutoUpdates"

    def __init__(
        self, parent=None, gui_id: str = None, window_title="Auto Update", *args, **kwargs
    ):
        super().__init__(parent=parent, gui_id=gui_id, window_title=window_title, **kwargs)

        self.dock_area = BECDockArea(
            parent=self,
            object_name="dock_area",
            enable_profile_management=False,
            startup_profile="skip",
        )
        self.setCentralWidget(self.dock_area)
        self._auto_update_selected_device: str | None = None

        self._default_dock = None  # type: ignore
        self.current_widget: BECWidget | None = None
        self.dock_name = None
        self._enabled = True
        self.start_auto_update()

    def start_auto_update(self):
        """
        Establish all connections for the auto updates.
        """
        self.bec_dispatcher.connect_slot(self._on_scan_status, MessageEndpoints.scan_status())

    def stop_auto_update(self):
        """
        Disconnect all connections for the auto updates.
        """
        self.bec_dispatcher.disconnect_slot(
            self._on_scan_status, MessageEndpoints.scan_status()  # type: ignore
        )

    @property
    def selected_device(self) -> str | None:
        """
        Get the selected device from the auto update config.

        Returns:
            str: The selected device. If no device is selected, None is returned.
        """
        return self._auto_update_selected_device

    @selected_device.setter
    def selected_device(self, value: str | None) -> None:
        """
        Set the selected device in the auto update config.

        Args:
            value(str): The selected device.
        """
        self._auto_update_selected_device = value

    @SafeSlot()
    def _on_scan_status(self, content: dict, metadata: dict) -> None:
        """
        Callback for scan status messages.
        """
        msg = ScanStatusMessage(**content, metadata=metadata)
        if not self.enabled:
            return

        self.enable_gui_highlights(True)

        match msg.status:
            case "open":
                self.on_scan_open(msg)
            case "closed":
                self.on_scan_closed(msg)
            case ["aborted", "halted"]:
                self.on_scan_abort(msg)
            case _:
                pass

    def start_default_dock(self):
        """
        Create a default dock for the auto updates.
        """
        self.dock_area.delete_all()
        self.dock_name = "update_dock"
        self.current_widget = self.dock_area.new("Waveform")
        docks = self.dock_area.dock_list()
        self._default_dock = docks[0] if docks else None

    @overload
    def set_dock_to_widget(self, widget: Literal["Waveform"]) -> Waveform: ...

    @overload
    def set_dock_to_widget(self, widget: Literal["Image"]) -> Image: ...

    @overload
    def set_dock_to_widget(self, widget: Literal["ScatterWaveform"]) -> ScatterWaveform: ...

    @overload
    def set_dock_to_widget(self, widget: Literal["MotorMap"]) -> MotorMap: ...

    @overload
    def set_dock_to_widget(self, widget: Literal["MultiWaveform"]) -> MultiWaveform: ...

    def set_dock_to_widget(
        self,
        widget: Literal["Waveform", "Image", "ScatterWaveform", "MotorMap", "MultiWaveForm"] | str,
    ) -> BECWidget:
        """
        Set the dock to the widget.

        Args:
            widget (str): The widget to set the dock to. Must be the name of a valid widget class.

        Returns:
            BECWidget: The widget that was set.
        """
        if self.current_widget is None:
            logger.warning(
                f"Auto Updates: No default dock found. Creating a new one with name {self.dock_name}"
            )
            self.start_default_dock()
        assert self.current_widget is not None

        if self.current_widget.__class__.__name__ != widget:
            self.dock_area.delete_all()
            self.current_widget = self.dock_area.new(widget)
            docks = self.dock_area.dock_list()
            self._default_dock = docks[0] if docks else None
        return self.current_widget

    def get_selected_device(
        self, monitored_devices, selected_device: str | None = None
    ) -> str | None:
        """
        Get the selected device for the plot. If no device is selected, the first
        device in the monitored devices list is selected.
        """

        if selected_device is None:
            selected_device = self.selected_device
        if selected_device:
            return selected_device
        if len(monitored_devices) > 0:
            sel_device = monitored_devices[0]
            return sel_device
        return None

    def enable_gui_highlights(self, enable: bool) -> None:
        """
        Enable or disable GUI highlights.

        Args:
            enable (bool): Whether to enable or disable the highlights.
        """
        if enable:
            title = self.dock_area.window().windowTitle()
            if " [Auto Updates]" in title:
                return
            self.dock_area.window().setWindowTitle(f"{title} [Auto Updates]")
        else:
            title = self.dock_area.window().windowTitle()
            self.dock_area.window().setWindowTitle(title.replace(" [Auto Updates]", ""))

    @property
    def enabled(self) -> bool:
        """
        Get the enabled status of the auto updates.
        """
        return self._enabled

    @enabled.setter
    def enabled(self, value: bool) -> None:
        """
        Set the enabled status of the auto updates.
        """
        if self._enabled == value:
            return
        self._enabled = value

        if value:
            self.start_auto_update()
            self.enable_gui_highlights(True)
            self.on_start()
        else:
            self.stop_auto_update()
            self.enable_gui_highlights(False)
            self.on_stop()

    def cleanup(self) -> None:
        """
        Cleanup procedure to run when the auto updates are disabled.
        """
        self.enabled = False
        self.stop_auto_update()
        self.dock_area.close()
        self.dock_area.deleteLater()
        self.dock_area = None
        super().cleanup()

    ########################################################################
    ################# Update Functions #####################################
    ########################################################################

    def simple_line_scan(self, info: ScanStatusMessage) -> None:
        """
        Simple line scan.

        Args:
            info (ScanStatusMessage): The scan status message.
        """

        # Set the dock to the waveform widget
        wf = self.set_dock_to_widget("Waveform")

        # Get the scan report devices reported by the scan
        dev_x = info.scan_report_devices[0]  # type: ignore

        # For the y axis, get the selected device
        dev_y = self.get_selected_device(info.readout_priority["monitored"])  # type: ignore
        if not dev_y:
            return

        # Clear the waveform widget and plot the data
        # with the scan number and device names
        # as the label and title
        wf.clear_all()
        wf.plot(
            device_x=dev_x,
            device_y=dev_y,
            label=f"Scan {info.scan_number} - {dev_y}",
            title=f"Scan {info.scan_number}",
            x_label=dev_x,
            y_label=dev_y,
        )

        logger.info(
            f"Auto Update [simple_line_scan]: Started plot with: device_x={dev_x}, device_y={dev_y}"
        )

    def simple_grid_scan(self, info: ScanStatusMessage) -> None:
        """
        Simple grid scan.

        Args:
            info (ScanStatusMessage): The scan status message.
        """
        # Set the dock to the scatter waveform widget
        scatter = self.set_dock_to_widget("ScatterWaveform")

        # Get the scan report devices reported by the scan
        dev_x, dev_y = info.scan_report_devices[0], info.scan_report_devices[1]  # type: ignore
        dev_z = self.get_selected_device(info.readout_priority["monitored"])  # type: ignore

        if None in (dev_x, dev_y, dev_z):
            return

        # Clear the scatter waveform widget and plot the data
        scatter.clear_all()
        scatter.plot(
            device_x=dev_x,
            device_y=dev_y,
            device_z=dev_z,
            label=f"Scan {info.scan_number} - {dev_z}",
        )

        logger.info(
            f"Auto Update [simple_grid_scan]: Started plot with: device_x={dev_x}, device_y={dev_y}, device_z={dev_z}"
        )

    def best_effort(self, info: ScanStatusMessage) -> None:
        """
        Best effort scan.

        Args:
            info (ScanStatusMessage): The scan status message.
        """

        # If the scan report devices are empty, there is nothing we can do
        if not info.scan_report_devices:
            return
        dev_x = info.scan_report_devices[0]  # type: ignore
        dev_y = self.get_selected_device(info.readout_priority["monitored"])  # type: ignore
        if not dev_y:
            return

        # Set the dock to the waveform widget
        wf = self.set_dock_to_widget("Waveform")

        # Clear the waveform widget and plot the data
        wf.clear_all()
        wf.plot(
            device_x=dev_x,
            device_y=dev_y,
            label=f"Scan {info.scan_number} - {dev_y}",
            title=f"Scan {info.scan_number}",
            x_label=dev_x,
            y_label=dev_y,
        )

        logger.info(
            f"Auto Update [best_effort]: Started plot with: device_x={dev_x}, device_y={dev_y}"
        )

    #######################################################################
    ################# GUI Callbacks #######################################
    #######################################################################

    def on_start(self) -> None:
        """
        Procedure to run when the auto updates are enabled.
        """
        self.start_default_dock()

    def on_stop(self) -> None:
        """
        Procedure to run when the auto updates are disabled.
        """

    def on_scan_open(self, msg: ScanStatusMessage) -> None:
        """
        Procedure to run when a scan starts.

        Args:
            msg (ScanStatusMessage): The scan status message.
        """
        if msg.scan_name == "line_scan" and msg.scan_report_devices:
            return self.simple_line_scan(msg)
        if msg.scan_name == "grid_scan" and msg.scan_report_devices:
            return self.simple_grid_scan(msg)
        if msg.scan_report_devices:
            return self.best_effort(msg)
        return None

    def on_scan_closed(self, msg: ScanStatusMessage) -> None:
        """
        Procedure to run when a scan ends.

        Args:
            msg (ScanStatusMessage): The scan status message.
        """

    def on_scan_abort(self, msg: ScanStatusMessage) -> None:
        """
        Procedure to run when a scan is aborted.

        Args:
            msg (ScanStatusMessage): The scan status message.
        """
