from concurrent.futures import Executor
from threading import Event
from unittest.mock import ANY, AsyncMock, MagicMock, PropertyMock, call, patch

import pytest
from blueapi.client.event_bus import BlueskyStreamingError
from blueapi.core import BlueskyContext
from blueapi.service.model import TaskRequest
from blueapi.worker import TaskStatus, WorkerState
from blueapi.worker.event import TaskError, TaskResult
from bluesky import RunEngine, RunEngineInterrupted
from bluesky import plan_stubs as bps

from mx_bluesky.common.parameters.components import (
    MxBlueskyParameters,
    get_param_version,
)
from mx_bluesky.common.parameters.constants import Status
from mx_bluesky.common.utils.exceptions import CrystalNotFoundError, SampleError
from mx_bluesky.hyperion._plan_runner_params import (
    RobotUnload,
    UDCCleanup,
    UDCDefaultState,
    Wait,
)
from mx_bluesky.hyperion.blueapi.parameters import LoadCentreCollectParams
from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
from mx_bluesky.hyperion.plan_runner import PlanError
from mx_bluesky.hyperion.supervisor import SupervisorRunner

TEST_VISIT = "cm12345-67"


@pytest.fixture
def mock_bluesky_context(run_engine: RunEngine):
    return BlueskyContext(run_engine=run_engine)


@pytest.fixture(autouse=True)
def mock_blueapi_client():
    with patch(
        "mx_bluesky.hyperion.supervisor._supervisor.BlueapiClient"
    ) as mock_class:
        mock_client = mock_class.from_config.return_value
        yield mock_client


@pytest.fixture(autouse=True)
def mock_blueapi_client_state(mock_blueapi_client):
    mock_state = PropertyMock(return_value=WorkerState.IDLE)
    type(mock_blueapi_client).state = mock_state
    yield mock_state


@pytest.fixture
def blueapi_config():
    return MagicMock()


@pytest.fixture
def runner(mock_bluesky_context, blueapi_config):
    runner = SupervisorRunner(mock_bluesky_context, blueapi_config, True)
    with patch.object(runner, "check_external_callbacks_are_alive"):
        yield runner


@patch("mx_bluesky.hyperion.supervisor._supervisor.TaskMonitor")
def test_decode_and_execute_load_centre_collect_executes_and_monitors_the_task_and_returns_the_visit(
    mock_task_monitor: MagicMock,
    mock_blueapi_client: MagicMock,
    runner: SupervisorRunner,
    external_load_centre_collect_params,
):
    parent = MagicMock()
    parent.attach_mock(mock_blueapi_client, "blueapi_client")
    parent.attach_mock(mock_task_monitor, "TaskMonitor")
    parent.attach_mock(mock_task_monitor.return_value, "task_monitor")

    result = runner.context.run_engine(
        runner.decode_and_execute(TEST_VISIT, [external_load_centre_collect_params])
    )

    expected_task_request = TaskRequest(
        name="load_centre_collect",
        params={"parameters": external_load_centre_collect_params},
        instrument_session=TEST_VISIT,
    )
    parent.assert_has_calls(
        [
            call.TaskMonitor(mock_blueapi_client, expected_task_request),
            call.task_monitor.__enter__(),
            call.blueapi_client.run_task(
                expected_task_request,
                on_event=mock_task_monitor.return_value.__enter__.return_value.on_blueapi_event,
            ),
            call.task_monitor.__exit__(None, None, None),
        ]
    )

    assert result.plan_result == TEST_VISIT  # type: ignore


def test_decode_and_execute_wait(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner
):
    mock_sleep = AsyncMock()
    runner.context.run_engine.register_command("sleep", mock_sleep)
    try:
        runner.context.run_engine(
            runner.decode_and_execute(
                TEST_VISIT,
                [
                    Wait.model_validate(
                        {
                            "parameter_model_version": get_param_version(),
                            "duration_s": 10,
                        }
                    )
                ],
            )
        )
    finally:
        runner.context.run_engine.register_command(
            "sleep", runner.context.run_engine._sleep
        )

    mock_blueapi_client.run_task.assert_not_called()
    mock_sleep.assert_awaited_once()


def test_decode_and_execute_default_state(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner
):
    runner.context.run_engine(
        runner.decode_and_execute(
            TEST_VISIT,
            [
                UDCDefaultState.model_validate(
                    {"parameter_model_version": get_param_version()}
                )
            ],
        )
    )

    mock_blueapi_client.run_task.assert_called_once_with(
        TaskRequest(
            name="move_to_udc_default_state", params={}, instrument_session=TEST_VISIT
        ),
        on_event=ANY,
    )


def test_decode_and_execute_robot_unload(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner
):
    runner.context.run_engine(runner.decode_and_execute(TEST_VISIT, [RobotUnload()]))

    mock_blueapi_client.run_task.assert_called_once_with(
        TaskRequest(
            name="robot_unload",
            params={"visit": TEST_VISIT},
            instrument_session=TEST_VISIT,
        ),
        on_event=ANY,
    )


def test_decode_and_execute_udc_cleanup(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner
):
    runner.context.run_engine(
        runner.decode_and_execute(
            TEST_VISIT,
            [
                UDCCleanup.model_validate(
                    {"parameter_model_version": get_param_version()}
                )
            ],
        )
    )

    mock_blueapi_client.run_task.assert_called_once_with(
        TaskRequest(
            name="clean_up_udc",
            params={},
            instrument_session=TEST_VISIT,
        ),
        on_event=ANY,
    )


