import threading
import time as ttime
import traceback

import numpy as np
from bec_lib import messages
from bec_lib.endpoints import MessageEndpoints
from bec_lib.logger import bec_logger
from ophyd import Component as Cpt
from ophyd import Device, DeviceStatus, Kind, OphydObject, PositionerBase, Signal, Staged

from ophyd_devices.interfaces.base_classes.psi_pseudo_motor_base import PSIPseudoMotorBase
from ophyd_devices.sim.sim_camera import SimCamera
from ophyd_devices.sim.sim_positioner import SimPositioner
from ophyd_devices.sim.sim_signals import SetableSignal
from ophyd_devices.utils.bec_processed_signal import BECProcessedSignal, ProcessedSignalModel
from ophyd_devices.utils.bec_signals import (
    AsyncSignal,
    DynamicSignal,
    FileEventSignal,
    PreviewSignal,
    ProgressSignal,
)

logger = bec_logger.logger


class DummyControllerDevice(Device):
    USER_ACCESS = ["controller"]


class DummyController:
    USER_ACCESS = [
        "some_var",
        "some_var_property",
        "controller_show_all",
        "_func_with_args",
        "_func_with_args_and_kwargs",
        "_func_with_kwargs",
        "_func_without_args_kwargs",
    ]
    some_var = 10
    another_var = 20

    def __init__(self) -> None:
        self._some_var_property = None
        self.connected = False

    @property
    def some_var_property(self):
        return self._some_var_property

    def on(self):
        self.connected = True

    def off(self):
        self.connected = False

    def _func_with_args(self, *args):
        return args

    def _func_with_args_and_kwargs(self, *args, **kwargs):
        return args, kwargs

    def _func_with_kwargs(self, **kwargs):
        return kwargs

    def _func_without_args_kwargs(self):
        return None

    def controller_show_all(self):
        """dummy controller show all

        Raises:
            in: _description_
            LimitError: _description_

        Returns:
            _type_: _description_
        """
        print(self.some_var)


class SimDeviceWithStatusStageUnstage(Device):
    """SimDevice with stage and unstage methods that return a status object.

    Methods resolve once the stage_thread_event or unstage_thread_event is set.
    Stop always resolves immediately.
    """

    def __init__(self, name, **kwargs):
        super().__init__(name=name, **kwargs)
        self.stage_thread = None
        self.stage_thread_event = None
        self.stopped = False

    def stage(self) -> DeviceStatus:
        """Stage the device and return a status object."""
        self.stopped = False
        if self.stage_thread is not None:
            self.stage_thread_event.set()
            self.stage_thread.join()
        self.stage_thread_event = threading.Event()
        status = DeviceStatus(self)

        def stage_device(status: DeviceStatus):

            self.stage_thread_event.wait()

            if self.stopped is True:
                exc = RuntimeError(f"Device {self.name} was stopped")
                status.set_exception(exc)
            else:
                self._staged = Staged.yes
                status.set_finished()

        self.stage_thread = threading.Thread(target=stage_device, args=(status,))
        self.stage_thread.start()
        return status

    def unstage(self) -> DeviceStatus:
        """Unstage the device and return a status object."""
        self.stopped = False
        return super().unstage()

    def stop(self, success: bool = False):
        """Stop the device and set the stopped flag."""
        self.stopped = True
        if self.stage_thread_event:
            self.stage_thread_event.set()
        super().stop(success=success)


class SynController(OphydObject):
    def on(self, timeout: int = 10):
        pass

    def off(self):
        pass


