import json
from typing import Any, Dict, List, Union

from event_model import DocumentRouter

from .exceptions import BlueskyDataParserError


class DocumentsDataParser(DocumentRouter):
    def __init__(self, bluesky_documents_list: List[List[Union[str, Dict[str, Any]]]], *args, **kwargs):
        """Constructor to instantiate a Bluesky Documents DataParser.

        Args:
            bluesky_documents_list (List[List[Union[str, Dict[str, Any]]]]): 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:
            BlueskyDataParserError: Failure to read bluesky documents.
        """
        super().__init__(*args, **kwargs)

        self._start_document = None
        self._descriptor_documents = {}

        # Stream name --> Subdevice configurations
        self._configurations = {}
        # Stream name --> Device data keys
        self._data_keys = {}
        # Stream name --> Device data key --> (Device data reading, Reading timestamp)
        self._event_data = {}

        try:
            # print("Init DocumentsDataParser")
            # Iterate over json data correctly parse data based on event type
            for event_name, document_data_dict in bluesky_documents_list:
                # print(event_name, document_data_dict)
                self(event_name, document_data_dict)
        except Exception as e:
            raise BlueskyDataParserError(
                f"Failed to read bluesky documents from structure: {bluesky_documents_list}"
            ) from e

    def start(self, doc: Dict):
        """DocumentRouter method related to processing bluesky START documments

        Args:
            doc (Dict): Dictionary containing bluesky documment data.
        """
        # print("========== START ==========")
        # print(doc)
        # print("===========================\n")
        self._start_document = doc

    def descriptor(self, doc: Dict):
        """DocumentRouter method related to processing bluesky DESCRIPTOR documments

        Args:
            doc (Dict): Dictionary containing bluesky documment data.
        """
        # print("========== DESC ==========")
        # print(doc)
        # print("===========================\n")
        self._descriptor_documents[doc["uid"]] = doc

        self._configurations[doc["name"]] = doc["configuration"]
        self._data_keys[doc["name"]] = doc["data_keys"]
        # print("DESC DOCS")
        # print(self._descriptor_documents)

    def event(self, doc: Dict):
        """DocumentRouter method related to processing bluesky EVENT documments

        Args:
            doc (Dict): Dictionary containing bluesky documment data.
        """
        # print("========== EVENT ==========")
        # print(doc)
        # print("===========================\n")
        stream_name = self._descriptor_documents[doc["descriptor"]]["name"]

        if stream_name not in self._event_data:
            self._event_data[stream_name] = {}

        if doc["seq_num"] not in self._event_data[stream_name]:
            self._event_data[stream_name][doc["seq_num"]] = {}

        for dev_name, val in doc["data"].items():
            self._ensure_dict_key(dev_name, self._event_data[stream_name][doc["seq_num"]])

            self._event_data[stream_name][doc["seq_num"]][dev_name] = {
                "value": val,
                "timestamps": doc["timestamps"][dev_name],
            }

    def get_stream_names(self) -> List[str]:
        """Return a List of streams names within parsed documents.

        Returns:
            List[str]: List of stream names found on parsed documents.
        """
        return [stream_name for stream_name in self._data_keys]

    def all_data_keys(self) -> Dict[str, List]:
        """Return all dictionary keys to access data fields from each stream into the data dictionary.

        PS: If no device_data_key is present for a stream, it will return an empty list for it.

        Returns:
            Dict[str, List]: Dictionary containing the stream name and a List with its respetive data_keys
        """
        data_keys = {}
        for stream, devices in self._data_keys.items():
            # Even if stream has no device within it, return an empty list for it
            if stream not in data_keys:
                data_keys[stream] = []
            for device_data_key, device_values in devices.items():
                data_keys[stream].append(device_data_key)
        return data_keys

    def data_key_to_pv_name(self, data_key: str) -> str:
        """Convert data_key to the respective PV name related to it.

        Args:
            data_key (str): data key from data dictionary that will be converted

        Returns:
            str: PV Name respective to passed data_key
        """
        pv_names = {}
        for stream, devices in self._data_keys.items():
            for device_data_key, device_values in devices.items():
                if device_data_key == data_key:
                    pv_names[stream] = ":".join(device_values["source"].split(":")[1:])
                    break
        return pv_names

    def _ensure_dict_key(self, key_name: Any, target_dict: Dict) -> None:
        """
        Ensures that a specified key exists in the given dictionary.
        If the key does not exist, it initializes it with an empty dictionary.

        Args:
            key_name (Any): The key to ensure exists in the dictionary.
            target_dict (Dict): The dictionary to check and modify.
        """
        if key_name not in target_dict:
            target_dict[key_name] = {}

    def config_data_key_to_pv_name(self, data_key: str) -> str:
        """Convert a configuration data_key to the respective PV name related to it.

        Args:
            data_key (str): data key from configuration dictionary from descriptor event
            that will be converted

        Returns:
            str: PV Name respective to passed configuration data_key
        """
        pv_names = {}
        for stream, devices in self._configurations.items():
            for device_name in devices:
                for config_data_key, config_values in devices[device_name]["data_keys"].items():
                    if config_data_key == data_key:
                        pv_names[stream] = ":".join(config_values["source"].split(":")[1:])
                        break
        return pv_names

    def pv_name_to_data_key(self, pv_name: str) -> str:
        """Convert PV name to the respective data_key related to it.

        Args:
            pv_name (str): PV Name that will be converted

        Returns:
            str: Data key respective to passed PV Name
        """
        data_keys = {}
        formatted_pv_name = f"PV:{pv_name}"
        for stream, devices in self._data_keys.items():
            for device_data_key, device_values in devices.items():
                if formatted_pv_name in device_values["source"]:
                    data_keys[stream] = device_data_key
                    break

        return data_keys

    def pv_name_to_config_data_key(self, pv_name: str) -> str:
        """Convert PV name to the respective configuration data_key related to it.

        Args:
            pv_name (str): PV Name that will be converted

        Returns:
            str: Configuration Data key respective to passed PV Name
        """
        data_keys = {}
        formatted_pv_name = f"PV:{pv_name}"
        for stream, devices in self._configurations.items():
            for device_name in devices:
                for config_data_key, config_values in devices[device_name]["data_keys"].items():
                    if formatted_pv_name in config_values["source"]:
                        data_keys[stream] = config_data_key
                        break

        return data_keys

    def query_config_by_pv_name(self, pv_name: str, stream_name: str) -> Dict[str, Dict[str, Any]]:
        """Return configuration information related to a PV name.

        The returned dictionary structure is the following:
            {
                CONFIG_PV_NAME_1: {
                        "value": CONFIG_VALUE,
                        "timestamps": CONFIG_VALUE_ACQ_TIMESTAMP
                },
                CONFIG_PV_NAME_2: {
                        "value": CONFIG_VALUE,
                        "timestamps": CONFIG_VALUE_ACQ_TIMESTAMP
                    }
                    ...
            }

        Args:
            pv_name (str): PV name that will be used to query configuration information.
            stream_name (str): Stream name that will be used to query configuration information.

        Returns:
            Dict[str, Dict[str, Any]]: Dictionary containing configuration information related to passed PV name
            structured as described above. If no config is found it returns an empty dict.

        """

        config_data = {}

        if stream_name not in self._configurations:
            return config_data

        for device_name in self._configurations[stream_name]:
            for config_data_key, config_values in self._configurations[stream_name][device_name]["data_keys"].items():
                config_pv_name = ":".join(config_values["source"].split(":")[1:])

                self._ensure_dict_key(config_pv_name, config_data)

                config_data[config_pv_name]["value"] = self._configurations[stream_name][device_name]["data"][
                    config_data_key
                ]
                config_data[config_pv_name]["timestamps"] = self._configurations[stream_name][device_name][
                    "timestamps"
                ][config_data_key]

        return config_data

    def query_data_by_stream_name(
        self, event_stream_name: str, with_config: bool = True
    ) -> Dict[int, Dict[str, Dict[str, Any]]]:
        """Return all information from bluesky docs from a specific event stream as described below.

        The organization structure returned is the following:
        {
            SEQ_NUM_1: {
                PV_NAME_A: {
                    "data": {"value": PV_VALUE_1, "timestamps": PV_VALUE_1_TIMESTAMP},
                    "configs": <Structure from query_config_by_pv_name() method>
                },
                PV_NAME_B: {
                    "data": {"value": PV_VALUE_1, "timestamps": PV_VALUE_1_TIMESTAMP},
                    "configs": <Structure from query_config_by_pv_name() method>
                },
            },
            SEQ_NUM_2: {
                PV_NAME_A: {
                        "data": {"value": PV_VALUE_2, "timestamps": PV_VALUE_2_TIMESTAMP},
                        "configs": <Structure from query_config_by_pv_name() method>
                    },
                    PV_NAME_B: {
                        "data": {"value": PV_VALUE_2, "timestamps": PV_VALUE_2_TIMESTAMP},
                        "configs": <Structure from query_config_by_pv_name() method>
                    },
        }

        PS: If stream exist but no device data is found, it returns an empty dict.
        PS: PV values may be list

        Args:
            event_stream_name (str): Query's target event stream name.
            with_config (bool): Flag to control if data will add configuration
            information or not. Defaults to True.

        Returns:
            Dict[int, Dict[str, Dict[str, Union[List[Dict[str, Any]], Dict[str, Any]]]]]: Dictionary
            with data following structure described above. If no data is found it returns an empty dict.
        """
        # Stream name --> Seq_num --> PV prefix --> Collected data
        # data = {}

        data = {}

        data_keys = self.all_data_keys()
        if len(data_keys) == 0 or event_stream_name not in data_keys.keys():
            return data

        data_keys = data_keys[event_stream_name]

        for data_key in data_keys:
            pv_name = self.data_key_to_pv_name(data_key)[event_stream_name]

            device_data = self._event_data.get(event_stream_name, {})
            for seq_num in device_data:

                self._ensure_dict_key(seq_num, data)
                self._ensure_dict_key(pv_name, data[seq_num])
                data[seq_num][pv_name]["data"] = device_data[seq_num].get(data_key, {})

                if with_config:
                    data[seq_num][pv_name]["configs"] = self.query_config_by_pv_name(pv_name, event_stream_name)

        return data

    def query_all_data(
        self, with_config: bool = True
    ) -> Dict[str, Dict[int, Dict[str, Dict[str, Union[List[Dict[str, Any]], Dict[str, Any]]]]]]:
        """Return all information from bluesky docs as described below.

        The organization structure returned is the following:
        {
            STREAM_NAME_1: <Structure from query_data_by_stream_name() method>,
            STREAM_NAME_2: <Structure from query_data_by_stream_name() method>,
            ...
        }

        PS: PV values are ALWAYS inside a List, even if it is a single value. It is also ALWAYS
        returned in a Tuple where the first position is the acquired value and the second the
        acquisition timestamp.

        Args:
            with_config (bool): Flag to control if data will add configuration
            information or not. Defaults to True.

        Returns:
            Dict[str, Dict[int, Dict[str, Dict[str, Union[List[Dict[str, Any]], Dict[str, Any]]]]]]: Dictionary
            with data following structure described above. If no data is found it returns an empty dict.
        """
        # Stream name --> structured from query_data_by_stream_name()
        data = {}

        data_keys = self.all_data_keys()
        if len(data_keys) == 0:
            return data

        for stream in data_keys.keys():
            # print(f"Stream '{stream}':")
            data[stream] = {}
            stream_data = self.query_data_by_stream_name(stream, with_config=with_config)
            data[stream] = stream_data
        return data

    def load_json(self, path: str) -> Dict[str, Any]:
        """Load data from passed json file

        Args:
            path (str): File path to json file containing data collect from bluesky.

        Returns:
            Dict[str, Any]: Dictionary containing retrieved data from json file.
        """
        contents = None
        with open(path, "r") as _f:
            contents = json.load(_f)

        return contents
