""" This file contains the BaseStore class from which other storage
    modules inherit.
"""
from __future__ import annotations

import logging
from types import TracebackType
from typing import Any, Dict, List, Optional, Sequence, TypeVar, Union

from pandas import Timestamp

from openghg.objectstore import get_object_from_json, exists, set_object_from_json
from openghg.objectstore.metastore import DataClassMetaStore
from openghg.types import DatasourceLookupError
from openghg.util import timestamp_now, to_lowercase


T = TypeVar("T", bound="BaseStore")


logger = logging.getLogger("openghg.store")
logger.setLevel(logging.DEBUG)  # Have to set level for logger as well as handler


class BaseStore:
    _registry: dict[str, type[BaseStore]] = {}
    _data_type = ""
    _root = "root"
    _uuid = "root_uuid"

    def __init__(self, bucket: str) -> None:
        self._creation_datetime = str(timestamp_now())
        self._stored = False
        # Hashes of previously uploaded files
        self._file_hashes: Dict[str, str] = {}
        # Hashes of previously stored data from other data platforms
        self._retrieved_hashes: Dict[str, Dict] = {}
        # Where we'll store this object's metastore
        self._metakey = ""

        if exists(bucket=bucket, key=self.key()):
            data = get_object_from_json(bucket=bucket, key=self.key())
            # Update myself
            self.__dict__.update(data)

        self._metastore = DataClassMetaStore(bucket=bucket, data_type=self._data_type)
        self._bucket = bucket
        self._datasource_uuids = self._metastore.select("uuid")

    def __init_subclass__(cls) -> None:
        BaseStore._registry[cls._data_type] = cls

    def __enter__(self) -> BaseStore:
        return self

    def __exit__(
        self,
        exc_type: Optional[BaseException],
        exc_val: Optional[BaseException],
        exc_tb: Optional[TracebackType],
    ) -> None:
        if exc_type is not None:
            logger.error(msg=f"{exc_type}, {exc_tb}")
        else:
            self.save()

    @classmethod
    def metakey(cls) -> str:
        return str(cls._metakey)

    @classmethod
    def key(cls) -> str:
        return f"{cls._root}/uuid/{cls._uuid}"

    def save(self) -> None:
        self._metastore.close()
        set_object_from_json(bucket=self._bucket, key=self.key(), data=self.to_data())

    def to_data(self) -> Dict:
        # We don't need to store the metadata store, it has its own location
        # QUESTION - Is this cleaner than the previous specifying
        DO_NOT_STORE = ["_metastore", "_bucket", "_datasource_uuids"]
        return {k: v for k, v in self.__dict__.items() if k not in DO_NOT_STORE}

    def read_data(self, *args: Any, **kwargs: Any) -> Optional[dict]:
        raise NotImplementedError

    def read_file(self, *args: Any, **kwargs: Any) -> dict:
        raise NotImplementedError

    def store_data(self, *args: Any, **kwargs: Any) -> Optional[dict]:
        raise NotImplementedError

    def transform_data(self, *args: Any, **kwargs: Any) -> dict:
        raise NotImplementedError

    def assign_data(
        self,
        data: Dict,
        overwrite: bool,
        data_type: str,
        required_keys: Sequence[str],
        min_keys: Optional[int] = None,
        update_keys: Optional[List] = None,
    ) -> Dict[str, Dict]:
        """Assign data to a Datasource. This will either create a new Datasource
        Create or get an existing Datasource for each gas in the file

            Args:
                data: Dictionary containing data and metadata for species
                overwrite: If True overwrite current data stored
                data_type: Type of data, timeseries etc
                required_keys: Required minimum keys to lookup unique Datasource
                min_keys: Minimum number of metadata keys needed to uniquely match a Datasource
            Returns:
                dict: Dictionary of UUIDs of Datasources data has been assigned to keyed by species name
        """
        from openghg.store.base import Datasource

        uuids = {}

        self._metastore.acquire_lock()
        try:
            lookup_results = self.datasource_lookup(data=data, required_keys=required_keys, min_keys=min_keys)
            # TODO - remove this when the lowercasing of metadata gets removed
            # We currently lowercase all the metadata and some keys we don't want to change, such as paths to the object store
            skip_keys = ["object_store"]

            for key, parsed_data in data.items():
                metadata = parsed_data["metadata"]
                _data = parsed_data["data"]

                # Our lookup results and gas data have the same keys
                uuid = lookup_results[key]

                # Add the read metadata to the Dataset attributes being careful
                # not to overwrite any attributes that are already there
                def convert_to_netcdf4_types(value: Any) -> Union[int, float, str, list]:
                    """Attributes in a netCDF file can be strings, numbers, or sequences:
                    http://unidata.github.io/netcdf4-python/#attributes-in-a-netcdf-file

                    This function converts any data whose type is not int, float, str, or list
                    to strings.
                    Booleans are converted to strings, even though they are a subtype of int.
                    """
                    if isinstance(value, (int, float, str, list)) and not isinstance(value, bool):
                        return value
                    else:
                        return str(value)

                to_add = {k: convert_to_netcdf4_types(v) for k, v in metadata.items() if k not in _data.attrs}
                _data.attrs.update(to_add)

                # If we have a UUID for this Datasource load the existing object
                # from the object store
                # If we haven't stored data with this metadata before we create a new Datasource
                # and add the metadata to our metastore

                # Take a copy of the metadata so we can update it
                meta_copy = metadata.copy()
                new_ds = uuid is False

                if new_ds:
                    datasource = Datasource()
                    uid = datasource.uuid()
                    meta_copy["uuid"] = uid
                    # Make sure all the metadata is lowercase for easier searching later
                    # TODO - do we want to do this or should be just perform lowercase comparisons?
                    meta_copy = to_lowercase(d=meta_copy, skip_keys=skip_keys)
                else:
                    datasource = Datasource.load(bucket=self._bucket, uuid=uuid)

                # Add the dataframe to the datasource
                datasource.add_data(
                    metadata=meta_copy,
                    data=_data,
                    overwrite=overwrite,
                    data_type=data_type,
                    skip_keys=skip_keys,
                )
                # Save Datasource to object store
                datasource.save(bucket=self._bucket)

                # Add the metadata to the metastore and make sure it's up to date with the metadata stored
                # in the Datasource
                datasource_metadata = datasource.metadata()
                if new_ds:
                    self._metastore.insert(datasource_metadata)
                else:
                    self._metastore.update(where={"uuid": datasource.uuid()}, to_update=datasource_metadata)

                uuids[key] = {"uuid": datasource.uuid(), "new": new_ds}
        finally:
            self._metastore.release_lock()

        return uuids

    def datasource_lookup(
        self, data: Dict, required_keys: Sequence[str], min_keys: Optional[int] = None
    ) -> Dict:
        """Search the metadata store for a Datasource UUID using the metadata in data. We expect the required_keys
        to be present and will require at least min_keys of these to be present when searching.

        As some metadata value might change (such as data owners etc) we don't want to do an exact
        search on *all* the metadata so we extract a subset (the required keys) and search for these.

        Args:
            metastore: Metadata database
            data: Combined data dictionary of form {key: {data: Dataset, metadata: Dict}}
            required_keys: Iterable of keys to extract from metadata
            min_keys: The minimum number of required keys, if not given it will be set
            to the length of required_keys
        Return:
            dict: Dictionary of datasource information
        """
        from openghg.util import to_lowercase

        if min_keys is None:
            min_keys = len(required_keys)

        results = {}
        for key, _data in data.items():
            metadata = _data["metadata"]

            required_metadata = {
                k.lower(): to_lowercase(v) for k, v in metadata.items() if k in required_keys
            }

            if len(required_metadata) < min_keys:
                raise ValueError(
                    f"The given metadata doesn't contain enough information, we need: {required_keys}"
                )

            required_result = self._metastore.search(required_metadata)

            if not required_result:
                results[key] = False
            elif len(required_result) > 1:
                raise DatasourceLookupError("More than once Datasource found for metadata, refine lookup.")
            else:
                results[key] = required_result[0]["uuid"]

        return results

    def uuid(self) -> str:
        """Return the UUID of this object

        Returns:
            str: UUID of object
        """
        return self._uuid

    def datasources(self) -> List[str]:
        """Return the list of Datasources UUIDs associated with this object

        Returns:
            list: List of Datasource UUIDs
        """
        return self._datasource_uuids

    def get_rank(self, uuid: str, start_date: Timestamp, end_date: Timestamp) -> Dict:
        """Get the rank for the given Datasource for a given date range

        Args:
            uuid: UUID of Datasource
            start_date: Start date
            end_date: End date
        Returns:
            dict: Dictionary of rank and daterange covered by that rank
        """
        raise NotImplementedError("Ranking is being reworked and will be reactivated in a future release.")
        from collections import defaultdict
        from openghg.util import create_daterange_str, daterange_overlap

        if uuid not in self._rank_data:
            return {}

        search_daterange = create_daterange_str(start=start_date, end=end_date)

        rank_data = self._rank_data[uuid]

        ranked = defaultdict(list)
        # Check if this Datasource is ranked for the dates passed
        for daterange, rank in rank_data.items():
            if daterange_overlap(daterange_a=search_daterange, daterange_b=daterange):
                ranked[rank].append(daterange)

        return ranked

    def clear_rank(self, uuid: str) -> None:
        """Clear the ranking data for a Datasource

        Args:
            uuid: UUID of Datasource
        Returns:
            None
        """
        raise NotImplementedError("Ranking is being reworked and will be reactivated in a future release.")
        if uuid in self._rank_data:
            del self._rank_data[uuid]
            self.save()
        else:
            raise ValueError("No ranking data set for that UUID.")

    def set_rank(
        self,
        uuid: str,
        rank: Union[int, str],
        date_range: Union[str, List[str]],
        overwrite: Optional[bool] = False,
    ) -> None:
        """Set the rank of a Datasource associated with this object.

        This function performs checks to ensure multiple ranks aren't set for
        overlapping dateranges.

        Passing a daterange and rank to this function will overwrite any current
        daterange stored for that rank.

        Args:
            uuid: UUID of Datasource
            rank: Rank of data
            date_range: Daterange(s)
            overwrite: Overwrite current ranking data
        Returns:
            None
        """
        raise NotImplementedError("Ranking is being reworked and will be reactivated in a future release.")
        from copy import deepcopy

        from openghg.util import (
            combine_dateranges,
            daterange_contains,
            daterange_overlap,
            sanitise_daterange,
            split_encompassed_daterange,
            trim_daterange,
        )

        rank = int(rank)

        if not 1 <= rank <= 10:
            raise TypeError("Rank can only take values 1 to 10 (for unranked). Where 1 is the highest rank.")

        if not isinstance(date_range, list):
            date_range = [date_range]

        # Make sure the dateranges passed are correct and are tz-aware
        date_range = [sanitise_daterange(d) for d in date_range]
        # Combine in case we have overlappng dateranges
        date_range = combine_dateranges(date_range)

        # Used to store dateranges that need to be trimmed to ensure no daterange overlap
        to_update = []
        # Non-overlapping dateranges that can be stored directly
        to_add = []

        if uuid in self._rank_data:
            rank_data = self._rank_data[uuid]
            # Check this source isn't ranked differently for the same dates
            for new_daterange in date_range:
                overlap = False
                # Check for overlapping dateranges and add
                for existing_daterange, existing_rank in rank_data.items():
                    if daterange_overlap(daterange_a=new_daterange, daterange_b=existing_daterange):
                        overlap = True

                        if rank != existing_rank and overwrite:
                            # Save the daterange we need to update
                            to_update.append((existing_daterange, new_daterange))
                            continue
                        # If the ranks are the same we just want to combine the dateranges
                        elif rank == existing_rank:
                            to_combine = [new_daterange, existing_daterange]
                            combined = combine_dateranges(dateranges=to_combine)[0]
                            to_update.append((existing_daterange, combined))
                        else:
                            raise ValueError(
                                f"This datasource has rank {existing_rank} for dates that overlap the ones given. \
                                                Overlapping dateranges are {new_daterange} and {existing_daterange}"
                            )
                # Otherwise we just want to add the new daterange to the dict
                if not overlap:
                    to_add.append(new_daterange)

            # If we've got dateranges to update and ranks to overwrite we need to trim the
            # previous ranking daterange down so we don't have overlapping dateranges
            if to_update:
                # Here we first take a backup of the old ranking data, update
                # it and then write it back
                ranking_backup = deepcopy(rank_data)

                for existing, new in to_update:
                    # Remove the existing daterange key
                    # Here we pass if it doesn't exist as if we have multiple new dateranges
                    # that overlap the existing daterange it might have been deleted during
                    # a previous iteration
                    try:
                        del ranking_backup[existing]
                    except KeyError:
                        pass

                    # If we want to overwrite an existing rank we need to trim that daterange and
                    # rewrite it back to the dictionary
                    rank_copy = rank_data[existing]

                    if overwrite:
                        if existing == new:
                            ranking_backup[new] = rank_copy
                        # If the existing daterange contains the new daterange
                        # we need to split it into parts and save those
                        elif daterange_contains(container=existing, contained=new):
                            result = split_encompassed_daterange(container=existing, contained=new)

                            existing_start = result["container_start"]
                            ranking_backup[existing_start] = rank_copy

                            updated_new = result["contained"]
                            ranking_backup[updated_new] = rank

                            # We might only end up with two dateranges
                            try:
                                existing_end = result["container_end"]
                                ranking_backup[existing_end] = rank_copy
                            except KeyError:
                                pass
                        # If the new daterange contains the existing we can just overwrite it
                        elif daterange_contains(container=new, contained=existing):
                            ranking_backup[new] = rank
                        else:
                            trimmed = trim_daterange(to_trim=existing, overlapping=new)
                            ranking_backup[trimmed] = rank_copy
                            ranking_backup[new] = rank
                    elif rank_copy == rank:
                        # If we're not overwriting we just need to update to use the new combined
                        ranking_backup[new] = rank_copy

                self._rank_data[uuid] = ranking_backup

            # Finally, store the dateranges that didn't overlap
            for d in to_add:
                self._rank_data[uuid][d] = rank
        else:
            for d in date_range:
                self._rank_data[uuid][d] = rank

        self.save()

    def rank_data(self) -> Dict:
        """Return a dictionary of rank data keyed
        by UUID

            Returns:
                dict: Dictionary of rank data
        """
        raise NotImplementedError("Ranking is being reworked and will be reactivated in a future release.")
        rank_dict: Dict = self._rank_data.to_dict()
        return rank_dict

    def clear_datasources(self) -> None:
        """Remove all Datasources from the object

        Returns:
            None
        """
        self._datasource_uuids.clear()
        self._file_hashes.clear()


def get_data_class(data_type: str) -> type[BaseStore]:
    """Return data class corresponding to given data type.

    Args:
        data_type: one of "surface", "column", "emissions", "footprints",
    "boundary_conditions", or "eulerian_model"

    Returns:
        Data class, one of `ObsSurface`, `ObsColumn`, `Emissions`, `EulerianModel`,
    `Footprints`, `BoundaryConditions`.
    """
    try:
        data_class = BaseStore._registry[data_type]
    except KeyError:
        raise ValueError(f"No data class for data type {data_type}.")
    else:
        return data_class
