import asyncio
import time
from contextlib import AbstractContextManager
from typing import Any

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

from ophyd_async.core import (
    AsyncConfigurable,
    AsyncReadable,
    SignalDatatypeT,
    SignalR,
    Table,
)


def approx_value(value: Any):
    return ApproxTable(value) if isinstance(value, Table) else pytest.approx(value)


async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
    """Assert a signal's value and compare it an expected signal.

    Parameters
    ----------
    signal:
        signal with get_value.
    value:
        The expected value from the signal.

    Notes
    -----
    Example usage::
        await assert_value(signal, value)

    """
    actual_value = await signal.get_value()
    assert approx_value(value) == actual_value


async def assert_reading(
    readable: AsyncReadable, expected_reading: dict[str, Reading]
) -> None:
    """Assert readings from readable.

    Parameters
    ----------
    readable:
        Callable with readable.read function that generate readings.

    reading:
        The expected readings from the readable.

    Notes
    -----
    Example usage::
        await assert_reading(readable, reading)

    """
    actual_reading = await readable.read()
    approx_expected_reading = {
        k: dict(v, value=approx_value(expected_reading[k]["value"]))
        for k, v in expected_reading.items()
    }
    assert actual_reading == approx_expected_reading


async def assert_configuration(
    configurable: AsyncConfigurable,
    configuration: dict[str, Reading],
) -> None:
    """Assert readings from Configurable.

    Parameters
    ----------
    configurable:
        Configurable with Configurable.read function that generate readings.

    configuration:
        The expected readings from configurable.

    Notes
    -----
    Example usage::
        await assert_configuration(configurable configuration)

    """
    actual_configuration = await configurable.read_configuration()
    approx_expected_configuration = {
        k: dict(v, value=approx_value(configuration[k]["value"]))
        for k, v in configuration.items()
    }
    assert actual_configuration == approx_expected_configuration


async def assert_describe_signal(signal: SignalR, /, **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: dict[str, list[dict]], **numbers: int):
    """Assert emitted document generated by running a Bluesky plan

    Parameters
    ----------
    Doc:
        A dictionary

    numbers:
        expected emission in kwarg from

    Notes
    -----
    Example usage::
        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


class ApproxTable:
    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):
    def __init__(self, signal: SignalR):
        self.signal = signal
        self.updates: asyncio.Queue[dict[str, Reading]] = asyncio.Queue()
        self.signal.subscribe(self.updates.put_nowait)

    async def assert_updates(self, expected_value):
        # Get an update, value and reading
        expected_type = type(expected_value)
        expected_value = approx_value(expected_value)
        update = await self.updates.get()
        value = await self.signal.get_value()
        reading = await self.signal.read()
        # Check they match what we expected
        assert value == expected_value
        assert type(value) is expected_type
        expected_reading = {
            self.signal.name: {
                "value": expected_value,
                "timestamp": pytest.approx(time.time(), rel=0.1),
                "alarm_severity": 0,
            }
        }
        assert reading == update == expected_reading

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

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