import os
import re
from decimal import Decimal
from functools import partial
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from aiohttp import ClientResponse
from dodal.beamlines import i03
from dodal.devices.oav.oav_parameters import OAVConfig
from ophyd_async.core import AsyncStatus
from ophyd_async.testing import set_mock_value
from PIL import Image

from mx_bluesky.hyperion.parameters.constants import CONST

# Map all the case-sensitive column names from their normalised versions
DATA_COLLECTION_COLUMN_MAP = {
    s.lower(): s
    for s in [
        "dataCollectionId",
        "BLSAMPLEID",
        "SESSIONID",
        "experimenttype",
        "dataCollectionNumber",
        "startTime",
        "endTime",
        "runStatus",
        "axisStart",
        "axisEnd",
        "axisRange",
        "overlap",
        "numberOfImages",
        "startImageNumber",
        "numberOfPasses",
        "exposureTime",
        "imageDirectory",
        "imagePrefix",
        "imageSuffix",
        "imageContainerSubPath",
        "fileTemplate",
        "wavelength",
        "resolution",
        "detectorDistance",
        "xBeam",
        "yBeam",
        "comments",
        "printableForReport",
        "CRYSTALCLASS",
        "slitGapVertical",
        "slitGapHorizontal",
        "transmission",
        "synchrotronMode",
        "xtalSnapshotFullPath1",
        "xtalSnapshotFullPath2",
        "xtalSnapshotFullPath3",
        "xtalSnapshotFullPath4",
        "rotationAxis",
        "phiStart",
        "kappaStart",
        "omegaStart",
        "chiStart",
        "resolutionAtCorner",
        "detector2Theta",
        "DETECTORMODE",
        "undulatorGap1",
        "undulatorGap2",
        "undulatorGap3",
        "beamSizeAtSampleX",
        "beamSizeAtSampleY",
        "centeringMethod",
        "averageTemperature",
        "ACTUALSAMPLEBARCODE",
        "ACTUALSAMPLESLOTINCONTAINER",
        "ACTUALCONTAINERBARCODE",
        "ACTUALCONTAINERSLOTINSC",
        "actualCenteringPosition",
        "beamShape",
        "dataCollectionGroupId",
        "POSITIONID",
        "detectorId",
        "FOCALSPOTSIZEATSAMPLEX",
        "POLARISATION",
        "FOCALSPOTSIZEATSAMPLEY",
        "APERTUREID",
        "screeningOrigId",
        "flux",
        "strategySubWedgeOrigId",
        "blSubSampleId",
        "processedDataFile",
        "datFullPath",
        "magnification",
        "totalAbsorbedDose",
        "binning",
        "particleDiameter",
        "boxSize",
        "minResolution",
        "minDefocus",
        "maxDefocus",
        "defocusStepSize",
        "amountAstigmatism",
        "extractSize",
        "bgRadius",
        "voltage",
        "objAperture",
        "c1aperture",
        "c2aperture",
        "c3aperture",
        "c1lens",
        "c2lens",
        "c3lens",
        "startPositionId",
        "endPositionId",
        "flux",
        "bestWilsonPlotPath",
        "totalExposedDose",
        "nominalMagnification",
        "nominalDefocus",
        "imageSizeX",
        "imageSizeY",
        "pixelSizeOnImage",
        "phasePlate",
        "dataCollectionPlanId",
    ]
}


@pytest.fixture(autouse=True)
def use_dev_ispyb_unless_overridden_by_environment():
    ispyb_config_path = os.environ.get(
        "ISPYB_CONFIG_PATH", CONST.SIM.DEV_ISPYB_DATABASE_CFG
    )
    with patch.dict(os.environ, {"ISPYB_CONFIG_PATH": ispyb_config_path}):
        yield ispyb_config_path


@pytest.fixture
def undulator_for_system_test(undulator):
    set_mock_value(undulator.current_gap, 1.11)
    return undulator


@pytest.fixture
def oav_for_system_test(test_config_files):
    parameters = OAVConfig(
        test_config_files["zoom_params_file"], test_config_files["display_config"]
    )
    oav = i03.oav(connect_immediately=True, mock=True, params=parameters)
    set_mock_value(oav.cam.array_size_x, 1024)
    set_mock_value(oav.cam.array_size_y, 768)

    # Grid snapshots
    set_mock_value(oav.grid_snapshot.x_size, 1024)
    set_mock_value(oav.grid_snapshot.y_size, 768)
    set_mock_value(oav.grid_snapshot.top_left_x, 50)
    set_mock_value(oav.grid_snapshot.top_left_y, 100)
    size_in_pixels = 0.1 * 1000 / 1.25
    set_mock_value(oav.grid_snapshot.box_width, size_in_pixels)

    # Rotation snapshots
    @AsyncStatus.wrap
    async def trigger_with_test_image(self):
        with Image.open(
            "tests/test_data/test_images/generate_snapshot_input.png"
        ) as image:
            await self.post_processing(image)

    oav.snapshot.trigger = MagicMock(
        side_effect=partial(trigger_with_test_image, oav.snapshot)
    )
    oav.grid_snapshot.trigger = MagicMock(
        side_effect=partial(trigger_with_test_image, oav.grid_snapshot)
    )

    empty_response = AsyncMock(spec=ClientResponse)
    empty_response.read.return_value = b""

    with (
        patch(
            "dodal.devices.areadetector.plugins.MJPG.ClientSession.get", autospec=True
        ) as mock_get,
    ):
        mock_get.return_value.__aenter__.return_value = empty_response
        set_mock_value(oav.zoom_controller.level, "1.0")
        zoom_levels_list = ["1.0x", "3.0x", "5.0x", "7.5x", "10.0x", "15.0x"]
        oav.zoom_controller._get_allowed_zoom_levels = AsyncMock(
            return_value=zoom_levels_list
        )
        yield oav


def compare_actual_and_expected(
    id, expected_values, fetch_datacollection_attribute, column_map: dict | None = None
):
    results = "\n"
    for k, v in expected_values.items():
        actual = fetch_datacollection_attribute(
            id, column_map[k.lower()] if column_map else k
        )
        if isinstance(actual, Decimal):
            actual = float(actual)
        if isinstance(v, float):
            actual_v = actual == pytest.approx(v)
        elif isinstance(v, str) and v.startswith("regex:"):
            actual_v = re.match(v.removeprefix("regex:"), str(actual))  # type: ignore
        else:
            actual_v = actual == v
        if not actual_v:
            results += f"expected {k} {v} == {actual}\n"
    assert results == "\n", results


def compare_comment(
    fetch_datacollection_attribute, data_collection_id, expected_comment
):
    actual_comment = fetch_datacollection_attribute(
        data_collection_id, DATA_COLLECTION_COLUMN_MAP["comments"]
    )
    match = re.search(" Zocalo processing took", actual_comment)
    truncated_comment = actual_comment[: match.start()] if match else actual_comment
    assert truncated_comment == expected_comment
