import asyncio
import pprint
import sys
from collections.abc import Generator
from functools import partial
from pathlib import Path
from typing import Any, cast
from unittest.mock import MagicMock, patch

import pytest
from _pytest.fixtures import FixtureRequest
from bluesky.run_engine import RunEngine
from dodal.beamlines import i03
from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
from dodal.devices.backlight import Backlight
from dodal.devices.beamlines.i24.commissioning_jungfrau import CommissioningJungfrau
from dodal.devices.beamsize.beamsize import BeamsizeBase
from dodal.devices.detector.detector_motion import DetectorMotion
from dodal.devices.eiger import EigerDetector
from dodal.devices.fast_grid_scan import PandAFastGridScan, ZebraFastGridScanThreeD
from dodal.devices.flux import Flux
from dodal.devices.hutch_shutter import ShutterState
from dodal.devices.mx_phase1.beamstop import Beamstop
from dodal.devices.oav.oav_detector import OAV
from dodal.devices.oav.pin_image_recognition import PinTipDetection
from dodal.devices.robot import BartRobot
from dodal.devices.s4_slit_gaps import S4SlitGaps
from dodal.devices.smargon import Smargon
from dodal.devices.synchrotron import Synchrotron, SynchrotronMode
from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutterState
from dodal.devices.zocalo import ZocaloResults
from event_model import RunStart
from event_model.documents import Event
from ophyd_async.core import (
    AsyncStatus,
    AutoMaxIncrementingPathProvider,
    PathInfo,
    PathProvider,
    completed_status,
    init_devices,
    set_mock_value,
)
from ophyd_async.fastcs.panda import HDFPanda

from mx_bluesky.common.experiment_plans.beamstop_check import BeamstopCheckDevices
from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
    BeamlineSpecificFGSFeatures,
)
from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
    ZocaloCallback,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
    GridscanISPyBCallback,
    generate_start_info_from_omega_map,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
    GridscanNexusFileCallback,
)
from mx_bluesky.common.external_interaction.ispyb.data_model import (
    DataCollectionGroupInfo,
)
from mx_bluesky.common.external_interaction.ispyb.ispyb_store import (
    IspybIds,
    StoreInIspyb,
)
from mx_bluesky.common.parameters.constants import (
    DocDescriptorNames,
    EnvironmentConstants,
    PlanNameConstants,
)
from mx_bluesky.common.parameters.device_composites import (
    FlyScanEssentialDevices,
    GridDetectThenXRayCentreComposite,
)
from mx_bluesky.common.parameters.gridscan import GridCommon, SpecifiedThreeDGridScan
from mx_bluesky.hyperion.parameters.device_composites import (
    HyperionGridDetectThenXRayCentreComposite,
)
from tests.conftest import raw_params_from_file

_ALLOWED_PYTEST_TASKS = {"async_finalizer", "async_setup", "async_teardown"}


def _error_and_kill_pending_tasks(
    loop: asyncio.AbstractEventLoop, test_name: str, test_passed: bool
) -> set[asyncio.Task]:
    """Cancels pending tasks in the event loop for a test. Raises an exception if
    the test hasn't already.

    Args:
        loop: The event loop to check for pending tasks.
        test_name: The name of the test.
        test_passed: Indicates whether the test passed.

    Returns:
        set[asyncio.Task]: The set of unfinished tasks that were cancelled.

    Raises:
        RuntimeError: If there are unfinished tasks and the test didn't fail.
    """
    unfinished_tasks = {
        task
        for task in asyncio.all_tasks(loop)
        if hasattr(coro := task.get_coro(), "__name__")
        and coro.__name__ not in _ALLOWED_PYTEST_TASKS  # type: ignore
        and not task.done()
    }
    for task in unfinished_tasks:
        task.cancel()

    # We only raise an exception here if the test didn't fail anyway.
    # If it did then it makes sense that there's some tasks we need to cancel,
    # but an exception will already have been raised.
    if unfinished_tasks and test_passed:
        raise RuntimeError(
            f"Not all tasks closed during test {test_name}:\n"
            f"{pprint.pformat(unfinished_tasks, width=88)}"
        )

    return unfinished_tasks


@pytest.fixture(autouse=True, scope="function")
async def fail_test_on_unclosed_tasks(request: FixtureRequest):
    """
    Used on every test to ensure failure if there are pending tasks
    by the end of the test.
    """

    fail_count = request.session.testsfailed
    loop = asyncio.get_running_loop()

    loop.set_debug(True)

    request.addfinalizer(
        lambda: _error_and_kill_pending_tasks(
            loop, request.node.name, request.session.testsfailed == fail_count
        )
    )


