from typing import Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, Union, cast

from starkware.cairo.lang.vm.memory_dict_backend import MemoryDictBackend
from starkware.cairo.lang.vm.relocatable import MaybeRelocatable, RelocatableValue

ADDR_SIZE_IN_BYTES = 8


class UnknownMemoryError(KeyError):
    def __init__(self, addr):
        self.addr = addr
        super().__init__(f"Unknown value for memory cell at address {addr}.")

    __str__: Callable[[BaseException], str] = Exception.__str__


class InconsistentMemoryError(Exception):
    def __init__(self, addr, old_value, new_value):
        self.addr = addr
        self.old_value = old_value
        self.new_value = new_value
        super().__init__(
            f"Inconsistent memory assignment at address {addr}. {old_value} != {new_value}."
        )


MemoryDictInitializer = Optional[
    Union[
        Mapping[MaybeRelocatable, MaybeRelocatable],
        Mapping[int, int],
        Iterable[Tuple[MaybeRelocatable, MaybeRelocatable]],
    ]
]

RelocateValueFunc = Callable[[MaybeRelocatable], MaybeRelocatable]


class MemoryDict:
    """
    Dictionary used for VM memory. Adds the following checks:
    * Checks that all memory addresses are valid.
    * getitem: Checks that the memory address is initialized.
    * setitem: Checks that memory value is not changed.
    """

    def __init__(
        self,
        values: MemoryDictInitializer = None,
        backend=MemoryDictBackend,
    ):
        if values is None:
            values = []
        elif isinstance(values, dict):
            values = values.items()
        self.data = backend(values)

        self._frozen: bool = False

        # A dict of segment relocation rules mapping a segment index to a RelocatableValue.
        # See add_relocation_rule for more details.
        self.relocation_rules: Dict[int, RelocatableValue] = {}

    def get(
        self, addr, default_value: Optional[MaybeRelocatable] = None
    ) -> Optional[MaybeRelocatable]:
        return self.relocate_value(self.data.get(addr, default_value))

    def items(self) -> Iterable[Tuple[MaybeRelocatable, MaybeRelocatable]]:
        return self.data.items()

    def keys(self):
        return self.data.keys()

    def __iter__(self):
        return iter(self.data)

    def __repr__(self):
        return repr(self.data)

    def __eq__(self, other):
        if not isinstance(other, MemoryDict):
            return NotImplemented
        return self.data == other.data

    def setdefault(self, key, default=None):
        # Copied from _collections_abc.MutableMapping.
        try:
            return self[key]
        except KeyError:
            self[key] = default
        return default

    def __contains__(self, addr):
        return addr in self.data

    def __len__(self):
        return len(self.data)

    def _check_element(self, num: MaybeRelocatable, name: str, exc_type: Type[Exception]):
        """
        Checks that num is a valid Cairo value: positive int or relocatable.
        Currently, does not check that value < prime.
        """
        if isinstance(num, RelocatableValue):
            return
        if not isinstance(num, int):
            raise exc_type(f"{name} must be an int, not {type(num).__name__}.")
        if num < 0:
            raise exc_type(f"{name} must be nonnegative. Got {num}.")

    def add_relocation_rule(self, src_ptr: RelocatableValue, dest_ptr: RelocatableValue):
        """
        Adds a relocation rule that moves values from the 'src_ptr' segment to 'dest_ptr'.

        'src_ptr' must point to the start of a temporary segment (negative index with offset 0).
        Once a relocation rule is set the memory dict relocates value on the fly,
        allowing the VM to execute the assertion 'src_ptr' = 'dest_ptr'.

        Note that relocation rules are not applied to addresses during execution
        and consequently adding a relocation rule does not allow the VM to
        read the value at memory['dest_ptr'].
        """
        assert src_ptr.segment_index < 0, f"src_ptr.segment_index must be < 0, src_ptr={src_ptr}."
        assert src_ptr.offset == 0, f"src_ptr.offset must be 0, src_ptr={src_ptr}."
        segment_index = src_ptr.segment_index
        assert (
            segment_index not in self.relocation_rules
        ), f"The segment with index {segment_index} already has a relocation rule."

        self.relocation_rules[segment_index] = dest_ptr

    def freeze(self):
        """
        Freezes the memory - no changes can be made from now on.
        """
        self._frozen = True

    def unfreeze_for_testing(self):
        """
        This function should only be used in tests.
        """
        self._frozen = False

    def is_frozen(self) -> bool:
        return self._frozen

    def relocate_value(self, value):
        """
        Relocates a value according to the relocation rules.

        The original value is returned if the relocation rules do not apply to value.
        """
        if not isinstance(value, RelocatableValue):
            return value

        segment_idx = value.segment_index
        if segment_idx >= 0:
            return value

        relocation = self.relocation_rules.get(segment_idx)
        if relocation is None:
            return value

        return self.relocate_value(relocation) + value.offset

    def relocate_memory(self):
        """
        Relocates the memory according to the relocation rules and clears self.relocation_rules.
        """
        assert not self._frozen, "Memory is frozen and cannot be changed."

        if len(self.relocation_rules) == 0:
            return

        relocated_memory: Dict[MaybeRelocatable, MaybeRelocatable] = {}
        for addr, value in self.items():
            relocated_addr = self.relocate_value(addr)
            relocated_value = self.relocate_value(value)
            current = relocated_memory.setdefault(relocated_addr, relocated_value)
            self.verify_same_value(addr=relocated_addr, current=current, value=relocated_value)

        self.data = relocated_memory
        self.relocation_rules = {}

    def __getitem__(self, addr: MaybeRelocatable) -> MaybeRelocatable:
        self._check_element(addr, "Memory address", KeyError)
        try:
            value = self.data[addr]
        except KeyError:
            raise UnknownMemoryError(addr) from None

        return self.relocate_value(value)

    def __setitem__(self, addr: MaybeRelocatable, value: MaybeRelocatable):
        assert not self._frozen, "Memory is frozen and cannot be changed."
        self._check_element(addr, "Memory address", KeyError)
        self._check_element(value, "Memory value", ValueError)

        # Additionally, check that address doesn't have a negative offset.
        if isinstance(addr, RelocatableValue) and addr.offset < 0:
            raise ValueError(
                f"The offset of a relocatable value must be nonnegative. Found: {addr}."
            )

        current = self.data.setdefault(addr, value)
        self.verify_same_value(addr, current, value)

    def verify_same_value(self, addr, current, value):
        """
        Verifies that 'current' and 'value' are the same and throws an exception otherwise.
        This function can be overridden by subclasses.
        """
        if current != value:
            raise InconsistentMemoryError(addr, current, value)

    def set_without_checks(self, addr: MaybeRelocatable, value: MaybeRelocatable):
        """
        Same as __setitem__() except that no checks are performed on addr and value.
        This function should only be used in tests.
        """
        self.data[addr] = value

    def serialize(self, field_bytes, relocate_value: Optional[RelocateValueFunc] = None):
        assert (
            len(self.relocation_rules) == 0
        ), "Cannot serialize a MemoryDict with active segment relocation rules."

        if relocate_value is None:
            relocate_value = lambda val: val

        return b"".join(
            RelocatableValue.to_bytes(relocate_value(addr), ADDR_SIZE_IN_BYTES, "little")
            + RelocatableValue.to_bytes(relocate_value(value), field_bytes, "little")
            for addr, value in self.items()
        )

    def get_range(self, addr, size) -> List[MaybeRelocatable]:
        return [self[addr + i] for i in range(size)]

    def get_range_as_ints(self, addr, size) -> List[int]:
        """
        Similar to get_range but asserts that all the values in the range are ints.
        """
        mem_range = self.get_range(addr=addr, size=size)
        assert all(isinstance(arg, int) for arg in mem_range)

        return cast(List[int], mem_range)

    @classmethod
    def deserialize(cls, data, field_bytes):
        pair_size = ADDR_SIZE_IN_BYTES + field_bytes
        assert (
            len(data) % (pair_size) == 0
        ), f"Data must consist of pairs of address (8 bytes) and value ({field_bytes} bytes)."
        pair_stream = (
            data[pair_size * i : pair_size * (i + 1)] for i in range(len(data) // pair_size)
        )
        return cls(
            (
                RelocatableValue.from_bytes(pair[:ADDR_SIZE_IN_BYTES], "little"),
                RelocatableValue.from_bytes(pair[ADDR_SIZE_IN_BYTES:], "little"),
            )
            for pair in pair_stream
        )
