import asyncio
from collections.abc import Mapping, Sequence
from contextlib import AbstractContextManager
from typing import Any, cast
from unittest.mock import Mock, call

import pytest
from bluesky.protocols import Reading
from event_model import DataKey

from ophyd_async.core import (
    AsyncConfigurable,
    AsyncReadable,
    Device,
    Signal,
    SignalDatatypeT,
    SignalR,
    Table,
    WatchableAsyncStatus,
    Watcher,
    get_mock,
)

from ._utils import T


def partial_reading(val: Any) -> Mapping[str, Any]:
    """Helper function for building expected reading or configuration mapping.

    :param val: Value to be wrapped in mapping with "value" as the key.
    :return: The mapping that has wrapped the val with key "value".
    """
    return {"value": val}


def approx_value(value: Any):
    """Allow any value to be compared to another in tests.

    This is needed because numpy arrays give a numpy array back when compared,
    not a bool. This means that you can't ``assert array1==array2``. Numpy
    arrays can be wrapped with `pytest.approx`, but this doesn't work for
    `Table` instances: in this case we use `ApproxTable`.
    """
    return ApproxTable(value) if isinstance(value, Table) else pytest.approx(value)


async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
    """Assert that a Signal has the given value.

    :param signal: Signal with get_value.
    :param value: The expected value from the signal.
    """
    actual_value = await signal.get_value()
    assert approx_value(value) == actual_value


async def assert_reading(
    readable: AsyncReadable,
    expected_reading: Mapping[str, Mapping[str, Any]],
    full_match: bool = True,
) -> None:
    """Assert that a readable Device has the given reading.

    :param readable: Device with an async ``read()`` method to get the reading from.
    :param expected_reading: The expected reading from the readable.
    :param full_match: if expected_reading keys set is same as actual keys set.
        true: exact match
        false: expected_reading keys is subset of actual reading keys
    """
    actual_reading = await readable.read()
    _assert_readings_approx_equal(expected_reading, actual_reading, full_match)


def _approx_reading(expected: Mapping[str, Any], actual: Reading) -> Reading:
    ret = dict(
        expected,
        value=approx_value(expected["value"]),
        timestamp=pytest.approx(expected["timestamp"], rel=0.1)
        if "timestamp" in expected
        else actual["timestamp"],
    )
    if "alarm_severity" in actual and "alarm_severity" not in expected:
        ret["alarm_severity"] = actual["alarm_severity"]
    return cast(Reading, ret)


def _assert_readings_approx_equal(
    expected: Mapping[str, Mapping[str, Any]],
    actual: Mapping[str, Reading],
    full_match: bool = True,
):
    # expand the expected keys to include actual if we allow partial matches
    if not full_match:
        expected = dict(actual, **expected)
    # now make them approximate if they are in actual so we get a nicer diff
    approx_expected = {
        k: _approx_reading(v, actual[k]) if k in actual else v
        for k, v in expected.items()
    }
    # now we can compare them
    assert actual == approx_expected


async def assert_configuration(
    configurable: AsyncConfigurable,
    expected_configuration: Mapping[str, Mapping[str, Any]],
    full_match: bool = True,
) -> None:
    """Assert that a configurable Device has the given configuration.

    :param configurable:
        Device with an async ``read_configuration()`` method to get the
        configuration from.
    :param expected_configuration: The expected configuration from the configurable.
    :param full_match: if expected_reading keys set is same as actual keys set.
        true: exact match
        false: expected_reading keys is subset of actual reading keys
    """
    actual_configuration = await configurable.read_configuration()
    _assert_readings_approx_equal(
        expected_configuration, actual_configuration, full_match
    )


async def assert_describe_signal(signal: SignalR, /, **metadata):
    """Assert the describe of a signal matches the expected metadata.

    :param signal: The signal to describe.
    :param metadata: The expected metadata.
    """
    actual_describe = await signal.describe()
    assert list(actual_describe) == [signal.name]
    (actual_datakey,) = actual_describe.values()
    expected_datakey = DataKey(source=signal.source, **metadata)
    assert actual_datakey == expected_datakey