BASIC_PRE_SETUP_DOC = {
    "undulator-current_gap": 0,
    "synchrotron-synchrotron_mode": SynchrotronMode.USER,
    "s4_slit_gaps-xgap": 0,
    "s4_slit_gaps-ygap": 0,
    "gonio-x": 10.0,
    "gonio-y": 20.0,
    "gonio-z": 30.0,
}

BASIC_POST_SETUP_DOC = {
    "aperture_scatterguard-selected_aperture": ApertureValue.OUT_OF_BEAM,
    "attenuator-actual_transmission": 0,
    "flux-flux_reading": 10,
    "dcm-energy_in_keV": 11.105,
    "beamsize-x_um": 50.0,
    "beamsize-y_um": 20.0,
    "eiger_cam_roi_mode": True,
    "eiger-ispyb_detector_id": 78,
}


def assert_event(mock_call, expected):
    actual = mock_call.args[0]
    if "data" in actual:
        actual = actual["data"]
    for k, v in expected.items():
        assert actual[k] == v, f"Mismatch in key {k}, {actual} <=> {expected}"


def create_gridscan_callbacks() -> tuple[
    GridscanNexusFileCallback, GridscanISPyBCallback
]:
    return (
        GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan),
        GridscanISPyBCallback(
            param_type=SpecifiedThreeDGridScan,
            emit=ZocaloCallback(
                PlanNameConstants.DO_FGS,
                EnvironmentConstants.ZOCALO_ENV,
                generate_start_info_from_omega_map,
            ),
        ),
    )


@pytest.fixture
def use_beamline_t01():
    """Beamline t01 is a beamline for unit tests that just contains a baton, so that
    loading the beamline context does not require importing lots of modules and instantiating
    many devices"""
    with patch.dict("os.environ", {"BEAMLINE": "t01"}):
        import tests.unit_tests.t01

        with (
            patch.dict(sys.modules, {"dodal.beamlines.t01": tests.unit_tests.t01}),
            patch("mx_bluesky.hyperion.in_process_runner.move_to_udc_default_state"),
            patch(
                "mx_bluesky.hyperion.in_process_runner.device_composite_from_context"
            ),
        ):
            yield tests.unit_tests.t01


@pytest.fixture
def mock_subscriptions(test_fgs_params):
    with (
        patch(
            "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
        ),
        patch(
            "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb"
        ) as mock_store_in_ispyb,
    ):
        mock_store_in_ispyb.return_value.begin_deposition.return_value = IspybIds(
            data_collection_ids=(100, 200), data_collection_group_id=0
        )
        mock_store_in_ispyb.return_value.update_deposition.return_value = IspybIds(
            data_collection_ids=(100, 200),
            data_collection_group_id=0,
            grid_ids=(0, 0),
        )
        nexus_callback, ispyb_callback = create_gridscan_callbacks()
        ispyb_callback.ispyb = MagicMock(spec=StoreInIspyb)

        yield (nexus_callback, ispyb_callback)


ReWithSubs = tuple[RunEngine, tuple[GridscanNexusFileCallback | GridscanISPyBCallback]]


@pytest.fixture
def run_engine_with_subs(
    run_engine: RunEngine,
    mock_subscriptions: tuple[GridscanNexusFileCallback | GridscanISPyBCallback],
) -> Generator[ReWithSubs, Any, None]:
    for cb in list(mock_subscriptions):
        run_engine.subscribe(cb)
    yield run_engine, mock_subscriptions


@pytest.fixture
def test_fgs_params(tmp_path, patch_beamline_env_variable):
    return SpecifiedThreeDGridScan(
        **raw_params_from_file(
            "tests/test_data/parameter_json_files/good_test_parameters.json", tmp_path
        )
    )


def mock_zocalo_trigger(zocalo: ZocaloResults, result):
    @AsyncStatus.wrap
    async def mock_complete(results):
        await zocalo._put_results(results, {"dcid": 0, "dcgid": 0})

    zocalo.trigger = MagicMock(side_effect=partial(mock_complete, result))


def modified_store_grid_scan_mock(*args, dcids=(0, 0), dcgid=0, **kwargs):
    mock = MagicMock(spec=StoreInIspyb)
    mock.begin_deposition.return_value = IspybIds(
        data_collection_ids=dcids, data_collection_group_id=dcgid
    )
    mock.update_deposition.return_value = IspybIds(
        data_collection_ids=dcids, data_collection_group_id=dcgid, grid_ids=(0, 0)
    )
    return mock


def make_event_doc(data, descriptor="abc123") -> Event:
    return {
        "time": 0,
        "timestamps": {"a": 0},
        "seq_num": 0,
        "uid": "not so random uid",
        "descriptor": descriptor,
        "data": data,
    }


