from typing import Any, Dict, List, Union

from assonant._bluesky import DataParserFactory
from assonant.naming_standards import AcquisitionType

from .assonant_data_retriever_interface import IAssonantDataRetriever
from .exceptions import AssonantDataRetrieverError


class BlueskyDataRetriever(IAssonantDataRetriever):
    """Bluesky Data Retriever."""

    retrieved_data = None

    @property
    def bluesky_data_parser(self):
        """Bluesky Data Parser getter."""
        return self._bluesky_data_parser

    @bluesky_data_parser.setter
    def bluesky_data_parser(self, parser):
        """Bluesky Data Parser setter. Invalidzgds retrieved data cache when set"""
        self._bluesky_data_parser = parser
        # Always invalidate cache when parser changes
        self.retrieved_data = None

    def __init__(self, bluesky_data_source: Union[List[List[Union[str, Dict[str, Any]]]], str]):
        """Instantiates BlueskyDataRetriever to retrieve data from given data source.

        Args:
            bluesky_data_source (Union[List[List[Union[str, Dict[str, Any]]]], str]): DataSource containing
            the bluesky data. Currently, supported formats are: JSON file containing bluesky events or a
            List of bluesky documents data presented also as  a list. The list with the document data has
            2 positions, the 1st is the bluesky identification for the document
            (start, event, stop, descriptor, ...) and the 2nd position the dict containing the document
            data.

        Raises:
            AssonantDataRetrieverError: Failed to create the suitable Bluesky Data Parser for given data source.
        """
        try:
            self.bluesky_data_parser = DataParserFactory.create_data_parser(bluesky_data_source=bluesky_data_source)
        except Exception as e:
            raise AssonantDataRetrieverError(
                "BlueskyDataRetriever failed to create Bluesky DataParser for given data source."
            ) from e

    def _convert_acquisition_type_into_event_stream_name(self, acquisition_type: AcquisitionType) -> List[str]:
        """Convert AcquisitionType value to its equivalent EventStreamName objects.

        PS: On Bluesky, the existing DataStreams can be mapped to a respective AcquisitionType. Current mapping is:
            * SNAPSHOT and VARY-> 'baseline'
            * SCAN -> All streams except 'baseline'

        Args:
            acquisition_type (AcquisitionType): AcquisitionType Enum object to be converted

        Raises:
            AssonantDataRetrieverError: Raised if there is no valid convertion for the passed AcquisitionType object.

        Returns:
            List[str]: List of stream names related to the passed Acquisition Type Enum object.
        """
        if acquisition_type in [AcquisitionType.SNAPSHOT_AT_START, AcquisitionType.SNAPSHOT_AT_END]:
            return ["baseline"]
        elif acquisition_type == AcquisitionType.SCAN:
            stream_names = self.bluesky_data_parser.get_stream_names()
            if "baseline" in stream_names:
                stream_names.remove("baseline")
            return stream_names  # Except baseline stream, all other streams fits SCAN AcquisitionType.
        else:
            raise AssonantDataRetrieverError(
                f"There is no valid EventStream convertion value for '{acquisition_type.value}' AcquisitionType"
            )

    def _merge_by_seq_num(
        self, input_dict: Dict[int, Dict[str, List[Any]]]
    ) -> Dict[str, Dict[str, Union[Any, List[Any]]]]:
        """Merge retrieved data by sequence number.

        Output structure:
        {
            PV_NAME_1: {
                'data': {
                    'value': [PV_VALUE_SEQ_NUM_1, PV_VALUE_SEQ_NUM_2, ...]
                    'timestamps': [PV_TIMESTAMP_SEQ_NUM_1, PV_TIMESTAMP_SEQ_NUM_2, ...]
                },
                'configs': PV_CONFIGS
            },
            ...
        }

        Args:
            input_dict (Dict[int, Dict[str, List[Any]]]): Dict returned by BlueskyDataParser
            with data stored by sequence number and PV name. Check BlueskyDataParser
            for details on the structure of the input dict.

        Returns:
            Dict[str, Dict[str, Union[Any, List[Any]]]]: Dict with information merged
            as the output structure described above.
        """
        merged = {}

        # Sort sequences
        for seq_num in sorted(input_dict.keys()):
            seq_data = input_dict[seq_num]

            for pv_name, pv_info in seq_data.items():
                value = pv_info["data"]["value"]
                timestamp = pv_info["data"]["timestamps"]
                configs = pv_info["configs"]

                if pv_name not in merged:
                    merged[pv_name] = {"data": {"value": [], "timestamps": []}, "configs": configs}

                merged[pv_name]["data"]["value"].append(value)
                merged[pv_name]["data"]["timestamps"].append(timestamp)

        return merged

    def get_pv_data_by_acquisition_type(self, acquisition_type: AcquisitionType) -> Dict[str, Union[Any, List[Any]]]:
        """Return collected PV data related to specified AcquisitionType.

        The returned value is a dictionary as follow:

        {
            PV_NAME_1: [PV_VALUE_1, PV_VALUE_2, ...],
            PV_NAME_2: [PV_VALUE_1],
            ...
        }

        PS: PV values are always returned within a list, even if it is a single value.

        Args:
            acquisition_type (AcquisitionType): Target Acquisition Type used to select
            data from which Acquisition Type will be retrieved.

        Returns:
            Dict[str, Union[Any, List[Any]]]: Dictionary containing retrieved PV data.
        """
        event_stream_names = self._convert_acquisition_type_into_event_stream_name(acquisition_type)
        data_retrived_by_query = {}

        # Cache retrieved data to avoid multiple retrievals from the same data source
        if self.retrieved_data is None:
            self.retrieved_data = self.bluesky_data_parser.query_all_data(with_config=True)

        # Filter data to fit the requested streams based on AcquisitionType
        for event_stream_name in event_stream_names:

            retrieved_stream_data = self.retrieved_data.get(event_stream_name, {event_stream_name: {}})
            # Special processing for SNAPSHOT cases
            if event_stream_name == "baseline":
                if acquisition_type == AcquisitionType.SNAPSHOT_AT_START:
                    # Instead of directly indexing seq_num, get the lower seq_num value may change
                    # depending in the order of acquisition
                    first_seq_num = min(retrieved_stream_data.keys())
                    collapsed_retrieved_stream_data = self._merge_by_seq_num(
                        {first_seq_num: retrieved_stream_data[first_seq_num]}
                    )
                elif acquisition_type == AcquisitionType.SNAPSHOT_AT_END:
                    # Instead of directly indexing seq_num, get the higher seq_num returned as seq_num value may change
                    # depending in the order of acquisition
                    last_seq_num = max(retrieved_stream_data.keys())
                    collapsed_retrieved_stream_data = self._merge_by_seq_num(
                        {last_seq_num: retrieved_stream_data[last_seq_num]}
                    )
                else:
                    raise AssonantDataRetrieverError(
                        f"Invalid AcquisitionType '{acquisition_type.value}' for 'baseline' Event Stream"
                    )
            else:
                collapsed_retrieved_stream_data = self._merge_by_seq_num(retrieved_stream_data)

            data_retrived_by_query = data_retrived_by_query | collapsed_retrieved_stream_data

        # Convert data representation to fit the expected return structure
        formatted_retrieved_data = {}

        for pv_name, pv_info_dict in data_retrived_by_query.items():
            formatted_retrieved_data[pv_name] = pv_info_dict["data"]["value"]

        return formatted_retrieved_data
