import asyncio
import re
import sys
from asyncio import CancelledError

import pytest

from ophyd_async.core import (
    DEFAULT_TIMEOUT,
    Device,
    gather_dict,
    init_devices,
    SoftSignalBackend,
    NotConnectedError,
    SignalRW,
    AsyncStatus,
    get_mock_put,
    set_mock_value,
    soft_signal_rw,
)
from ophyd_async.epics.core import epics_signal_rw
from ophyd_async.epics.motor import Motor


class ValueErrorBackend(SoftSignalBackend):
    def __init__(self, exc_text=""):
        self.exc_text = exc_text
        super().__init__(datatype=int)

    async def connect(self, timeout: float = DEFAULT_TIMEOUT):
        raise ValueError(self.exc_text)


class WorkingDummyChildDevice(Device):
    def __init__(self, name: str = "working_dummy_child_device") -> None:
        self.working_signal = soft_signal_rw(int)
        super().__init__(name=name)


class TimeoutDummyChildDeviceCA(Device):
    def __init__(self, name: str = "timeout_dummy_child_device_ca") -> None:
        self.timeout_signal = epics_signal_rw(int, "ca://A_NON_EXISTENT_SIGNAL")
        super().__init__(name=name)


class TimeoutDummyChildDevicePVA(Device):
    def __init__(self, name: str = "timeout_dummy_child_device_pva") -> None:
        self.timeout_signal = epics_signal_rw(int, "pva://A_NON_EXISTENT_SIGNAL")
        super().__init__(name=name)


class ValueErrorDummyChildDevice(Device):
    def __init__(
        self, name: str = "value_error_dummy_child_device", exc_text=""
    ) -> None:
        self.value_error_signal = SignalRW(ValueErrorBackend(exc_text=exc_text))
        super().__init__(name=name)


class DummyDeviceOneWorkingOneTimeout(Device):
    def __init__(self, name: str = "dummy_device_one_working_one_timeout") -> None:
        self.working_child_device = WorkingDummyChildDevice()
        self.timeout_child_device = TimeoutDummyChildDeviceCA()
        super().__init__(name=name)


ONE_WORKING_ONE_TIMEOUT_OUTPUT = NotConnectedError(
    {
        "timeout_child_device": NotConnectedError(
            {"timeout_signal": NotConnectedError("ca://A_NON_EXISTENT_SIGNAL")}
        )
    }
)


class DummyDeviceTwoWorkingTwoTimeOutTwoValueError(Device):
    def __init__(
        self,
        name: str = "dummy_device_two_working_one_timeout_two_value_error",
    ) -> None:
        self.working_child_device1 = WorkingDummyChildDevice()
        self.working_child_device2 = WorkingDummyChildDevice()
        self.timeout_child_device_ca = TimeoutDummyChildDeviceCA()
        self.timeout_child_device_pva = TimeoutDummyChildDevicePVA()
        self.value_error_child_device1 = ValueErrorDummyChildDevice(
            exc_text="Some ValueError text"
        )
        self.value_error_child_device2 = ValueErrorDummyChildDevice()
        super().__init__(name=name)


TWO_WORKING_TWO_TIMEOUT_TWO_VALUE_ERROR_OUTPUT = NotConnectedError(
    {
        "timeout_child_device_ca": NotConnectedError(
            {
                "timeout_signal": NotConnectedError("ca://A_NON_EXISTENT_SIGNAL"),
            }
        ),
        "timeout_child_device_pva": NotConnectedError(
            {"timeout_signal": NotConnectedError("pva://A_NON_EXISTENT_SIGNAL")}
        ),
        "value_error_child_device1": NotConnectedError(
            {"value_error_signal": ValueError("Some ValueError text")}
        ),
        "value_error_child_device2": NotConnectedError(
            {
                "value_error_signal": ValueError(),
            }
        ),
    }
)


