# (C) Copyright 2019- ECMWF.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

import collections
from typing import Any
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import Mapping
from typing import MutableMapping
from typing import Sequence
from typing import Set
from typing import Union

import eccodes  # type: ignore
import numpy as np
import pandas as pd  # type: ignore

from pdbufr.core.filters import BufrFilter
from pdbufr.core.filters import filters_match
from pdbufr.core.keys import COMPUTED_KEYS
from pdbufr.core.keys import UncompressedBufrKey
from pdbufr.core.structure import MessageWrapper

from . import Reader


def extract_message(
    message: Mapping[str, Any],
    filters: Dict[str, BufrFilter] = {},
    base_observation: Dict[str, Any] = {},
    required_columns: Set[str] = set(),
    header_keys: Set[str] = set(),
) -> Iterator[Dict[str, Any]]:
    value_cache = {}
    try:
        is_compressed = bool(message["compressedData"])
    except KeyError:
        is_compressed = False

    if is_compressed:
        is_uncompressed = False
        subset_count = message["numberOfSubsets"]
    else:
        is_uncompressed = int(message["numberOfSubsets"]) > 1
        subset_count = 1

    # For messages with uncompressed subsets constnider this:
    # - for each data key we have a single value
    # - there is no way to identify the subset from the key
    # - we cannot directly iterate over a given subset
    # - if we iterate over the keys a new subset is indicated by the
    #   appearance of the "subsetNumber" key, which contains the same array
    #   of values each time (the subset index for all the subsets). This key is
    #   generated by ecCodes and does not contain any ranking so its name is
    #   always "subsetNumber".

    uncompressed_keys: Dict["str", Any]
    uncompressed_keys = dict()
    skip_keys = {
        "unexpandedDescriptors",
        "shortDelayedDescriptorReplicationFactor",
        "delayedDescriptorReplicationFactor",
        "extendedDelayedDescriptorReplicationFactor",
        "delayedDescriptorAndDataRepetitionFactor",
        "extendedDelayedDescriptorAndDataRepetitionFactor" "associatedFieldSignificance",
        "dataPresentIndicator",
        "operator",
    }

    for subset in range(subset_count):
        filters_match = {k: False for k in filters.keys()}
        required_columns_match = {k: False for k in required_columns}
        current_observation: Dict[str, Any]
        current_observation = collections.OrderedDict(base_observation)

        uncompressed_subset = 0
        for key in message:
            name = key.rpartition("#")[2]
            if name in skip_keys or "->" in key:
                continue

            if is_uncompressed and key == "subsetNumber":
                if uncompressed_subset > 0:
                    if (
                        current_observation
                        and all(filters_match.values())
                        and all(required_columns_match.values())
                    ):
                        yield dict(current_observation)

                    # header keys appear only once so we need to keep the match info for them
                    for k in filters_match:
                        if k not in header_keys:
                            filters_match[k] = False

                    for k in required_columns_match:
                        if k not in header_keys:
                            required_columns_match[k] = False

                    # only keep the header keys in the result because they are the
                    # same for all the subsets
                    while len(current_observation) and "subsetNumber" in current_observation:
                        # OrderedDict.popitem uses LIFO orde
                        current_observation.popitem()

                    for v in uncompressed_keys.values():
                        v.adjust_ref_rank()

                uncompressed_subset += 1

            if is_compressed:
                if key not in value_cache:
                    try:
                        value_cache[key] = message[key]
                    except KeyError:
                        value_cache[key] = None
                value = value_cache[key]
            else:
                value = message[key]

            # extract compressed BUFR values. They are either numpy arrays (for numeric types)
            # or lists of strings
            if (
                is_compressed
                and name != "unexpandedDescriptors"
                and isinstance(value, (np.ndarray, list))
                and len(value) == subset_count
            ):
                value = value[subset]

            if isinstance(value, float) and value == eccodes.CODES_MISSING_DOUBLE:
                value = None
            elif isinstance(value, int) and value == eccodes.CODES_MISSING_LONG:
                value = None

            if is_uncompressed:
                if name not in uncompressed_keys:
                    uncompressed_keys[name] = UncompressedBufrKey.from_key(key)
                else:
                    uncompressed_keys[name].update_rank(key)
                key = uncompressed_keys[name].relative_key

                # subsetNumber is an array and we need the current value
                if key == "subsetNumber":
                    value = uncompressed_subset

            if name in filters:
                if filters[name].match(value):
                    filters_match[name] = True

            if name in required_columns:
                required_columns_match[name] = True

            current_observation[key] = value

        # yield the last observation
        if current_observation and all(filters_match.values()) and all(required_columns_match.values()):
            yield dict(current_observation)


def test_computed_keys(
    observation: Dict[str, Any],
    filters: Dict[str, BufrFilter] = {},
    prefix: str = "",
) -> bool:
    for keys, computed_key, getter in COMPUTED_KEYS:
        if computed_key in filters:
            computed_value = None
            try:
                computed_value = getter(observation, prefix, keys)
            except Exception:
                return False
            if computed_value is not None:
                if not filters[computed_key].match(computed_value):
                    return False
            else:
                return False
    return True


