import subprocess
import time
from concurrent.futures.thread import ThreadPoolExecutor
from functools import partial
from os import environ, getcwd
from pathlib import Path
from threading import Event
from time import sleep
from unittest.mock import MagicMock, call, patch

import pytest
from blueapi.client.event_bus import AnyEvent, EventBusClient
from blueapi.config import ApplicationConfig, ConfigLoader
from blueapi.core import BlueskyContext, DataEvent
from blueapi.service.model import TaskRequest
from blueapi.worker import WorkerEvent, WorkerState
from bluesky import RunEngine, RunEngineInterrupted
from bluesky import plan_stubs as bps
from bluesky_stomp.messaging import MessageContext

from mx_bluesky.common.external_interaction.alerting import Metadata
from mx_bluesky.common.parameters.components import get_param_version
from mx_bluesky.common.parameters.constants import Status
from mx_bluesky.hyperion._plan_runner_params import RobotUnload
from mx_bluesky.hyperion.blueapi.parameters import LoadCentreCollectParams
from mx_bluesky.hyperion.plan_runner import PlanError
from mx_bluesky.hyperion.supervisor import SupervisorRunner

from ....conftest import (
    TEST_CONTAINER,
    TEST_SAMPLE_ID,
    TEST_VISIT,
    raw_params_from_file,
)
from .dummy_plans import (
    BEAMLINE_ERROR_SAMPLE_ID,
    CRYSTAL_NOT_FOUND_SAMPLE_ID,
    WaitForFeedbackParams,
)

BLUEAPI_SERVER_CONFIG = (
    "tests/system_tests/hyperion/supervisor/system_test_blueapi.yaml"
)


@pytest.fixture(scope="module")
def tpe():
    return ThreadPoolExecutor(max_workers=1)


@pytest.fixture(scope="module")
def mock_blueapi_server():
    with subprocess.Popen(
        [
            "blueapi",
            "--config",
            BLUEAPI_SERVER_CONFIG,
            "serve",
        ],
        env=environ | {"PYTHONPATH": getcwd() + "/tests"},
    ) as blueapi_server:
        try:
            yield blueapi_server
        finally:
            blueapi_server.terminate()


@pytest.fixture
def mock_bluesky_context(run_engine: RunEngine):
    loader = ConfigLoader(ApplicationConfig)
    loader.use_values_from_yaml(
        Path("tests/system_tests/hyperion/supervisor/supervisor_config.yaml")
    )
    supervisor_config = loader.load()
    yield BlueskyContext(configuration=supervisor_config, run_engine=run_engine)


@pytest.fixture
def client_config() -> ApplicationConfig:
    loader = ConfigLoader(ApplicationConfig)
    loader.use_values_from_yaml(
        Path("tests/system_tests/hyperion/supervisor/client_config.yaml")
    )
    return loader.load()


def get_event_bus_client(supervisor: SupervisorRunner) -> EventBusClient:
    return supervisor.blueapi_client._events  # type: ignore


@pytest.fixture
def supervisor_runner(supervisor_runner_no_ping: SupervisorRunner):
    supervisor_runner_no_ping.reset_callback_watchdog_timer()
    yield supervisor_runner_no_ping


@pytest.fixture
def supervisor_runner_no_ping(
    mock_bluesky_context: BlueskyContext,
    client_config: ApplicationConfig,
    mock_blueapi_server,
    monkeypatch,
):
    monkeypatch.setenv("BEAMLINE", "i03")
    runner = SupervisorRunner(mock_bluesky_context, client_config, True)
    timeout = time.monotonic() + 30
    while time.monotonic() < timeout:
        if runner.is_connected():
            return runner
        sleep(1)
    else:
        raise AssertionError("Failed to connect to blueapi")


def handle_event(plan_started: Event, event_payload: AnyEvent, context: MessageContext):
    match event_payload:
        case DataEvent() as data_event:
            if (
                data_event.name == "start"
                and data_event.doc["plan_name"] == "robot_unload"
            ):
                plan_started.set()


@pytest.mark.system_test
def test_supervisor_connects_to_blueapi_and_stomp(
    mock_blueapi_server,
    mock_bluesky_context: BlueskyContext,
    client_config: ApplicationConfig,
    supervisor_runner: SupervisorRunner,
):
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )
    ebc = get_event_bus_client(supervisor_runner)

    received_message_event = Event()

    ebc.subscribe_to_all_events(partial(handle_event, received_message_event))
    assert supervisor_runner.current_status == Status.IDLE
    supervisor_runner.run_engine(
        supervisor_runner.decode_and_execute(TEST_VISIT, [params])
    )
    received_message_event.wait()


@pytest.mark.skip(reason="https://github.com/DiamondLightSource/blueapi/issues/1312")
def test_supervisor_continues_to_next_instruction_on_warning_error(
    supervisor_runner: SupervisorRunner,
):
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )
    supervisor_runner.run_engine(
        supervisor_runner.decode_and_execute("raise_warning_error", [params])
    )
    assert supervisor_runner.current_status == Status.FAILED