class SynFlyerLamNI(Device, PositionerBase):
    def __init__(
        self,
        *,
        name,
        readback_func=None,
        value=0,
        delay=0,
        speed=1,
        update_frequency=2,
        precision=3,
        parent=None,
        labels=None,
        kind=None,
        device_manager=None,
        **kwargs,
    ):
        if readback_func is None:

            def readback_func(x):
                return x

        self.sim_state = {}
        self._readback_func = readback_func
        self.delay = delay
        self.precision = precision
        self.tolerance = kwargs.pop("tolerance", 0.5)
        self.device_manager = device_manager

        # initialize values
        self.sim_state["readback"] = readback_func(value)
        self.sim_state["readback_ts"] = ttime.time()

        super().__init__(name=name, parent=parent, labels=labels, kind=kind, **kwargs)
        self.controller = SynController(name="SynController")

    def kickoff(self, metadata, num_pos, positions, exp_time: float = 0):
        positions = np.asarray(positions)

        def produce_data(device, metadata):
            buffer_time = 0.2
            elapsed_time = 0
            bundle = messages.BundleMessage()
            for ii in range(num_pos):
                bundle.append(
                    messages.DeviceMessage(
                        signals={
                            "flyer_samx": {"value": positions[ii, 0], "timestamp": 0},
                            "flyer_samy": {"value": positions[ii, 1], "timestamp": 0},
                        },
                        metadata={"point_id": ii, **metadata},
                    )
                )
                ttime.sleep(exp_time)
                elapsed_time += exp_time
                if elapsed_time > buffer_time:
                    elapsed_time = 0
                    device.device_manager.connector.set_and_publish(
                        MessageEndpoints.device_read(device.name), bundle
                    )
                    bundle = messages.BundleMessage()
                    device.device_manager.connector.set(
                        MessageEndpoints.device_status(device.name),
                        messages.DeviceStatusMessage(
                            device=device.name, status=1, metadata={"point_id": ii, **metadata}
                        ),
                    )
            device.device_manager.connector.send(MessageEndpoints.device_read(device.name), bundle)
            device.device_manager.connector.set(
                MessageEndpoints.device_status(device.name),
                messages.DeviceStatusMessage(
                    device=device.name, status=0, metadata={"point_id": num_pos, **metadata}
                ),
            )
            print("done")

        flyer = threading.Thread(target=produce_data, args=(self, metadata))
        flyer.start()


class SimDeviceWithSignalDelay(Device):
    """
    Simulated device with a delay. Can be used to simulate a slow rampup after an initial delay.
    One use case is to simulate a beam loss and its recovery over time.

    Usage:
    - Set the min_val and max_val to define the range of the rampup.
    - Set the delay to define how long to wait before starting the rampup.
    - Set the rampup_time to define how long the rampup should take from min_val to max_val.
    - Trigger the rampup by setting the start signal to 1. The value signal will then update over time according to the defined delay and rampup_time.
    """

    start = Cpt(SetableSignal, value=0, kind=Kind.config)
    min_val = Cpt(SetableSignal, value=0, kind=Kind.config)
    max_val = Cpt(SetableSignal, value=10, kind=Kind.config)
    delay = Cpt(SetableSignal, value=1.0, kind=Kind.config)  # Delay in seconds before rampup starts
    rampup_time = Cpt(SetableSignal, value=5.0, kind=Kind.config)  # Total rampup time in seconds
    value = Cpt(SetableSignal, value=0, kind=Kind.hinted)  # Current value signal

    SUB_VALUE = "value"
    _default_sub = SUB_VALUE

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._rampup_thread = None
        self._stop_event = threading.Event()
        self._running = False

        # Subscribe to start signal changes
        self.start.subscribe(self._on_start_changed)
        self.value.subscribe(self._on_value_changed)

    def _on_start_changed(self, value, **kwargs):
        """Callback when start signal changes."""
        if value == 1 and not self._running:
            self._start_rampup()

    def _on_value_changed(self, value, **kwargs):
        self._run_subs(self.SUB_VALUE, value=value, **kwargs)

    def _start_rampup(self):
        """Start the rampup simulation in a background thread."""
        if self._running:
            return

        self._stop_event.clear()
        self._running = True
        self._rampup_thread = threading.Thread(target=self._rampup_worker, daemon=True)
        self._rampup_thread.start()

        # Reset start signal to 0 after launching
        self.start.put(0)

    def _stop_rampup(self):
        """Stop the rampup simulation."""
        if not self._running:
            return

        self._stop_event.set()
        if self._rampup_thread is not None:
            self._rampup_thread.join(timeout=1.0)
        self._running = False
        self.value.put(self.min_val.get())

    def _rampup_worker(self):
        """Worker thread that handles delay and rampup."""
        try:
            min_value = self.min_val.get()
            max_value = self.max_val.get()
            delay_time = max(self.delay.get(), 0)  # Ensure non-negative delay
            rampup_duration = max(self.rampup_time.get(), 0)  # Ensure non-negative rampup duration

            if rampup_duration == 0:
                # If the rampup duration is zero, set value to max immediately after delay
                if self._stop_event.wait(timeout=delay_time):
                    return  # Stopped during delay
                self.value.put(max_value)
                return

            # Wait for delay period
            if self._stop_event.wait(timeout=delay_time):
                return  # Stopped during delay

            # Set initial value to min_val
            self.value.put(min_value)

            # Rampup phase
            start_time = ttime.time()
            update_interval = 0.1  # Update every 100ms

            while not self._stop_event.is_set():
                elapsed = ttime.time() - start_time

                if elapsed >= rampup_duration:
                    # Rampup complete
                    self.value.put(max_value)
                    break

                # Calculate current value based on linear interpolation
                progress = elapsed / rampup_duration
                current_value = min_value + (max_value - min_value) * progress
                self.value.put(current_value)

                # Wait for next update or stop signal
                if self._stop_event.wait(timeout=update_interval):
                    break

        except Exception as exc:
            logger.error(f"Error in rampup worker for {self.name}: {exc}")
        finally:
            self._running = False

    def stop(self, success=False):
        """Stop the device and any ongoing rampup."""
        self._stop_rampup()
        super().stop(success=success)