def run_generic_ispyb_handler_setup(
    ispyb_handler: GridscanISPyBCallback,
    params: SpecifiedThreeDGridScan,
):
    """This is useful when testing 'run_gridscan_and_move(...)' because this stuff
    happens at the start of the outer plan."""

    ispyb_handler.active = True
    ispyb_handler.activity_gated_start(
        {
            "subplan_name": PlanNameConstants.GRIDSCAN_OUTER,
            "mx_bluesky_parameters": params.model_dump_json(),
        }  # type: ignore
    )
    ispyb_handler.activity_gated_descriptor(
        {"uid": "123abc", "name": DocDescriptorNames.HARDWARE_READ_PRE}  # type: ignore
    )
    ispyb_handler.activity_gated_event(
        make_event_doc(
            BASIC_PRE_SETUP_DOC,
            descriptor="123abc",
        )
    )
    ispyb_handler.activity_gated_descriptor(
        {"uid": "abc123", "name": DocDescriptorNames.HARDWARE_READ_DURING}  # type: ignore
    )
    ispyb_handler.activity_gated_event(
        make_event_doc(
            BASIC_POST_SETUP_DOC,
            descriptor="abc123",
        )
    )


@pytest.fixture
async def zebra_fast_grid_scan():
    zebra_fast_grid_scan = i03.zebra_fast_grid_scan.build(
        connect_immediately=True, mock=True
    )
    set_mock_value(zebra_fast_grid_scan.device_scan_invalid, 0.0)
    set_mock_value(zebra_fast_grid_scan.x_scan_valid, 1.0)
    set_mock_value(zebra_fast_grid_scan.y_scan_valid, 1.0)
    set_mock_value(zebra_fast_grid_scan.z_scan_valid, 1.0)
    set_mock_value(zebra_fast_grid_scan.position_counter, 0)
    return zebra_fast_grid_scan


@pytest.fixture
async def fake_fgs_composite(
    smargon: Smargon,
    test_fgs_params: SpecifiedThreeDGridScan,
    attenuator,
    xbpm_feedback,
    synchrotron,
    aperture_scatterguard,
    zocalo,
    panda,
    backlight,
):
    fake_composite = FlyScanEssentialDevices(
        # We don't use the eiger fixture here because .unstage() is used in some tests
        eiger=i03.eiger.build(mock=True),
        gonio=smargon,
        synchrotron=synchrotron,
    )

    fake_composite.eiger.stage = MagicMock(side_effect=lambda: completed_status())
    # unstage should be mocked on a per-test basis because several rely on unstage
    fake_composite.eiger.set_detector_parameters(test_fgs_params.detector_params)
    fake_composite.eiger.stop_odin_when_all_frames_collected = MagicMock()
    fake_composite.eiger.odin.check_and_wait_for_odin_state = lambda timeout: True

    test_result = {
        "centre_of_mass": [6, 6, 6],
        "max_voxel": [5, 5, 5],
        "max_count": 123456,
        "n_voxels": 321,
        "total_count": 999999,
        "bounding_box": [[3, 3, 3], [9, 9, 9]],
        "sample_id": 12345,
    }

    @AsyncStatus.wrap
    async def mock_complete(result):
        await zocalo._put_results([result], {"dcid": 0, "dcgid": 0})

    zocalo.trigger = MagicMock(side_effect=partial(mock_complete, test_result))  # type: ignore
    zocalo.timeout_s = 3
    set_mock_value(fake_composite.gonio.x.max_velocity, 10)

    return fake_composite


@pytest.fixture
def dummy_rotation_data_collection_group_info():
    return DataCollectionGroupInfo(
        visit_string="cm31105-4",
        experiment_type="SAD",
        sample_id=364758,
    )


@pytest.fixture
def beamline_specific(
    zebra_fast_grid_scan: ZebraFastGridScanThreeD,
) -> BeamlineSpecificFGSFeatures:
    return BeamlineSpecificFGSFeatures(
        setup_trigger_plan=MagicMock(),
        tidy_plan=MagicMock(),
        set_flyscan_params_plan=MagicMock(),
        fgs_motors=zebra_fast_grid_scan,
        read_pre_flyscan_plan=MagicMock(),
        read_during_collection_plan=MagicMock(),
    )


@pytest.fixture
def test_full_grid_scan_params(tmp_path):
    params = raw_params_from_file(
        "tests/test_data/parameter_json_files/good_test_grid_with_edge_detect_parameters.json",
        tmp_path,
    )
    return GridCommon(**params)