def test_supervisor_raises_request_abort_when_shutdown_requested(
    supervisor_runner: SupervisorRunner, tpe: ThreadPoolExecutor
):
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )
    ebc = get_event_bus_client(supervisor_runner)
    plan_aborted = Event()
    plan_called = Event()

    def handle_abort(event_payload: AnyEvent, context: MessageContext):
        match event_payload:
            case WorkerEvent() as worker_event:
                if worker_event.state == WorkerState.ABORTING:
                    plan_aborted.set()

    ebc.subscribe_to_all_events(partial(handle_event, plan_called))
    ebc.subscribe_to_all_events(handle_abort)

    def shutdown_in_background():
        plan_called.wait(10)
        assert supervisor_runner.current_status == Status.BUSY
        assert supervisor_runner.blueapi_client.state == WorkerState.RUNNING
        supervisor_runner.shutdown()
        assert supervisor_runner.current_status == Status.ABORTING

    fut = tpe.submit(shutdown_in_background)

    def execute_remotely():
        yield from supervisor_runner.decode_and_execute("wait_for_abort", [params])
        # Simulate supervisor going round in a loop, since async abort request may not be processed before requested execution
        # completes

        yield from bps.sleep(10)

    with pytest.raises(RunEngineInterrupted):
        supervisor_runner.run_engine(execute_remotely())

    assert supervisor_runner.blueapi_client.state == WorkerState.IDLE
    assert plan_aborted.wait(10)
    fut.result()


def test_supervisor_raises_plan_error_when_plan_fails_with_other_exception(
    supervisor_runner: SupervisorRunner,
):
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )
    with pytest.raises(PlanError, match="Exception raised during plan execution:"):
        supervisor_runner.run_engine(
            supervisor_runner.decode_and_execute("raise_other_error", [params])
        )
    assert supervisor_runner.current_status == Status.FAILED


def test_shutdown_raises_run_engine_interrupted_when_idle(
    supervisor_runner: SupervisorRunner, tpe: ThreadPoolExecutor
):
    plan_started = Event()

    def idle_plan():
        plan_started.set()
        while True:
            yield from bps.sleep(1)

    def shutdown_in_background():
        plan_started.wait(10)
        supervisor_runner.shutdown()

    fut = tpe.submit(shutdown_in_background)

    with pytest.raises(RunEngineInterrupted):
        supervisor_runner.run_engine(idle_plan())

    fut.result()


async def test_supervisor_checks_for_external_callback_ping(
    supervisor_runner_no_ping: SupervisorRunner, tpe: ThreadPoolExecutor
):
    ebc = get_event_bus_client(supervisor_runner_no_ping)
    received_message_event = Event()
    ebc.subscribe_to_all_events(partial(handle_event, received_message_event))
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )

    def run_test_in_background():
        sleep(1)
        assert not received_message_event.is_set()
        supervisor_runner_no_ping.reset_callback_watchdog_timer()
        received_message_event.wait(1)

    fut = tpe.submit(run_test_in_background)
    supervisor_runner_no_ping.run_engine(
        supervisor_runner_no_ping.decode_and_execute(TEST_VISIT, [params])
    )
    fut.result()


def test_supervisor_raises_plan_error_when_external_callbacks_watchdog_expired(
    supervisor_runner_no_ping: SupervisorRunner,
):
    runner = supervisor_runner_no_ping
    runner.EXTERNAL_CALLBACK_WATCHDOG_TIMER_S = 0.5  # type: ignore
    runner.reset_callback_watchdog_timer()
    params = RobotUnload.model_validate(
        {"parameter_model_version": get_param_version()}
    )
    # Allow callback watchdog to expire
    sleep(1)
    with pytest.raises(PlanError, match="External callback watchdog timer expired.*"):
        runner.run_engine(runner.decode_and_execute(TEST_VISIT, [params]))


def test_supervisor_calls_load_centre_collect(
    supervisor_runner: SupervisorRunner, tmp_path
):
    params = LoadCentreCollectParams(
        **raw_params_from_file(
            "tests/test_data/parameter_json_files/external_load_centre_collect_params.json",
            tmp_path,
        )
    )

    supervisor_runner.run_engine(
        supervisor_runner.decode_and_execute(TEST_VISIT, [params])
    )


