"""Utilities to mock and test devices."""

import threading
from contextlib import contextmanager
from functools import partial
from time import sleep
from typing import TYPE_CHECKING, Callable, Generator, TypeVar
from unittest import mock

import ophyd
from bec_lib.devicemanager import ScanInfo
from bec_lib.logger import bec_logger
from bec_lib.utils.import_utils import lazy_import_from
from ophyd import Device
from ophyd.utils.epics_pvs import AlarmSeverity, AlarmStatus

if TYPE_CHECKING:
    from bec_lib.messages import ScanStatusMessage
else:
    # TODO: put back normal import when Pydantic gets faster
    ScanStatusMessage = lazy_import_from("bec_lib.messages", ("ScanStatusMessage",))

logger = bec_logger.logger

T = TypeVar("T", bound=Device)


@contextmanager
def patched_device(
    device_type: type[T],
    *args,
    _mock_pv_initial_value=0,
    mask_stage_sigs: bool = True,
    mask_trigger_sigs: bool = True,
    **kwargs,
) -> Generator[T, None, None]:
    """Context manager to yield a patched ophyd device with certain initialisation args.
    *args and **kwargs are passed directly through to the device constructor. In addition,
    stage and trigger signals will be patched to avoid side effects for PVs with enums and 'string' types
    for which the MockPV patch fails to imitate the expected behaviour. The 'string=True' converts any set
    value to a string, but then misses to properly handle enum types, where both the string and integer
    representation are allowed. This leads to tests hanging indefinitely in a 'set' call if stage_sigs or trigger_signals
    are defined. Most often this is the case for AreaDetector classes.

    Example:
        @pytest.fixture(scope="function")
        def mock_ddg():
            with patched_device(DelayGenerator, name="ddg", prefix="X12SA-CPCL-DDG3:") as ddg:
                yield ddg

    """
    with mock.patch.object(ophyd, "cl") as mock_cl:
        mock_cl.get_pv = partial(MockPV, _mock_pv_initial_value=_mock_pv_initial_value)
        mock_cl.thread_class = threading.Thread
        device = device_type(*args, **kwargs)
        patch_dual_pvs(device)
        patch_functions_required_for_connection(device)

        # NOTE Patch stage signals to avoid issues with enums and string types
        # These enums cannot be known as they may depend on individual PV configurations
        # Therefore we simply remove them for the purpose of testing
        if mask_stage_sigs is True:
            for _, sub_device in device.walk_subdevices(include_lazy=True):
                sub_device.stage_sigs = {}  # Remove stage signals
                if hasattr(sub_device, "_plugin_type") and hasattr(sub_device, "plugin_type"):
                    # Patch plugin_type configured/misconfigure checks in AD plugins
                    sub_device.plugin_type._read_pv.mock_data = sub_device._plugin_type

        # NOTE Patch trigger signals for the same reason as stage_sigs
        if mask_trigger_sigs is True:
            for cpt_walk in device.walk_components():
                cpt_walk.item.trigger_value = None  # Remove any trigger value indicators
        yield device


def patch_dual_pvs(device):
    """Patch dual PVs"""
    patch_functions_required_for_connection(device)
    try:
        device.wait_for_connection(all_signals=True)
    except TypeError:
        # Maybe no all_signals in kwargs,
        device.wait_for_connection()

    for walk in device.walk_signals():
        if not hasattr(walk.item, "_read_pv"):
            continue
        if not hasattr(walk.item, "_write_pv"):
            continue
        if walk.item._read_pv.pvname != walk.item._write_pv.pvname:
            walk.item._read_pv = walk.item._write_pv


def patch_functions_required_for_connection(device):
    """Patch functions required for connection. This will run the subs for all sub devices and devices.
    This is needed to ensure that the wait_for_connection method of required for connections methods are properly patched.
    """
    for event in device.event_types:
        device._run_subs(sub_type=event, value=0, timestamp=0)
    for name, dev in device.walk_subdevices(include_lazy=True):
        for event in dev.event_types:
            dev._run_subs(sub_type=event, value=0, timestamp=0)