@pytest.fixture
async def grid_detect_xrc_devices(
    aperture_scatterguard: ApertureScatterguard,
    backlight: Backlight,
    beamstop_phase1: Beamstop,
    beamsize: BeamsizeBase,
    detector_motion: DetectorMotion,
    eiger: EigerDetector,
    smargon: Smargon,
    oav: OAV,
    ophyd_pin_tip_detection: PinTipDetection,
    zocalo: ZocaloResults,
    synchrotron: Synchrotron,
    fast_grid_scan: ZebraFastGridScanThreeD,
    s4_slit_gaps: S4SlitGaps,
    flux: Flux,
    zebra,
    zebra_shutter,
    xbpm_feedback,
    attenuator,
    undulator,
    dcm,
):
    yield GridDetectThenXRayCentreComposite(
        aperture_scatterguard=aperture_scatterguard,
        attenuator=attenuator,
        backlight=backlight,
        beamstop=beamstop_phase1,
        beamsize=beamsize,
        detector_motion=detector_motion,
        eiger=eiger,
        zebra_fast_grid_scan=fast_grid_scan,
        flux=flux,
        oav=oav,
        pin_tip_detection=ophyd_pin_tip_detection,
        gonio=smargon,
        synchrotron=synchrotron,
        s4_slit_gaps=s4_slit_gaps,
        undulator=undulator,
        xbpm_feedback=xbpm_feedback,
        zebra=zebra,
        zocalo=zocalo,
        dcm=dcm,
        robot=MagicMock(spec=BartRobot),
        sample_shutter=zebra_shutter,
    )


@pytest.fixture
async def hyperion_grid_detect_xrc_devices(grid_detect_xrc_devices):
    composite = cast(HyperionGridDetectThenXRayCentreComposite, grid_detect_xrc_devices)
    composite.panda = MagicMock(spec=HDFPanda)
    composite.panda_fast_grid_scan = MagicMock(spec=PandAFastGridScan)
    return composite


class _BasePathProvider(PathProvider):
    """Minimal adaption of BlueAPI's PathProvider.

    To use this PathProvider in a test, a run must be open
    with 'detector_file_template' attached as metadata. This is done
    automatically when using BlueAPI

    For longer term plans to test plans with BlueAPI in CI,
    see https://github.com/DiamondLightSource/blueapi/issues/1206"""

    def __init__(self, tmp_path: Path) -> None:
        self.path = tmp_path
        self._docs: list[RunStart] = []

    def run_start(self, name: str, start_document: RunStart) -> None:
        if name == "start":
            self._docs.append(start_document)

    def __call__(self, device_name: str | None = None) -> PathInfo:
        template = self._docs[-1].get("detector_file_template")
        if not template:
            raise ValueError("detector_file_template must be set in metadata")
        sub_path = template.format_map(self._docs[-1] | {"device_name": device_name})
        return PathInfo(directory_path=self.path, filename=sub_path)


# See https://github.com/DiamondLightSource/dodal/issues/1455
@pytest.fixture
def jungfrau(tmp_path: Path, run_engine: RunEngine) -> CommissioningJungfrau:
    with init_devices(mock=True):
        base_provider = _BasePathProvider(tmp_path)
        path = AutoMaxIncrementingPathProvider(base_provider)
        detector = CommissioningJungfrau("", "", path)
    set_mock_value(detector._writer.writer_ready, 1)
    run_engine.subscribe(base_provider.run_start, "start")

    return detector


@pytest.fixture
async def beamstop_check_devices(
    aperture_scatterguard,
    attenuator,
    backlight,
    baton,
    detector_motion,
    ipin,
    zebra_shutter,
    xbpm_feedback,
    sim_run_engine,
    run_engine,
):
    async def noop(_):
        await asyncio.sleep(0)

    run_engine.register_command("sleep", noop)
    try:
        async with init_devices(mock=True):
            beamstop = Beamstop("", MagicMock())

        devices = BeamstopCheckDevices(
            aperture_scatterguard=aperture_scatterguard,
            attenuator=attenuator,
            backlight=backlight,
            baton=baton,
            beamstop=beamstop,
            detector_motion=detector_motion,
            ipin=ipin,
            sample_shutter=zebra_shutter,
            xbpm_feedback=xbpm_feedback,
        )
        sim_run_engine.add_read_handler_for(
            devices.sample_shutter, ZebraShutterState.CLOSE
        )
        sim_run_engine.add_handler(
            "locate",
            lambda msg: {"readback": ShutterState.CLOSED},
            "detector_motion-shutter",
        )
        sim_run_engine.add_read_handler_for(ipin.pin_readback, 0.1)

        return devices
    finally:
        run_engine.register_command("sleep", run_engine._sleep)


@pytest.fixture
async def ipin():
    yield i03.ipin.build(connect_immediately=True, mock=True)