def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int):
    """Assert emitted document generated by running a Bluesky plan.

    :param docs: A mapping of document type -> list of documents that have been emitted.
    :param numbers: The number of each document type expected.

    :example:
    ```python
    docs = defaultdict(list)
    RE.subscribe(lambda name, doc: docs[name].append(doc))
    RE(my_plan())
    assert_emitted(docs, start=1, descriptor=1, event=1, stop=1)
    ```
    """
    assert list(docs) == list(numbers)
    actual_numbers = {name: len(d) for name, d in docs.items()}
    assert actual_numbers == numbers


def assert_has_calls(device: Device | Signal, calls: Sequence[Any], reset_after=True):
    """Check that the device connected in mock mode has seen the supplied calls.

    :param device: Device or Signal connected in mock mode to check.
    :param calls: Expected sequence of mock calls.
    :param reset_after:
        Whether to reset the mock after assertion so subsequent calls don't see
        the same calls again.

    :example:
    ```python device.trigger() assert_has_calls(
        device, [
            call.device.num_frames.put(1),
            call.device.start_writing.put(None),
        ],
    )
    ```
    """
    mock = get_mock(device)
    assert list(mock.mock_calls) == list(calls)
    if reset_after:
        mock.reset_mock()


class ApproxTable:
    """For approximating two tables are equivalent.

    :param expected: The expected table.
    :param rel: The relative tolerance.
    :param abs: The absolute tolerance.
    :param nan_ok: Whether NaNs are allowed.
    """

    def __init__(self, expected: Table, rel=None, abs=None, nan_ok: bool = False):
        self.expected = expected
        self.rel = rel
        self.abs = abs
        self.nan_ok = nan_ok

    def __eq__(self, value):
        approx_fields = {
            k: pytest.approx(v, self.rel, self.abs, self.nan_ok)
            for k, v in self.expected
        }
        expected = type(self.expected).model_construct(**approx_fields)  # type: ignore
        return expected == value


class MonitorQueue(AbstractContextManager):
    """Monitors a `Signal` and stores its updates."""

    def __init__(self, signal: SignalR):
        self.signal = signal
        self.updates: asyncio.Queue[dict[str, Reading]] = asyncio.Queue()

    async def assert_updates(self, expected_value):
        # Get an update, value and reading
        update = await asyncio.wait_for(self.updates.get(), timeout=5)
        await assert_value(self.signal, expected_value)
        expected_reading = {
            self.signal.name: {
                "value": expected_value,
            }
        }
        await assert_reading(self.signal, expected_reading)
        _assert_readings_approx_equal(expected_reading, update)

    def __enter__(self):
        self.signal.subscribe_reading(self.updates.put_nowait)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.signal.clear_sub(self.updates.put_nowait)


class StatusWatcher(Watcher[T]):
    """Watches an `AsyncStatus`, storing the calls within."""

    def __init__(self, status: WatchableAsyncStatus) -> None:
        self._event = asyncio.Event()
        self.mock = Mock()
        """Mock that stores watcher updates from the status."""
        status.watch(self)

    def __call__(
        self,
        current: T | None = None,
        initial: T | None = None,
        target: T | None = None,
        name: str | None = None,
        unit: str | None = None,
        precision: int | None = None,
        fraction: float | None = None,
        time_elapsed: float | None = None,
        time_remaining: float | None = None,
    ) -> Any:
        self.mock(
            current=current,
            initial=initial,
            target=target,
            name=name,
            unit=unit,
            precision=precision,
            fraction=fraction,
            time_elapsed=time_elapsed,
            time_remaining=time_remaining,
        )
        self._event.set()

    async def wait_for_call(
        self,
        current: T | None = None,
        initial: T | None = None,
        target: T | None = None,
        name: str | None = None,
        unit: str | None = None,
        precision: int | None = None,
        fraction: float | None = None,
        # Any so we can use pytest.approx
        time_elapsed: float | Any = None,
        time_remaining: float | Any = None,
    ):
        await asyncio.wait_for(self._event.wait(), timeout=1)
        assert self.mock.call_count == 1
        assert self.mock.call_args == call(
            current=current,
            initial=initial,
            target=target,
            name=name,
            unit=unit,
            precision=precision,
            fraction=fraction,
            time_elapsed=time_elapsed,
            time_remaining=time_remaining,
        )
        self.mock.reset_mock()
        self._event.clear()