@patch("mx_bluesky.hyperion.supervisor._task_monitor.TaskMonitor.DEFAULT_TIMEOUT_S", 1)
@patch("mx_bluesky.hyperion.supervisor._task_monitor.get_alerting_service")
def test_supervisor_alerts_repeatedly_when_waiting_for_beam(
    mock_get_alerting_service: MagicMock,
    supervisor_runner: SupervisorRunner,
):
    mock_alerting_service = mock_get_alerting_service.return_value
    params = WaitForFeedbackParams(
        time_for_beam_stable=2.5,
        time_in_plan=0.5,
        sample_id=TEST_SAMPLE_ID,
        visit=TEST_VISIT,
        sample_puck=TEST_CONTAINER,
    )
    task_request = TaskRequest(
        name="wait_for_feedback",
        params={"parameters": params},
        instrument_session=TEST_VISIT,
    )
    supervisor_runner._run_task_remotely(task_request)

    mock_alerting_service.raise_alert.assert_has_calls(
        [
            call(
                "Hyperion is paused waiting for beam on i03.",
                "Hyperion has been paused waiting for beam for 0 minutes.",
                {
                    Metadata.SAMPLE_ID: TEST_SAMPLE_ID,
                    Metadata.VISIT: TEST_VISIT,
                    Metadata.CONTAINER: TEST_CONTAINER,
                },
            ),
            call(
                "Hyperion is paused waiting for beam on i03.",
                "Hyperion has been paused waiting for beam for 0 minutes.",
                {
                    Metadata.SAMPLE_ID: TEST_SAMPLE_ID,
                    Metadata.VISIT: TEST_VISIT,
                    Metadata.CONTAINER: TEST_CONTAINER,
                },
            ),
        ]
    )


@patch("mx_bluesky.hyperion.supervisor._task_monitor.TaskMonitor.DEFAULT_TIMEOUT_S", 1)
@patch("mx_bluesky.hyperion.supervisor._task_monitor.get_alerting_service")
def test_supervisor_alerts_with_error_and_aborts_task_when_stuck_not_waiting_for_beam(
    mock_get_alerting_service: MagicMock,
    supervisor_runner: SupervisorRunner,
):
    mock_alerting_service = mock_get_alerting_service.return_value
    params = WaitForFeedbackParams(
        time_for_beam_stable=0.5,
        time_in_plan=1.5,
        sample_id=TEST_SAMPLE_ID,
        visit=TEST_VISIT,
        sample_puck=TEST_CONTAINER,
    )

    task_request = TaskRequest(
        name="wait_for_feedback",
        params={"parameters": params},
        instrument_session=TEST_VISIT,
    )

    with pytest.raises(PlanError, match="Exception raised during plan execution"):
        supervisor_runner._run_task_remotely(task_request)

    mock_alerting_service.raise_alert.assert_has_calls(
        [
            call(
                "UDC encountered an error on i03",
                "Hyperion Supervisor detected that BlueAPI was stuck for 1 seconds.",
                {
                    Metadata.SAMPLE_ID: TEST_SAMPLE_ID,
                    Metadata.VISIT: TEST_VISIT,
                    Metadata.CONTAINER: TEST_CONTAINER,
                },
            ),
        ]
    )


@patch("mx_bluesky.hyperion.supervisor._task_monitor.TaskMonitor.DEFAULT_TIMEOUT_S", 1)
@patch("mx_bluesky.hyperion.supervisor._task_monitor.get_alerting_service")
def test_supervisor_no_alerts_when_not_stuck(
    mock_get_alerting_service: MagicMock,
    supervisor_runner: SupervisorRunner,
):
    mock_alerting_service = mock_get_alerting_service.return_value
    params = WaitForFeedbackParams(
        time_for_beam_stable=0.5,
        time_in_plan=0.5,
        sample_id=TEST_SAMPLE_ID,
        visit=TEST_VISIT,
        sample_puck=TEST_CONTAINER,
    )

    task_request = TaskRequest(
        name="wait_for_feedback",
        params={"parameters": params},
        instrument_session=TEST_VISIT,
    )

    supervisor_runner._run_task_remotely(task_request)

    mock_alerting_service.raise_alert.assert_not_called()


def test_supervisor_decode_and_execute_raises_planerror_if_blueapi_plan_raises_exception(
    supervisor_runner: SupervisorRunner, tmp_path
):
    params = LoadCentreCollectParams(
        **raw_params_from_file(
            "tests/test_data/parameter_json_files/external_load_centre_collect_params.json",
            tmp_path,
        )
    )
    params.sample_id = BEAMLINE_ERROR_SAMPLE_ID
    with pytest.raises(PlanError, match="Simulated beamline error"):
        supervisor_runner.run_engine(
            supervisor_runner.decode_and_execute(TEST_VISIT, [params])
        )


def test_supervisor_decode_and_execute_continues_if_blueapi_plan_raises_sample_error(
    supervisor_runner: SupervisorRunner, tmp_path
):
    params = LoadCentreCollectParams(
        **raw_params_from_file(
            "tests/test_data/parameter_json_files/external_load_centre_collect_params.json",
            tmp_path,
        )
    )
    params.sample_id = CRYSTAL_NOT_FOUND_SAMPLE_ID
    supervisor_runner.run_engine(
        supervisor_runner.decode_and_execute(TEST_VISIT, [params, params])
    )