class DummyDeviceCombiningTopLevelSignalAndSubDevice(Device):
    def __init__(
        self, name: str = "dummy_device_combining_top_level_signal_and_sub_device"
    ) -> None:
        self.timeout_signal = epics_signal_rw(int, "ca://A_NON_EXISTENT_SIGNAL")
        self.sub_device = ValueErrorDummyChildDevice(exc_text="Some ValueError text")
        super().__init__(name=name)


async def test_error_handling_connection_timeout(caplog):
    caplog.set_level(10)

    dummy_device_one_working_one_timeout = DummyDeviceOneWorkingOneTimeout()

    # This should work since the error is a connection timeout
    with pytest.raises(NotConnectedError) as exc:
        await dummy_device_one_working_one_timeout.connect(timeout=0.01)

    assert str(exc.value) == str(ONE_WORKING_ONE_TIMEOUT_OUTPUT)

    logs = caplog.get_records("call")

    # See https://github.com/bluesky/ophyd-async/issues/519
    # assert len(logs) == 3

    assert "signal ca://A_NON_EXISTENT_SIGNAL timed out" == logs[-1].message
    assert logs[-1].levelname == "DEBUG"


async def test_error_handling_value_errors(caplog):
    """Checks that NotConnectedError is aggregated correctly across Devices."""

    caplog.set_level(10)

    dummy_device_two_working_one_timeout_two_value_error = (
        DummyDeviceTwoWorkingTwoTimeOutTwoValueError("dsf")
    )

    # This should fail since the error is a ValueError
    with pytest.raises(NotConnectedError) as exc:
        await dummy_device_two_working_one_timeout_two_value_error.connect(timeout=0.01)

    assert str(exc.value) == str(TWO_WORKING_TWO_TIMEOUT_TWO_VALUE_ERROR_OUTPUT)

    logs = caplog.get_records("call")
    logs = [
        log
        for log in logs
        if "ophyd_async" in log.pathname and "_signal" not in log.pathname
    ]
    assert len(logs) == 4

    for i in range(0, 2):
        assert (
            "device `value_error_signal` raised unexpected exception ValueError"
            == logs[i].message
        )
        assert logs[i].levelname == "ERROR"

    assert logs[2].levelname == "DEBUG"
    assert logs[3].levelname == "DEBUG"
    # These messages could come in any order
    messages = [logs[idx].message for idx in (2, 3)]
    for protocol in ("pva", "ca"):
        assert f"signal {protocol}://A_NON_EXISTENT_SIGNAL timed out" in messages


class BadDatatypeDevice(Device):
    def __init__(self, name: str = "") -> None:
        self.sig = epics_signal_rw(object, "", "")
        super().__init__(name)


async def test_error_handling_init_devices_mock():
    with pytest.raises(NotConnectedError) as exc:
        async with init_devices(mock=True):
            device = BadDatatypeDevice()
            device2 = BadDatatypeDevice()
    expected_output = NotConnectedError(
        {
            "device": NotConnectedError(
                {"sig": TypeError(f"Can't make converter for {object}")}
            ),
            "device2": NotConnectedError(
                {"sig": TypeError(f"Can't make converter for {object}")}
            ),
        }
    )
    assert str(expected_output) == str(exc.value)


def test_introspecting_sub_errors():
    sub_error1 = NotConnectedError("bad")
    assert sub_error1.sub_errors == {}
    sub_error2 = ValueError("very bad")
    error = NotConnectedError({"child1": sub_error1, "child2": sub_error2})
    assert error.sub_errors == {"child1": sub_error1, "child2": sub_error2}