class SocketMock:
    """Socket Mock. Used for testing"""

    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.buffer_put = []
        self.buffer_recv = [b""]
        self.is_open = False
        self.sock = None
        self.open()

    def connect(self, timeout: int = 10):
        """Mock connect method"""
        print(f"connecting to {self.host} port {self.port}")

    def _put(self, msg_bytes):
        """Mock put method"""
        self.buffer_put.append(msg_bytes)
        print(self.buffer_put)

    # pylint: disable=unused-argument
    def _recv(self, buffer_length=1024):
        """Mock receive method"""
        print(self.buffer_recv)
        if isinstance(self.buffer_recv, list):
            if len(self.buffer_recv) > 0:
                ret_val = self.buffer_recv.pop(0)
            else:
                ret_val = b""
            return ret_val
        return self.buffer_recv

    def _initialize_socket(self):
        """Mock initialize socket method"""

    def put(self, msg):
        """Mock put method"""
        return self._put(msg)

    def receive(self, buffer_length=1024):
        """Mock receive method"""
        return self._recv(buffer_length=buffer_length)

    def open(self, timeout: int = 10):
        """Mock open method"""
        self._initialize_socket()
        self.is_open = True

    def close(self):
        """Mock close method"""
        self.sock = None
        self.is_open = False

    def flush_buffer(self):
        """Mock flush buffer method"""
        self.buffer_put = []
        self.buffer_recv = ""


