from __future__ import annotations

import os
from collections.abc import Iterator
from itertools import accumulate
from typing import Annotated, Any

from annotated_types import Len
from dodal.devices.aperturescatterguard import ApertureValue
from dodal.devices.detector import DetectorParams
from dodal.devices.zebra import (
    RotationDirection,
)
from dodal.log import LOGGER
from pydantic import Field, field_validator, model_validator
from scanspec.core import AxesPoints
from scanspec.core import Path as ScanPath
from scanspec.specs import Line

from mx_bluesky.hyperion.parameters.components import (
    DiffractionExperimentWithSample,
    IspybExperimentType,
    OptionalGonioAngleStarts,
    OptionalXyzStarts,
    RotationAxis,
    SplitScan,
    WithScan,
)
from mx_bluesky.hyperion.parameters.constants import (
    CONST,
    I03Constants,
)


class RotationScanPerSweep(OptionalGonioAngleStarts, OptionalXyzStarts):
    omega_start_deg: float = Field(default=0)  # type: ignore
    rotation_axis: RotationAxis = Field(default=RotationAxis.OMEGA)
    scan_width_deg: float = Field(default=360, gt=0)
    rotation_direction: RotationDirection = Field(default=RotationDirection.NEGATIVE)
    nexus_vds_start_img: int = Field(default=0, ge=0)


class RotationExperiment(DiffractionExperimentWithSample):
    shutter_opening_time_s: float = Field(default=CONST.I03.SHUTTER_TIME_S)
    rotation_increment_deg: float = Field(default=0.1, gt=0)
    ispyb_experiment_type: IspybExperimentType = Field(
        default=IspybExperimentType.ROTATION
    )

    def _detector_params(self, omega_start_deg: float):
        self.det_dist_to_beam_converter_path = (
            self.det_dist_to_beam_converter_path
            or CONST.PARAM.DETECTOR.BEAM_XY_LUT_PATH
        )
        optional_args = {}
        if self.run_number:
            optional_args["run_number"] = self.run_number
        assert self.detector_distance_mm is not None
        os.makedirs(self.storage_directory, exist_ok=True)
        return DetectorParams(
            detector_size_constants=I03Constants.DETECTOR,
            expected_energy_ev=self.demand_energy_ev,
            exposure_time=self.exposure_time_s,
            directory=self.storage_directory,
            prefix=self.file_name,
            detector_distance=self.detector_distance_mm,
            omega_start=omega_start_deg,
            omega_increment=self.rotation_increment_deg,
            num_images_per_trigger=self.num_images,
            num_triggers=1,
            use_roi_mode=False,
            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
            **optional_args,
        )

    @field_validator("selected_aperture")
    @classmethod
    def _set_default_aperture_position(cls, aperture_position: ApertureValue | None):
        if not aperture_position:
            default_aperture = CONST.PARAM.ROTATION.DEFAULT_APERTURE_POSITION
            LOGGER.warning(
                f"No aperture position selected. Defaulting to {default_aperture}"
            )
            return default_aperture
        else:
            return aperture_position


class RotationScan(WithScan, RotationScanPerSweep, RotationExperiment):
    @property
    def detector_params(self):
        return self._detector_params(self.omega_start_deg)

    @property
    def scan_points(self) -> AxesPoints:
        scan_spec = Line(
            axis="omega",
            start=self.omega_start_deg,
            stop=(
                self.omega_start_deg
                + (self.scan_width_deg - self.rotation_increment_deg)
            ),
            num=self.num_images,
        )
        scan_path = ScanPath(scan_spec.calculate())
        return scan_path.consume().midpoints

    @property
    def num_images(self) -> int:
        return int(self.scan_width_deg / self.rotation_increment_deg)


class MultiRotationScan(RotationExperiment, SplitScan):
    rotation_scans: Annotated[list[RotationScanPerSweep], Len(min_length=1)]

    def _single_rotation_scan(self, scan: RotationScanPerSweep) -> RotationScan:
        # self has everything from RotationExperiment
        params = self.model_dump()
        del params["rotation_scans"]
        # provided `scan` has everything from RotationScanPerSweep
        params.update(scan.model_dump())
        # together they have everything for RotationScan
        return RotationScan(**params)

    @model_validator(mode="after")
    @classmethod
    def correct_start_vds(cls, values: Any) -> Any:
        assert isinstance(values, MultiRotationScan)
        start_img = 0.0
        for scan in values.rotation_scans:
            scan.nexus_vds_start_img = int(start_img)
            start_img += scan.scan_width_deg / values.rotation_increment_deg
        return values

    @property
    def single_rotation_scans(self) -> Iterator[RotationScan]:
        for scan in self.rotation_scans:
            yield self._single_rotation_scan(scan)

    def _num_images_per_scan(self):
        return [
            int(scan.scan_width_deg / self.rotation_increment_deg)
            for scan in self.rotation_scans
        ]

    @property
    def num_images(self):
        return sum(self._num_images_per_scan())

    @property
    def scan_indices(self):
        return list(accumulate([0, *self._num_images_per_scan()]))

    @property
    def detector_params(self):
        return self._detector_params(self.rotation_scans[0].omega_start_deg)

    @property
    def ispyb_params(self):  # pyright: ignore
        raise ValueError("Please get ispyb params from one of the individual scans")