async def test_error_handling_init_devices(caplog):
    caplog.set_level(10)
    with pytest.raises(NotConnectedError) as exc:
        # flake8: noqa
        async with init_devices(timeout=0.1):
            dummy_device_two_working_one_timeout_two_value_error = (
                DummyDeviceTwoWorkingTwoTimeOutTwoValueError()
            )
            dummy_device_one_working_one_timeout = DummyDeviceOneWorkingOneTimeout()

    expected_output = NotConnectedError(
        {
            "dummy_device_two_working_one_timeout_two_value_error": (
                TWO_WORKING_TWO_TIMEOUT_TWO_VALUE_ERROR_OUTPUT
            ),
            "dummy_device_one_working_one_timeout": ONE_WORKING_ONE_TIMEOUT_OUTPUT,
        }
    )
    assert str(expected_output) == str(exc.value)

    logs = caplog.get_records("call")
    logs = [
        log
        for log in logs
        if "ophyd_async" in log.pathname and "_signal" not in log.pathname
    ]
    assert len(logs) == 5
    assert (
        logs[0].message
        == logs[1].message
        == "device `value_error_signal` raised unexpected exception ValueError"
    )
    assert logs[0].levelname == logs[1].levelname == "ERROR"

    assert logs[2].levelname == "DEBUG"
    assert logs[3].levelname == "DEBUG"
    # These messages could come in any order
    messages = [logs[idx].message for idx in (2, 3)]
    for protocol in ("pva", "ca"):
        assert f"signal {protocol}://A_NON_EXISTENT_SIGNAL timed out" in messages


def test_not_connected_error_output():
    assert str(TWO_WORKING_TWO_TIMEOUT_TWO_VALUE_ERROR_OUTPUT) == (
        "\ntimeout_child_device_ca: NotConnectedError:\n"
        "    timeout_signal: NotConnectedError: ca://A_NON_EXISTENT_SIGNAL\n"
        "timeout_child_device_pva: NotConnectedError:\n"
        "    timeout_signal: NotConnectedError: pva://A_NON_EXISTENT_SIGNAL\n"
        "value_error_child_device1: NotConnectedError:\n"
        "    value_error_signal: ValueError: Some ValueError text\n"
        "value_error_child_device2: NotConnectedError:\n"
        "    value_error_signal: ValueError\n"
    )


async def test_combining_top_level_signal_and_child_device():
    dummy_device1 = DummyDeviceCombiningTopLevelSignalAndSubDevice()
    with pytest.raises(NotConnectedError) as exc:
        await dummy_device1.connect(timeout=0.01)
    assert str(exc.value) == (
        "\ntimeout_signal: NotConnectedError: ca://A_NON_EXISTENT_SIGNAL\n"
        "sub_device: NotConnectedError:\n"
        "    value_error_signal: ValueError: Some ValueError text\n"
    )

    with pytest.raises(NotConnectedError) as exc:
        async with init_devices(timeout=0.1):
            dummy_device2 = DummyDeviceCombiningTopLevelSignalAndSubDevice()
    assert str(exc.value) == (
        "\ndummy_device2: NotConnectedError:\n"
        "    timeout_signal: NotConnectedError: ca://A_NON_EXISTENT_SIGNAL\n"
        "    sub_device: NotConnectedError:\n"
        "        value_error_signal: ValueError: Some ValueError text\n"
    )


async def test_format_error_string_input():
    with pytest.raises(
        RuntimeError,
        match=("Unexpected type `<class 'int'>` expected `str` or `dict`"),
    ):
        not_connected = NotConnectedError(123)
        str(not_connected)

    with pytest.raises(
        RuntimeError, match=("Unexpected type `<class 'int'>`, expected an Exception")
    ):
        not_connected = NotConnectedError({"test": 123})
        str(not_connected)


def test_core_notconnected_emits_deprecation_warning():
    with pytest.deprecated_call():
        from ophyd_async.core import NotConnected  # noqa: F401


async def _async_value(val):
    """Helper function to create an awaitable that returns a value."""
    return val


async def test_gather_dict_empty():
    """Test gather_dict with an empty dictionary."""
    result = await gather_dict({})
    assert result == {}


async def test_gather_dict_duplicate_resolved_keys():
    """Test gather_dict when different awaitable keys resolve to the same value."""
    # Two coroutine objects that resolve to the same key - second should win
    result = await gather_dict(
        {
            _async_value("key"): "first",
            _async_value("key"): "second",
        }
    )
    assert result == {"key": "second"}