class MockPV:
    """
    MockPV class

    This class is used for mocking pyepics signals for testing purposes

    """

    _fmtsca = "<PV '%(pvname)s', count=%(count)i, type=%(typefull)s, access=%(access)s>"
    _fmtarr = "<PV '%(pvname)s', count=%(count)i/%(nelm)i, type=%(typefull)s, access=%(access)s>"
    _fields = (
        "pvname",
        "value",
        "char_value",
        "status",
        "ftype",
        "chid",
        "host",
        "count",
        "access",
        "write_access",
        "read_access",
        "severity",
        "timestamp",
        "posixseconds",
        "nanoseconds",
        "precision",
        "units",
        "enum_strs",
        "upper_disp_limit",
        "lower_disp_limit",
        "upper_alarm_limit",
        "lower_alarm_limit",
        "lower_warning_limit",
        "upper_warning_limit",
        "upper_ctrl_limit",
        "lower_ctrl_limit",
    )

    def __init__(
        self,
        pvname,
        *,
        _mock_pv_initial_value=0,
        callback=None,
        form="time",
        verbose=False,
        auto_monitor=None,
        count=None,
        connection_callback=None,
        connection_timeout=None,
        access_callback=None,
    ):
        self.pvname = pvname.strip()
        self.form = form.lower()
        self.verbose = verbose
        self._auto_monitor = auto_monitor
        self.ftype = None
        self.connected = True
        self.connection_timeout = connection_timeout
        self._user_max_count = count

        if self.connection_timeout is None:
            self.connection_timeout = 3
        self._args = {}.fromkeys(self._fields)
        self._args["pvname"] = self.pvname
        self._args["count"] = count
        self._args["nelm"] = -1
        self._args["type"] = "unknown"
        self._args["typefull"] = "unknown"
        self._args["access"] = "unknown"
        self._args["status"] = 0
        self.connection_callbacks = []
        self._mock_data = _mock_pv_initial_value

        if connection_callback is not None:
            self.connection_callbacks = [connection_callback]

        self.access_callbacks = []
        if access_callback is not None:
            self.access_callbacks = [access_callback]

        self.callbacks: dict[int, tuple[Callable, dict]] = {}
        self._put_complete = None
        self._put_complete_event: threading.Event | None = None
        self._monref = None  # holder of data returned from create_subscription
        self._monref_mask = None
        self._conn_started = False
        if isinstance(callback, (tuple, list)):
            for i, thiscb in enumerate(callback):
                if callable(thiscb):
                    self.callbacks[i] = (thiscb, {})
        elif callable(callback):
            self.callbacks[0] = (callback, {})

        self.chid = None
        self.context = mock.MagicMock()
        self._cache_key = (pvname, form, self.context)
        self._reference_count = 0
        for conn_cb in self.connection_callbacks:
            conn_cb(pvname=pvname, conn=True, pv=self)
        for acc_cb in self.access_callbacks:
            acc_cb(True, True, pv=self)

    @property
    def mock_data(self):
        """Get mock data"""
        return self._mock_data

    @mock_data.setter
    def mock_data(self, value):
        """Set mock data"""
        old_value = self._mock_data

        self._mock_data = value
        for callback, kw in self.callbacks.values():
            callback(value=value, old_value=old_value, obj=self, **kw)

    # pylint disable: unused-argument
    def wait_for_connection(self, timeout=None):
        """Wait for connection"""
        return self.connected

    # pylint disable: unused-argument
    def get_all_metadata_blocking(self, timeout):
        """Get all metadata blocking"""
        md = self._args.copy()
        md.pop("value", None)
        return md

    def get_all_metadata_callback(self, callback, *, timeout):
        """Get all metadata callback"""

        def get_metadata_thread(pvname):
            md = self.get_all_metadata_blocking(timeout=timeout)
            callback(pvname, md)

        get_metadata_thread(pvname=self.pvname)

    def set_severity(self, severity: AlarmSeverity):
        """Set severity for the mock PV"""
        self._args["severity"] = severity

    def set_alarm_status(self, status: AlarmStatus):
        """Set alarm status for the mock PV"""
        self._args["status"] = status

    def get_ctrlvars(self, timeout=5, warn=True):
        "get control values for variable"
        return self._args

    def get_timevars(self, timeout=5, warn=True):
        "get time values for variable"
        kwds = {
            "status": self._args["status"],
            "severity": self._args["severity"],
            "timestamp": self._args["timestamp"],
        }
        return kwds

    # pylint disable: unused-argument
    def put(
        self, value, wait=False, timeout=None, use_complete=False, callback=None, callback_data=None
    ):
        """MOCK PV, put function"""

        def put_complete():
            while True:
                if self._put_complete_event.is_set():
                    self._put_complete_event.clear()
                    callback()
                    break
                sleep(0.2)

        self.mock_data = value

        if callback is not None:
            if not self._put_complete_event:
                callback()
            else:
                threading.Thread(target=put_complete, daemon=True).start()

    # pylint: disable=unused-argument
    def add_callback(self, callback=None, index=None, run_now=False, with_ctrlvars=True, **kw):
        """Add callback"""
        if callback is None:
            logger.warning("Callback is None, cannot add callback")
            return
        if index is None:
            index = len(self.callbacks)
        self.callbacks[index] = (callback, kw)
        if run_now:
            callback(value=self.mock_data, old_value=self.mock_data, obj=self, **kw)
        return index

    # pylint: disable=unused-argument
    def get_with_metadata(
        self,
        count=None,
        as_string=False,
        as_numpy=True,
        timeout=None,
        with_ctrlvars=False,
        form=None,
        use_monitor=True,
        as_namespace=False,
    ):
        """Get MOCKPV data together with metadata"""

        return {"value": self.mock_data}

    def get(
        self,
        count=None,
        as_string=False,
        as_numpy=True,
        timeout=None,
        with_ctrlvars=False,
        use_monitor=True,
    ):
        """Get value from MOCKPV"""
        data = self.get_with_metadata(
            count=count,
            as_string=as_string,
            as_numpy=as_numpy,
            timeout=timeout,
            with_ctrlvars=with_ctrlvars,
            use_monitor=use_monitor,
        )
        return data["value"] if data is not None else None


def get_mock_scan_info(device: Device | None) -> ScanInfo:
    """
    Get a mock scan info object.
    """
    return ScanInfo(msg=fake_scan_status_msg(device=device))