class SimPositionerWithCommFailure(SimPositioner):
    fails = Cpt(SetableSignal, value=0)

    def move(self, value: float, **kwargs) -> DeviceStatus:
        if self.fails.get() == 1:
            raise RuntimeError("Communication failure")
        if self.fails.get() == 2:
            while not self._stopped:
                ttime.sleep(1)
            status = DeviceStatus(self)
            status.set_exception(RuntimeError("Communication failure"))
        return super().move(value, **kwargs)


class SimPositionerWithDescribeFailure(SimPositioner):
    """
    Simulated positioner that raises a RuntimeError if the 'e2e_test_<name>_fail' endpoint is set to 1 using
    a DeviceStatusMessage. This is used to test the behavior of the device when the describe method fails.
    """

    def __init__(self, *args, device_manager=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.device_manager = device_manager

    def _get_fail_state(self):
        msg = self.device_manager.connector.get(f"e2e_test_{self.name}_fail")
        if not msg:
            return 0
        return msg.status

    def describe(self):
        if self._get_fail_state() == 1:
            raise RuntimeError("Communication failure")
        return super().describe()


class SimPositionerWithController(SimPositioner):
    USER_ACCESS = ["sim", "readback", "dummy_controller", "registered_proxies"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dummy_controller = DummyController()


class SimCameraWithStageStatus(SimCamera):
    """Simulated camera which returns a status object when staged.

    Note: This is a minimum implementation for test purposes without refactoring
    the super().stage() method. In theory, the super().stage() method should take
    into account a thread event to stop the staging process if the device is stopped.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.raise_on_stage = False
        if "raise_on_stage" in kwargs:
            self.raise_on_stage = kwargs.pop("raise_on_stage")

    def stage(self):
        status = DeviceStatus(self)

        def _stage_device(obj, status: DeviceStatus):
            """Start thread to stage the device"""
            try:
                logger.info(f"Staging device {obj.name} with status object")
                super(SimCamera, obj).stage()
                if obj.raise_on_stage is True:
                    raise RuntimeError(f"Error during staging in {obj.name}")
            # pylint: disable=broad-except
            except Exception as exc:
                content = traceback.format_exc()
                logger.warning(f"Error in staging of {obj.name}; Traceback: {content}")
                status.set_exception(exc=exc)
            else:
                status.set_finished()

        thread = threading.Thread(target=_stage_device, args=(self, status), daemon=True)
        thread.start()
        return status

    def unstage(self):
        status = DeviceStatus(self)

        def _unstage_device(obj, status: DeviceStatus):
            """Start thread to stage the device"""
            try:
                logger.info(f"Unstaging device {obj.name} with status object")
                super(SimCamera, obj).unstage()
                if obj.raise_on_stage is True:
                    raise RuntimeError(f"Error during unstaging in {obj.name}")
            # pylint: disable=broad-except
            except Exception as exc:
                content = traceback.format_exc()
                logger.warning(f"Error in unstaging of {obj.name}; Traceback: {content}")
                status.set_exception(exc=exc)
            else:
                status.set_finished()

        thread = threading.Thread(target=_unstage_device, args=(self, status), daemon=True)
        thread.start()
        return status


class SimCameraWithPSIComponents(SimCamera):
    """Test Device for PSIComponents"""

    preview_2d = Cpt(
        PreviewSignal, ndim=2, doc="2D preview signal", num_rotation_90=2, transpose=True
    )
    preview_1d = Cpt(PreviewSignal, ndim=1, doc="1D preview signal")
    file_event = Cpt(FileEventSignal, doc="File event signal")
    progress = Cpt(ProgressSignal, doc="Progress signal")
    dynamic_signal = Cpt(
        DynamicSignal, doc="Dynamic signals", signals=["dyn_signal1", "dyn_signal2"], max_size=1000
    )
    async_signal = Cpt(AsyncSignal, name="async_signal", ndim=1, doc="Async signal", max_size=1000)
    # TODO Handling of AsyncComponents is postponed, issue #104 is created

    # Define signals for the async components
    # signal_dict = {
    #     "signal1": {"kind": Kind.hinted, "doc": "Signal 1"},
    #     "signal2": {"kind": Kind.normal, "doc": "Signal 2"},
    #     "signal3": {"kind": Kind.config, "doc": "Signal 3"},
    #     "signal4": {"kind": Kind.omitted, "doc": "Signal 4"},
    # }

    # async_1d = Async1DComponent(doc="Async 1D signal", signal_def=signal_dict)
    # async_2d = Async2DComponent(doc="Async 2D signal", signal_def=signal_dict)

    def __init__(self, name: str, scan_info=None, device_manager=None, **kwargs):
        super().__init__(name=name, scan_info=scan_info, device_manager=device_manager, **kwargs)
        self._triggers_received = 0

    def on_stage(self):
        """Stage device"""
        self.file_path = self.file_utils.get_full_path(
            scan_status_msg=self.scan_info.msg, name=self.name
        )
        self.frames.set(
            self.scan_info.msg.num_points * self.scan_info.msg.scan_parameters["frames_per_trigger"]
        ).wait()
        self.exp_time.set(self.scan_info.msg.scan_parameters["exp_time"]).wait()
        self.burst.set(self.scan_info.msg.scan_parameters["frames_per_trigger"]).wait()
        # Always emit a file event
        msg = messages.FileMessage(
            file_path=self.file_path, done=False, successful=False, device_name=self.name
        )
        self.file_event.put(file_path=self.file_path, done=False, successful=False)
        self.file_event.set(msg).wait()
        self._triggers_received = 0

    def on_trigger(self):
        """Trigger device"""
        self._triggers_received += 1

        def trigger_cam():
            for _ in range(self.burst.get()):
                data = self.image.get()
                self.preview_2d.put(data)
                # sum array in one dimension
                self.preview_1d.put(np.sum(data, 1))
                self.dynamic_signal.put(
                    {"preview_2d": {"value": data}, "preview_1d": {"value": np.sum(data, 1)}}
                )
                progress = {
                    "value": self._triggers_received,
                    "max_value": self.scan_info.msg.num_points,
                    "done": (self._triggers_received == self.scan_info.msg.num_points),
                }
                progress = messages.ProgressMessage(**progress)
                self.progress.put(progress)
                self.async_signal.put(
                    np.sum(data, 1), async_update={"type": "add", "max_shape": [None]}
                )

        status = self.task_handler.submit_task(trigger_cam)
        return status

    def on_unstage(self):
        """Unstage device"""
        self._triggers_received = 0

    def on_complete(self):
        """Complete device"""

        def complete_cam():
            """Complete the camera acquisition."""
            msg = messages.FileMessage(
                file_path=self.file_path if self.file_path else "",
                done=True,
                successful=True,
                device_name=self.name,
            )
            self.file_event.set(msg).wait()
            progress = {
                "value": self._triggers_received,
                "max_value": self.scan_info.msg.num_points,
                "done": True,
            }
            progress = messages.ProgressMessage(**progress)
            self.progress.set(progress).wait()

        status = self.task_handler.submit_task(complete_cam)
        return status


class VirtualSlitCenter(PSIPseudoMotorBase):
    """
    Alternative implementation of a Virtual Slit Center based on two sub-positioners.
    The positioners are sub-devices called left_edge and right_edge, and are used
    to calculate the center position of the slit.

    Args:
        name (str): The name of the pseudo motor device.
        device_manager (DeviceManagerBase): The device manager instance to fetch the positioner devices from
    """

    left_edge = Cpt(SimPositioner, name="left_edge", kind=Kind.normal)
    right_edge = Cpt(SimPositioner, name="right_edge", kind=Kind.normal)

    def __init__(self, name, device_manager, **kwargs):
        super().__init__(name, device_manager, **kwargs)
        positioners = {"left_edge": self.left_edge, "right_edge": self.right_edge}
        self.set_positioner_objects(positioners)

    def forward_calculation(self, left_edge: Signal, right_edge: Signal) -> float:
        """
        Forward calculation to compute the center position of the slit based on the positions of the left and right edges.

        Args:
            left_edge (Signal): The signal representing the position of the left edge positioner.
            right_edge (Signal): The signal representing the position of the right edge positioner.
        """
        return float((left_edge.get() + right_edge.get()) / 2)

    def inverse_calculation(
        self, position, left_edge: Signal, right_edge: Signal
    ) -> dict[str, float]:
        """
        Inverse calculation to compute the setpoints for the left and right edge positioners based on the desired center position.

        Args:
            position (float): The desired center position of the slit.
            left_edge (Signal): The signal representing the position of the left edge positioner.
            right_edge (Signal): The signal representing the position of the right edge positioner.
        Returns:
            dict[str, float]: A dictionary containing the setpoints for the left and right edge positioners.
        """
        left_pos = left_edge.root.readback.get()
        right_pos = right_edge.root.readback.get()
        width = right_pos - left_pos
        new_right_pos = position + width / 2
        new_left_pos = position - width / 2
        return {"left_edge": new_left_pos, "right_edge": new_right_pos}

    def motors_are_moving(self, left_edge: Signal, right_edge: Signal) -> int:
        """
        Check if either the left or right edge positioners are currently moving.

        Args:
            left_edge (Signal): The signal representing the position of the left edge positioner.
            right_edge (Signal): The signal representing the position of the right edge positioner.
        Returns:
            int: 1 if either motor is moving, 0 otherwise.
        """
        left_moving = left_edge.get()
        right_moving = right_edge.get()
        return int(left_moving or right_moving)


class BECPseudoSignal(BECProcessedSignal):
    """
    Example of a pseudo signal implementation based on two signals available in the device manager.
    The value of the signal is calculated based on the compute method.

    Args:
        name (str): The name of the pseudo signal.
        signal_1 (str): The name of the first signal in the device manager.
        signal_2 (str): The name of the second signal in the device manager.
        device_manager (DeviceManagerBase): The device manager instance to fetch the signals from.
    """

    def __init__(self, name, signal_1: str, signal_2: str, device_manager=None, **kwargs):
        super().__init__(name, model_config=None, device_manager=device_manager, **kwargs)
        signal_1 = self.get_device_object_from_bec(
            object_name=signal_1, signal_name=self.name, device_manager=device_manager
        )
        signal_2 = self.get_device_object_from_bec(
            object_name=signal_2, signal_name=self.name, device_manager=device_manager
        )
        self.set_compute_method(self.compute, signal_1=signal_1, signal_2=signal_2)

    def compute(self, signal_1: Signal, signal_2: Signal) -> float:
        """Compute the value of the pseudo signal based on the values of signal_1 and signal_2."""
        return float(signal_1.get() + signal_2.get())


if __name__ == "__main__":
    ring = SimDeviceWithSignalDelay(name="ring")
    ring.min_val.put(300)
    ring.max_val.put(400)
    ring.value.put(400)
    ring.delay.put(5)
    ring.rampup_time.put(20)
    ring.subscribe(lambda *args, **kwargs: print(f"Value changed to {kwargs['value']}"))
    ring.start.put(1)
    while ring._running:
        ttime.sleep(0.5)
    print("Rampup complete")