@pytest.mark.parametrize("keys_awaitable", ("", "b", "abc"))
@pytest.mark.parametrize("values_awaitable", ([], [2], [1, 2, 3]))
async def test_gather_dict(keys_awaitable, values_awaitable):
    """Test gather_dict with plain values, awaitable keys, and awaitable values."""
    expected = {"a": 1, "b": 2, "c": 3}

    input_dict = {}
    for k, v in expected.items():
        key = _async_value(k) if k in keys_awaitable else k
        value = _async_value(v) if v in values_awaitable else v
        input_dict[key] = value
    result = await gather_dict(input_dict)
    assert result == expected
    assert list(result.keys()) == list(expected.keys())


# Cancellation propagation is broken in asyncio.gather on Python 3.10, so the exception
# with the message is not reachable via the __cause__ chain however you will still get
# the message in the stack trace
@pytest.mark.skipif(
    sys.version_info < (3, 11), reason="https://github.com/python/cpython/issues/112534"
)
@pytest.mark.parametrize(
    "set_to_delay, expected_message, expected_results",
    # Test cases for different things taking too long
    [
        # plain awaitable (currently not handled)
        [
            "x",
            None,
            None,
        ],
        # Device with a custom set() method
        ["y", "CancelledError while awaiting .* on my_device", None],
        # Task wrapping an async function (currently not handled)
        ["z", None, None],
        # Setting a Motor on a Device
        ["omega", "BL03I-MO-OMEGA", None],
        # Everything completes successfully in a timely fashion
        [None, None, [1, None, 1, None]],
    ],
)
async def test_cancelled_error_message_for_gather_is_populated_on_timeout(
    set_to_delay: str, expected_message: str, expected_results: list
):
    async def async_func(name: str, value: int) -> int:
        if set_to_delay == name:
            await asyncio.sleep(20)
        return value

    plain_awaitable = async_func("x", 1)

    class GenericDevice(Device):
        def __init__(self, name):
            super().__init__(name)

        @AsyncStatus.wrap
        async def set(self, value: int):
            await async_func("y", value)

    class MotorDevice(Device):
        def __init__(self, name):
            self.omega = Motor("BL03I-MO-OMEGA")
            super().__init__(name)

        async def set_omega(self, value: int, *args, **kwargs):
            await async_func("omega", value)

    async_status_device = GenericDevice("my_device")
    task_wrapped = asyncio.create_task(async_func("z", 1), name="set z")
    motor_device = MotorDevice("my_motor_device")
    await motor_device.connect(mock=True)
    _patch_motor(motor_device.omega)
    get_mock_put(motor_device.omega.user_setpoint).side_effect = motor_device.set_omega

    gather_awaitable = asyncio.gather(
        plain_awaitable,
        async_status_device.set(1),
        task_wrapped,
        motor_device.omega.set(1),
    )
    if set_to_delay:
        with pytest.raises(asyncio.TimeoutError) as exc_info:
            await asyncio.wait_for(gather_awaitable, timeout=0.3)
        cause = exc_info.value.__cause__
        assert isinstance(cause, CancelledError)
        assert _contains_message(cause, expected_message)
    else:
        results = await gather_awaitable
        assert results == expected_results


def _contains_message(e: BaseException, expected_message: str) -> bool:
    return (
        expected_message is None
        or (e.args and re.search(expected_message, e.args[0]))
        or (e.__cause__ and _contains_message(e.__cause__, expected_message))
    )


def _patch_motor(motor: Motor, initial_position=0):
    set_mock_value(motor.user_setpoint, initial_position)
    set_mock_value(motor.user_readback, initial_position)
    set_mock_value(motor.deadband, 0.001)
    set_mock_value(motor.motor_done_move, 1)
    set_mock_value(motor.velocity, 3)
    set_mock_value(motor.max_velocity, 5)
    set_mock_value(motor.high_limit_travel, float("inf"))
    set_mock_value(motor.low_limit_travel, float("-inf"))