def fake_scan_status_msg(device: Device | None = None) -> ScanStatusMessage:
    """
    Create a fake scan status message.

    Args:
        device: The device creating the fake scan status message.

    """
    if device is None:
        device = Device(name="mock_device")
    logger.warning(
        (
            f"Device {device.name} is not connected to a Redis server. Fetching mocked ScanStatusMessage."
        )
    )
    return ScanStatusMessage(
        metadata={},
        scan_id="mock_scan_id",
        status="closed",
        scan_number=0,
        session_id=None,
        num_points=11,
        scan_name="mock_line_scan",
        scan_type="step",
        dataset_number=0,
        scan_report_devices=["samx"],
        user_metadata={},
        readout_priority={
            "monitored": ["bpm4a", "samx"],
            "baseline": ["eyex"],
            "async": ["waveform"],
            "continuous": [],
            "on_request": ["flyer_sim"],
        },
        scan_parameters={
            "exp_time": 0,
            "frames_per_trigger": 1,
            "settling_time": 0,
            "readout_time": 0,
            "optim_trajectory": None,
            "return_to_start": True,
            "relative": True,
            "system_config": {"file_suffix": None, "file_directory": None},
        },
        request_inputs={
            "arg_bundle": ["samx", -10, 10],
            "inputs": {},
            "kwargs": {
                "steps": 11,
                "relative": True,
                "system_config": {"file_suffix": None, "file_directory": None},
            },
        },
        info={
            "readout_priority": {
                "monitored": ["bpm4a", "samx"],
                "baseline": ["eyex"],
                "async": ["waveform"],
                "continuous": [],
                "on_request": ["flyer_sim"],
            },
            "file_suffix": None,
            "file_directory": None,
            "user_metadata": {},
            "RID": "a1d86f61-191c-4460-bcd6-f33c61b395ea",
            "scan_id": "3edb8219-75a7-4791-8f86-d5ca112b771a",
            "queue_id": "0f3639ee-899f-4ad1-9e71-f40514c937ef",
            "scan_motors": ["samx"],
            "num_points": 11,
            "positions": [
                [-10.0],
                [-8.0],
                [-6.0],
                [-4.0],
                [-2.0],
                [0.0],
                [2.0],
                [4.0],
                [6.0],
                [8.0],
                [10.0],
            ],
            "scan_name": "line_scan",
            "scan_type": "step",
            "scan_number": 1,
            "dataset_number": 1,
            "exp_time": 0.5,
            "frames_per_trigger": 1,
            "settling_time": 0,
            "readout_time": 0,
            "scan_report_devices": ["samx"],
            "monitor_sync": "bec",
            "scan_parameters": {
                "exp_time": 0.5,
                "frames_per_trigger": 1,
                "settling_time": 0,
                "readout_time": 0,
                "optim_trajectory": None,
                "return_to_start": False,
                "relative": False,
                "system_config": {"file_suffix": None, "file_directory": None},
            },
            "request_inputs": {
                "arg_bundle": ["samx", -2, 2],
                "inputs": {},
                "kwargs": {
                    "steps": 10,
                    "exp_time": 0.5,
                    "relative": False,
                    "system_config": {"file_suffix": None, "file_directory": None},
                },
            },
            "file_components": ["/tmp/bec/data/S00000-00999/S00001/S00001", "h5"],
            "scan_msgs": [
                "metadata={'file_suffix': None, 'file_directory': None, 'user_metadata': {}, 'RID': 'a87334b4-51f2-420e-8efd-fa8b1faba457'} scan_type='line_scan' parameter={'args': {'samx': [-2, 2]}, 'kwargs': {'steps': 10, 'exp_time': 0.5, 'relative': False, 'system_config': {'file_suffix': None, 'file_directory': None}}} queue='primary'"
            ],
            "args": {"samx": [-2, 2]},
            "kwargs": {
                "steps": 11,
                "exp_time": 0.5,
                "relative": False,
                "system_config": {"file_suffix": None, "file_directory": None},
            },
        },
        timestamp=1737100681.694211,
    )