def test_current_status_set_to_busy_during_execution(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner, executor: Executor
):
    task_executing = Event()
    check_complete = Event()

    def check_busy():
        task_executing.wait(timeout=1)
        assert runner.current_status == Status.BUSY
        check_complete.set()

    def wait_for_check(*args, **kwargs):
        task_executing.set()
        check_complete.wait(1)
        return TaskStatus(
            task_id="TASK_ID",
            result=TaskResult(result=None, type="None"),
            task_complete=True,
            task_failed=False,
        )

    mock_blueapi_client.run_task.side_effect = wait_for_check

    fut = executor.submit(check_busy)
    try:
        assert runner.current_status == Status.IDLE
        runner.context.run_engine(
            runner.decode_and_execute(
                TEST_VISIT,
                [
                    UDCDefaultState.model_validate(
                        {"parameter_model_version": get_param_version()}
                    )
                ],
            )
        )
    finally:
        fut.result(1)


def test_current_status_set_to_failed_on_exception_and_raise_plan_error(
    mock_blueapi_client: MagicMock, runner: SupervisorRunner
):
    status = TaskStatus(
        task_id="TASK_ID",
        result=TaskError(type="RuntimeError", message="Simulated exception"),
        task_complete=True,
        task_failed=True,
    )
    mock_blueapi_client.run_task.return_value = status

    with pytest.raises(
        PlanError, match="Exception raised.* message='Simulated exception'"
    ):
        runner.context.run_engine(
            runner.decode_and_execute(
                TEST_VISIT,
                [
                    UDCDefaultState.model_validate(
                        {"parameter_model_version": get_param_version()}
                    )
                ],
            )
        )
    assert runner.current_status == Status.FAILED


def test_is_connected_queries_blueapi_client(
    runner: SupervisorRunner,
    mock_blueapi_client: MagicMock,
    mock_blueapi_client_state: PropertyMock,
):
    assert runner.is_connected()
    mock_blueapi_client_state.assert_called_once()


def test_is_connected_returns_false_on_exception(
    runner: SupervisorRunner,
    mock_blueapi_client: MagicMock,
    mock_blueapi_client_state: PropertyMock,
):
    mock_blueapi_client_state.side_effect = RuntimeError("Simulated exception")
    assert not runner.is_connected()


def test_shutdown_sends_abort_to_run_engine_when_idle(
    runner: SupervisorRunner, executor: Executor
):
    runner.run_engine = MagicMock(spec=RunEngine)

    def request_shutdown():
        runner.shutdown()

    fut = executor.submit(request_shutdown)

    fut.result(1)
    runner.run_engine.abort.assert_called_once()


def test_shutdown_sends_abort_to_blueapi_client_when_running_then_aborts(
    runner: SupervisorRunner, executor: Executor
):
    task_running = Event()
    remote_abort_requested = Event()

    def request_shutdown():
        task_running.wait(1)
        runner.shutdown()

    def mock_run_task(*args, **kwargs):
        task_running.set()
        remote_abort_requested.wait(1)
        raise BlueskyStreamingError("Simulated abort exception")

    def mock_baton_handler_loop():
        yield from runner.decode_and_execute(
            TEST_VISIT,
            [
                UDCDefaultState.model_validate(
                    {"parameter_model_version": get_param_version()}
                )
            ],
        )
        while True:
            yield from bps.sleep(0.1)

    runner.blueapi_client.run_task.side_effect = mock_run_task  # type: ignore
    runner.blueapi_client.abort.side_effect = remote_abort_requested.set  # type: ignore

    fut = executor.submit(request_shutdown)
    try:
        with pytest.raises(RunEngineInterrupted):
            runner.context.run_engine(mock_baton_handler_loop())
    finally:
        fut.result(1)


def test_exception_during_callback_check_raises_plan_error(
    runner: SupervisorRunner, load_centre_collect_params: LoadCentreCollect
):
    with patch.object(
        runner,
        "check_external_callbacks_are_alive",
        side_effect=RuntimeError("Simulated exception"),
    ):
        with pytest.raises(PlanError, match="Exception raised.*: Simulated exception"):
            runner.context.run_engine(
                runner.decode_and_execute(TEST_VISIT, [load_centre_collect_params])
            )


def test_unrecognised_instruction_raises_assertion_error(runner: SupervisorRunner):
    with pytest.raises(AssertionError, match="Unsupported instruction"):
        runner.context.run_engine(
            runner.decode_and_execute(
                TEST_VISIT,
                [
                    MxBlueskyParameters.model_validate(
                        {"parameter_model_version": get_param_version()}
                    )
                ],
            )
        )


@pytest.mark.parametrize(
    "exception_type", [SampleError.__name__, CrystalNotFoundError.__name__]
)
def test_sample_error_skips_subsequent_instructions(
    runner: SupervisorRunner,
    external_load_centre_collect_params: LoadCentreCollectParams,
    mock_blueapi_client: MagicMock,
    exception_type: str,
):
    mock_blueapi_client.run_task.return_value = TaskStatus(
        task_id="TASK_ID",
        result=TaskError(type=exception_type, message="Simulated sample error"),
        task_complete=True,
        task_failed=True,
    )
    runner.context.run_engine(
        runner.decode_and_execute(
            TEST_VISIT,
            [external_load_centre_collect_params, external_load_centre_collect_params],
        )
    )
    mock_blueapi_client.run_task.assert_called_once_with(ANY, on_event=ANY)