class FlatReader(Reader):
    class ColumnInfo:
        def __init__(self) -> None:
            self.first_count = 0

    def __init__(
        self,
        path_or_messages,
        columns: Union[Sequence[str], str] = [],
        **kwargs: Any,
    ):
        class ColumnInfo:
            def __init__(self) -> None:
                self.first_count = 0

        self.column_info = ColumnInfo()
        super().__init__(path_or_messages, columns=columns, column_info=self.column_info, **kwargs)

    def execute(
        self,
        # columns: Union[Sequence[str], str] = [],
        # filters: Mapping[str, Any] = {},
        # required_columns: Union[bool, Iterable[str]] = True,
        # **kwargs,
    ) -> pd.DataFrame:
        # class ColumnInfo:
        #     def __init__(self) -> None:
        #         self.first_count = 0

        # column_info = ColumnInfo()
        df = super().execute(
            # columns=columns,
            # filters=filters,
            # required_columns=required_columns,
            # column_info=column_info,
            # **kwargs,
        )

        # compare the column count in the first record to that of the
        # dataframe. If the latter is larger, then there were non-aligned columns,
        # which were appended to the end of the dataframe columns.
        if self.column_info.first_count > 0 and self.column_info.first_count < len(df.columns):
            import warnings

            # temporarily overwrite warnings formatter
            ori_formatwarning = warnings.formatwarning
            warnings.formatwarning = lambda msg, *args, **kwargs: f"Warning: {msg}\n"
            warnings.warn(
                (
                    "not all BUFR messages/subsets have the same structure in the input file. "
                    "Non-overlapping columns (starting with column[{column_info.first_count-1}] ="
                    f"{df.columns[self.column_info.first_count-1]}) were added to end of the resulting dataframe"
                    "altering the original column order for these messages."
                )
            )
            warnings.formatwarning = ori_formatwarning

        return df

    def read_records(
        self,
        bufr_obj: Iterable[MutableMapping[str, Any]],
        columns: Union[Sequence[str], str],
        filters: Mapping[str, Any] = {},
        required_columns: Union[bool, Iterable[str]] = True,
        prefilter_headers: bool = False,
        column_info: Any = None,
    ) -> Iterator[Dict[str, Any]]:
        if isinstance(columns, str):
            columns = (columns,)

        if not isinstance(columns, Sequence):
            raise TypeError("invalid columns type")
        elif len(columns) == 0 or columns[0] == "":
            columns = ("all",)
        elif len(columns) != 1:
            raise ValueError("when columns is an iterable it can have maximum 1 element")

        if columns[0] not in ["all", "header", "data"]:
            raise ValueError("columns must be all, header or data")

        add_header = columns[0] in ["all", "header"]
        add_data = columns[0] in ["all", "data"]
        if not add_header and not add_data:
            raise ValueError("either header or data must be extracted")

        if isinstance(required_columns, bool):
            required_columns = set()
        elif isinstance(required_columns, str):
            required_columns = {required_columns}
        elif isinstance(required_columns, Iterable):
            required_columns = set(required_columns)
        else:
            raise TypeError("required_columns must be a bool, str or an iterable")

        # compile filters
        filters = dict(filters)
        value_filters = {k: BufrFilter.from_user(filters[k], key=k) for k in filters}

        # prepare count filter
        if "count" in value_filters:
            max_count = value_filters["count"].max()
        else:
            max_count = None

        count_filter = value_filters.pop("count", None)

        # prepare computed keys
        computed_keys = [x for _, x, _ in COMPUTED_KEYS]
        value_filters_without_computed = {k: v for k, v in value_filters.items() if k not in computed_keys}

        if column_info is not None:
            column_info.first_count = 0

        for count, msg in enumerate(bufr_obj, 1):
            # we use a context manager to automatically delete the handle of the BufrMessage.
            # We have to use a wrapper object here because a message can also be a dict
            with MessageWrapper.wrap(msg) as message:
                # count filter
                if count_filter is not None and not count_filter.match(count):
                    continue

                message_value_filters = value_filters_without_computed
                message_required_columns = required_columns

                header_keys = set()
                if not add_header or prefilter_headers or value_filters or required_columns:
                    header_keys = set(message.keys())

                    if required_columns:
                        message_required_columns = required_columns - set(header_keys)

                    # test header keys for failed matches before unpacking
                    if prefilter_headers:
                        if not filters_match(message, value_filters, required=False):
                            continue
                        # remove already tested filters
                        else:
                            message_value_filters = {
                                k: v for k, v in value_filters.items() if k not in header_keys
                            }

                message["skipExtraKeyAttributes"] = 1

                if add_data or message_value_filters or message_required_columns:
                    message["unpack"] = 1

                observation: Dict[str, Any] = {}

                for observation in extract_message(
                    message,
                    message_value_filters,
                    observation,
                    message_required_columns,
                    header_keys,
                ):
                    if header_keys:
                        if not add_header:
                            for key in header_keys:
                                observation.pop(key, None)
                        if not add_data:
                            data_keys = set(observation.keys()) - header_keys
                            for key in data_keys:
                                observation.pop(key, None)

                    if observation and test_computed_keys(observation, value_filters, "#1#"):
                        if column_info is not None and column_info.first_count == 0:
                            column_info.first_count = len(observation)
                        yield observation

                # optimisation: skip decoding messages above max_count
                if max_count is not None and count >= max_count:
                    break

    def adjust_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
        return df


reader = FlatReader
