import copy
from collections.abc import KeysView, ValuesView, ItemsView, MutableMapping
from contextlib import contextmanager
from typing import Dict, Optional, Iterator, cast, List, Any, overload, OrderedDict

import msgpack
import zlib

from compressedfhir.utilities.compressed_dict.v1.compressed_dict_access_error import (
    CompressedDictAccessError,
)
from compressedfhir.utilities.compressed_dict.v1.compressed_dict_storage_mode import (
    CompressedDictStorageMode,
    CompressedDictStorageType,
)
from compressedfhir.utilities.json_serializers.type_preservation_serializer import (
    TypePreservationSerializer,
)
from compressedfhir.utilities.ordered_dict_to_dict_converter.ordered_dict_to_dict_converter import (
    OrderedDictToDictConverter,
)


class CompressedDict[K, V](MutableMapping[K, V]):
    """
    A dictionary-like class that supports flexible storage options.

    It can store data in raw format, MessagePack format, or compressed MessagePack format.
    """

    # use slots to reduce memory usage
    __slots__ = [
        "_storage_mode",
        "_working_dict",
        "_raw_dict",
        "_serialized_dict",
        "_properties_to_cache",
        "_cached_properties",
        "_length",
        "_transaction_depth",
    ]

    def __init__(
        self,
        *,
        initial_dict: Dict[K, V] | OrderedDict[K, V] | None = None,
        storage_mode: CompressedDictStorageMode,
        properties_to_cache: List[K] | None,
    ) -> None:
        """
        Initialize a dictionary with flexible storage options
        Args:
            initial_dict: Initial dictionary to store
            storage_mode: Storage method for dictionary contents
                - 'raw': Store as original dictionary
                - 'msgpack': Store as MessagePack serialized bytes
                - 'compressed_msgpack': Store as compressed MessagePack bytes
        """
        # Storage configuration
        self._storage_mode: CompressedDictStorageMode = storage_mode

        # Working copy of the dictionary during context
        self._working_dict: Optional[OrderedDict[K, V]] = None

        # Private storage options
        self._raw_dict: OrderedDict[K, V] = OrderedDict[K, V]()
        self._serialized_dict: Optional[bytes] = None

        self._properties_to_cache: List[K] | None = properties_to_cache

        self._cached_properties: Dict[K, V] = {}

        self._length: int = 0

        self._transaction_depth: int = 0

        # Populate initial dictionary if provided
        if initial_dict:
            # Ensure we use an OrderedDict to maintain original order
            initial_dict_ordered = (
                initial_dict
                if isinstance(initial_dict, OrderedDict)
                else OrderedDict[K, V](initial_dict)
            )
            self.replace(value=initial_dict_ordered)

    @contextmanager
    def transaction(self) -> Iterator["CompressedDict[K, V]"]:
        """
        Context manager to safely access and modify the dictionary.

        Deserializes the dictionary before entering the context
        Serializes the dictionary after exiting the context

        Raises:
            CompressedDictAccessError: If methods are called outside the context
        """
        try:
            self.start_transaction()

            # Yield the working dictionary
            yield self

        finally:
            self.end_transaction()

    def start_transaction(self) -> "CompressedDict[K, V]":
        """
        Starts a transaction.  Use transaction() for a contextmanager for simpler usage.
        """
        # Increment transaction depth
        self._transaction_depth += 1
        # Ensure working dictionary is ready on first entry
        if self._transaction_depth == 1:
            self.ensure_working_dict()

        return self

    def end_transaction(self) -> "CompressedDict[K, V]":
        """
        Ends a transaction.  Use transaction() for a context_manager for simpler usage.

        """
        # Decrement transaction depth
        self._transaction_depth -= 1
        # Only update serialized dict when outermost transaction completes
        if self._transaction_depth == 0:
            self._update_serialized_dict(current_dict=self._working_dict)

            # Clear the working dictionary
            self._working_dict = None

        return self

    def ensure_working_dict(self) -> None:
        """
        Ensures that the working dictionary is initialized and deserialized.

        """
        if not self._working_dict:
            self._working_dict = self.create_working_dict()

    def create_working_dict(self) -> OrderedDict[K, V]:
        working_dict: OrderedDict[K, V]
        # Deserialize the dictionary before entering the context
        if self._storage_mode.storage_type == "raw":
            # For raw mode, create a deep copy of the existing dictionary
            working_dict = self._raw_dict
        else:
            # For serialized modes, deserialize
            working_dict = (
                self._deserialize_dict(
                    serialized_dict_bytes=self._serialized_dict,
                    storage_type=self._storage_mode.storage_type,
                )
                if self._serialized_dict
                else OrderedDict[K, V]()
            )
            assert isinstance(working_dict, OrderedDict)
        return working_dict

    @staticmethod
    def _serialize_dict(
        *, dictionary: OrderedDict[K, V], storage_type: CompressedDictStorageType
    ) -> bytes:
        """
        Serialize entire dictionary using MessagePack

        Args:
            dictionary: Dictionary to serialize
            storage_type: Storage type to use for serialization

        Returns:
            Serialized bytes
        """
        assert isinstance(dictionary, OrderedDict)
        if storage_type == "compressed":
            # Serialize to JSON and compress with zlib
            json_str = TypePreservationSerializer.serialize(
                dictionary, separators=(",", ":")
            )
            return zlib.compress(
                json_str.encode("utf-8"), level=zlib.Z_BEST_COMPRESSION
            )

        # Serialize using MessagePack
        packed = msgpack.packb(
            dictionary,
            use_bin_type=True,  # Preserve string/bytes distinction
            use_single_float=True,  # More compact float representation
        )

        # Optional compression
        if storage_type == "compressed_msgpack":
            packed = zlib.compress(packed, level=zlib.Z_BEST_COMPRESSION)

        return packed

    @staticmethod
    def _deserialize_dict(
        *,
        serialized_dict_bytes: bytes,
        storage_type: CompressedDictStorageType,
    ) -> OrderedDict[K, V]:
        """
        Deserialize entire dictionary from MessagePack

        Args:
            serialized_dict_bytes: Serialized dictionary bytes

        Returns:
            Deserialized dictionary
        """
        assert serialized_dict_bytes is not None, "Serialized dictionary cannot be None"
        assert isinstance(serialized_dict_bytes, bytes)

        if storage_type == "compressed":
            # Decompress and parse JSON
            decompressed_bytes: bytes = zlib.decompress(serialized_dict_bytes)
            decoded_text: str = decompressed_bytes.decode("utf-8")
            # noinspection PyTypeChecker
            decompressed_dict = TypePreservationSerializer.deserialize(
                decoded_text, object_pairs_hook=OrderedDict
            )
            assert isinstance(decompressed_dict, OrderedDict)
            return cast(OrderedDict[K, V], decompressed_dict)

        # Decompress if needed
        to_unpack = (
            zlib.decompress(serialized_dict_bytes)
            if storage_type == "compressed_msgpack"
            else serialized_dict_bytes
        )

        # Deserialize
        unpacked_dict = msgpack.unpackb(
            to_unpack,
            raw=False,  # Convert to strings
            strict_map_key=False,  # Handle potential key type variations
        )
        unpacked_dict = (
            unpacked_dict
            if isinstance(unpacked_dict, OrderedDict)
            else OrderedDict[K, V](unpacked_dict)
        )
        assert isinstance(unpacked_dict, OrderedDict)
        return cast(
            OrderedDict[K, V],
            unpacked_dict,
        )

    def _get_dict(self) -> OrderedDict[K, V]:
        """
        Get the dictionary, deserializing if necessary

        Returns:
            Current dictionary state
        """

        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary access is only allowed within an transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to access the dictionary."
                f"You tried to access it with storage type {self._storage_mode.storage_type}."
            )

        if self._storage_mode.storage_type == "raw":
            return self._raw_dict

        # For non-raw modes, do not keep deserialized dict
        return self._working_dict

    def __getitem__(self, key: K) -> V:
        """
        Retrieve a value

        Args:
            key: Dictionary key

        Returns:
            Value associated with the key
        """

        if self._properties_to_cache and key in self._properties_to_cache:
            return self._cached_properties[key]

        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary access is only allowed within an transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to access the dictionary."
            )
        return self._get_dict()[key]

    def __setitem__(self, key: K, value: V) -> None:
        """
        Set a value

        Args:
            key: Dictionary key
            value: Value to store
        """
        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary modification is only allowed within an transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
            )

        # Update the working dictionary
        self._working_dict[key] = value

        if self._properties_to_cache and key in self._properties_to_cache:
            # Update the cached properties if the key is in the list
            self._cached_properties[key] = value

    def _update_serialized_dict(self, current_dict: OrderedDict[K, V] | None) -> None:
        if current_dict is None:
            self._cached_properties.clear()
            self._length = 0
            self._serialized_dict = None
            self._raw_dict = OrderedDict[K, V]()
            return

        if self._properties_to_cache:
            for key in self._properties_to_cache:
                if key in current_dict:
                    self._cached_properties[key] = current_dict[key]

        self._length = len(current_dict)

        if self._working_dict is not None:
            # If the working dictionary is None, initialize it
            self._working_dict = current_dict

        if self._transaction_depth == 0:
            # If we're in a transaction,
            # The serialized dictionary will be updated after the transaction
            if self._storage_mode.storage_type == "raw":
                self._raw_dict = current_dict
            else:
                self._serialized_dict = (
                    self._serialize_dict(
                        dictionary=current_dict,
                        storage_type=self._storage_mode.storage_type,
                    )
                    if current_dict
                    else None
                )

    def __delitem__(self, key: K) -> None:
        """
        Delete an item

        Args:
            key: Key to delete
        """
        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary modification is only allowed within an transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
            )

        del self._working_dict[key]

    def __contains__(self, key: object) -> bool:
        """
        Check if a key exists

        Args:
            key: Key to check

        Returns:
            Whether the key exists
        """
        # first check if the key is in the cached properties
        if self._properties_to_cache and key in self._properties_to_cache:
            return self._cached_properties.__contains__(key)

        return self._get_dict().__contains__(key)

    def __len__(self) -> int:
        """
        Get the number of items

        Returns:
            Number of items in the dictionary
        """
        return self._length

    def __iter__(self) -> Iterator[K]:
        """
        Iterate over keys

        Returns:
            Iterator of keys
        """
        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary modification is only allowed within an transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
            )
        return iter(self._get_dict())

    def keys(self) -> KeysView[K]:
        """
        Get an iterator of keys

        Returns:
            Iterator of keys
        """
        return self._get_dict().keys()

    def values(self) -> ValuesView[V]:
        """
        Get an iterator of values

        Returns:
            Iterator of values
        """
        return self._get_dict().values()

    def items(self) -> ItemsView[K, V]:
        """
        Get an iterator of key-value pairs

        Returns:
            Iterator of key-value pairs
        """
        return self._get_dict().items()

    def dict(self, *, remove_nulls: bool = True) -> OrderedDict[K, V]:
        """
        Convert to a standard dictionary

        Returns:
            Standard dictionary with all values
        """
        if self._working_dict:
            return self._working_dict
        else:
            # if the working dict is None, return it but don't store it in the self._working_dict to keep memory low
            return self.create_working_dict()

    def __repr__(self) -> str:
        """
        String representation of the dictionary

        Returns:
            String representation
        """
        cached_property_list: List[str] = [
            f"{k}={v}" for k, v in self._cached_properties.items()
        ]
        return (
            (
                f"CompressedDict(storage_type='{self._storage_mode.storage_type}', keys={self._length}"
            )
            + (", " if cached_property_list else "")
            + ", ".join(cached_property_list)
            + ")"
        )

    def replace(
        self, *, value: Dict[K, V] | OrderedDict[K, V]
    ) -> "CompressedDict[K, V]":
        """
        Replace the current dictionary with a new one

        Args:
            value: New dictionary to store

        Returns:
            Self
        """
        if not value:
            self.clear()
            return self

        new_dict: OrderedDict[K, V] = (
            value if isinstance(value, OrderedDict) else OrderedDict[K, V](value)
        )
        self._update_serialized_dict(current_dict=new_dict)
        return self

    def clear(self) -> None:
        """
        Clear the dictionary
        """
        self._update_serialized_dict(current_dict=None)

    def __eq__(self, other: object) -> bool:
        """
        Check equality with another dictionary

        Args:
            other: Dictionary to compare with (CompressedDict or plain dict)

        Returns:
            Whether the dictionaries are equal in keys and values
        """
        # If other is not a dictionary-like object, return False
        if not isinstance(other, (CompressedDict, dict, OrderedDict)):
            return False

        # Get the dictionary representation of self
        self_dict = self.dict()

        # If other is a CompressedDict, use its _get_dict() method
        if isinstance(other, CompressedDict):
            other_dict = other.dict()
        else:
            # If other is a plain dict, use it directly
            if isinstance(other, OrderedDict):
                # If other is an OrderedDict, use it directly
                other_dict = other
            else:
                other_dict = OrderedDict[K, V](other)

        # Compare keys and values
        # Check that all keys in both dictionaries match exactly
        return set(self_dict.keys()) == set(other_dict.keys()) and all(
            self_dict[key] == other_dict[key] for key in self_dict
        )

    @overload
    def get(self, key: K) -> Optional[V]:
        """
        Get a value for an existing key

        :param key: Key to retrieve
        :return: Value or None if key is not found
        """
        ...

    # noinspection PyMethodOverriding
    @overload
    def get[_T](self, key: K, /, default: V | _T) -> V | _T:
        """
        Get a value with a default

        Args:
            key: Key to retrieve
            default: Default value if key is not found

        Returns:
            Value associated with the key or default
        """
        ...

    def get[_T](self, key: K, default: V | _T | None = None) -> V | _T | None:
        if key in self:
            return self[key]
        return default

    def __deepcopy__(self, memo: Dict[int, Any]) -> "CompressedDict[K, V]":
        """
        Create a deep copy of the dictionary

        Args:
            memo: Memoization dictionary for deep copy

        Returns:
            Deep copy of the dictionary
        """
        # Create a new instance with the same storage mode
        new_instance = CompressedDict(
            initial_dict=copy.deepcopy(self.dict()),
            storage_mode=self._storage_mode,
            properties_to_cache=self._properties_to_cache,
        )
        return new_instance

    def copy(self) -> "CompressedDict[K,V]":
        """
        Creates a copy of the BundleEntry object.

        :return: A new BundleEntry object with the same attributes.
        """
        return copy.deepcopy(self)

    def get_storage_mode(self) -> CompressedDictStorageMode:
        """
        Get the storage mode

        Returns:
            Storage mode
        """
        return self._storage_mode

    @overload
    def pop(self, key: K, /) -> V:
        """
        Remove and return a value for an existing key

        :param key: Key to remove
        :return: Removed value
        :raises KeyError: If key is not found
        """
        ...

    # noinspection PyMethodOverriding
    @overload
    def pop[_T](self, key: K, /, default: _T) -> V | _T:
        """
        Remove and return a value, or return default if key is not found

        :param key: Key to remove
        :param default: Default value to return if key is not found
        :return: Removed value or default
        """
        ...

    def pop[_T](self, key: K, /, default: V | _T | None = None) -> V | _T | None:
        """
        Remove and return a value

        :param key: Key to remove
        :param default: Optional default value if key is not found
        :return: Removed value or default
        :raises KeyError: If key is not found and no default is provided
        """
        if self._working_dict is None:
            raise CompressedDictAccessError(
                "Dictionary modification is only allowed within a transaction() block. "
                "Use 'with compressed_dict.transaction() as d:' to modify the dictionary."
            )

        # If no default is provided, use the standard dict.pop() behavior
        if default is None:
            return self._working_dict.pop(key)

        return self._working_dict.pop(key, default)

    def to_plain_dict(self) -> Dict[K, V]:
        """
        Get the plain dictionary representation

        Returns:
            Plain dictionary
        """
        return OrderedDictToDictConverter.convert(self.dict())

    @classmethod
    def from_json(cls, json_str: str) -> "CompressedDict[K, V]":
        """
        Create a FhirResource object from a JSON string.

        :param json_str: The JSON string to convert.
        :return: A FhirResource object.
        """
        data = TypePreservationSerializer.deserialize(json_str)
        return cls.from_dict(data)

    @classmethod
    def from_dict(
        cls,
        d: Dict[K, V],
        *,
        storage_mode: CompressedDictStorageMode = CompressedDictStorageMode.default(),
        properties_to_cache: List[K] | None = None,
    ) -> "CompressedDict[K, V]":
        """
        Creates a FhirResource object from a dictionary.

        :param d: The dictionary to convert.
        :param storage_mode: The storage mode for the CompressedDict.
        :param properties_to_cache: Optional list of properties to cache
        :return: A FhirResource object.
        """
        return cls(
            initial_dict=d,
            storage_mode=storage_mode,
            properties_to_cache=properties_to_cache,
        )
