"""
Optimized State.

.. contents:: Table of Contents
    :backlinks: none
    :local:

Introduction
------------

This module contains functions that can be monkey patched into the fork's
`state` module to use an optimized database backed state.
"""

import logging
from collections import defaultdict
from dataclasses import dataclass
from importlib import import_module
from typing import Any, ClassVar, Dict, List, Optional, Set, cast

try:
    import rust_pyspec_glue
except ImportError as e:
    # Add a message, but keep it an ImportError.
    raise e from Exception(
        "Install with `pip install 'ethereum[optimized]'` to enable this "
        "package"
    )

from ethereum_types.bytes import Bytes, Bytes20, Bytes32
from ethereum_types.numeric import U256, Uint

from ethereum.crypto.hash import Hash32

from .utils import add_item

Address = Bytes20
Root = Hash32
Account_ = Any  # noqa N806


class UnmodifiedType:
    """
    Sentinel type to represent a value that hasn't been modified.
    """

    pass


Unmodified = UnmodifiedType()


def get_optimized_state_patches(fork: str) -> Dict[str, Any]:
    """
    Get a dictionary of functions/objects to be monkey patched into the state
    to make it optimized.
    """
    patches: Dict[str, Any] = {}

    types_mod = cast(
        Any, import_module("ethereum.forks." + fork + ".fork_types")
    )
    state_mod = cast(Any, import_module("ethereum.forks." + fork + ".state"))
    Account = types_mod.Account  # noqa N806

    has_transient_storage = hasattr(state_mod, "TransientStorage")

    @add_item(patches)
    @dataclass
    class State:
        """
        The State, backed by a LMDB database.

        When created with `State()` store the db in a temporary directory. When
        created with `State(path)` open or create the db located at `path`.
        """

        default_path: ClassVar[Optional[str]] = None

        db: Any
        dirty_accounts: Dict[Address, Optional[Account_]]
        dirty_storage: Dict[Address, Dict[Bytes32, U256]]
        destroyed_accounts: Dict[Address, Uint]
        tx_restore_points: List[int]
        journal: List[Any]
        created_accounts: Set[Address]

        def __init__(self, path: Optional[str] = None) -> None:
            logging.info("using optimized state db at %s", path)

            if path is None:
                path = State.default_path

            self.db = rust_pyspec_glue.DB(path)
            self.dirty_accounts = {}
            self.dirty_storage = {}
            self.destroyed_accounts = defaultdict(lambda: Uint(0))
            self.tx_restore_points = []
            self.journal = []
            self.created_accounts = set()
            self.db.begin_mutable()

        def __eq__(self, other: object) -> bool:
            """
            Test for equality by comparing state roots.
            """
            if not isinstance(other, State):
                return NotImplemented
            return state_root(self) == state_root(other)

        def __enter__(self) -> "State":
            """Support with statements."""
            return self

        def __exit__(self, *args: Any) -> None:
            """Support with statements."""
            close_state(self)

    @add_item(patches)
    def close_state(state: State) -> None:
        """Close a state, releasing all resources it holds."""
        state.db.close()
        state.db = None
        del state.dirty_accounts
        del state.dirty_storage
        del state.destroyed_accounts
        del state.journal
        del state.created_accounts

    @add_item(patches)
    def get_metadata(state: State, key: Bytes) -> Optional[Bytes]:
        """Get a piece of metadata."""
        return state.db.get_metadata(key)

    @add_item(patches)
    def set_metadata(state: State, key: Bytes, value: Bytes) -> None:
        """Set a piece of metadata."""
        return state.db.set_metadata(key, value)

    @add_item(patches)
    def begin_db_transaction(state: State) -> None:
        """
        Start a database transaction. A transaction is automatically started
        when a `State` is created. Nesting of DB transactions is not supported
        (unlike non-db transactions).

        No operations are supported when not in a transaction.
        """
        state.db.begin_mutable()
        state.tx_restore_points = []
        state.journal = []

    @add_item(patches)
    def commit_db_transaction(state: State) -> None:
        """
        Commit the current database transaction.
        """
        if state.tx_restore_points:
            raise Exception("In a non-db transaction")
        flush(state)
        state.db.commit_mutable()

    @add_item(patches)
    def state_root(state: State) -> Root:
        """
        See `state`.
        """
        if state.tx_restore_points:
            raise Exception("In a non-db transaction")
        flush(state)
        return state.db.state_root()

    @add_item(patches)
    def storage_root(state: State, address: Address) -> Root:
        """
        See `state`.
        """
        if state.tx_restore_points:
            raise Exception("In a non-db transaction")
        flush(state)
        return state.db.storage_root(address)

    @add_item(patches)
    def flush(state: State) -> None:
        """
        Send everything in the internal caches to the Rust layer.
        """
        if state.tx_restore_points:
            raise Exception("In a non-db transaction")
        for address, count in state.destroyed_accounts.items():
            if count:
                state.db.destroy_storage(address)
        for address, account in state.dirty_accounts.items():
            state.db.set_account(address, account)
        for address, storage in state.dirty_storage.items():
            for key, value in storage.items():
                state.db.set_storage(address, key, value)
        state.destroyed_accounts.clear()
        state.dirty_accounts.clear()
        state.dirty_storage.clear()

    @add_item(patches)
    def rollback_db_transaction(state: State) -> None:
        """
        Rollback the current database transaction.
        """
        if state.tx_restore_points:
            raise Exception("In a non-db transaction")
        state.db.rollback_mutable()
        state.dirty_accounts.clear()
        state.dirty_storage.clear()
        state.destroyed_accounts.clear()

    def _begin_transaction(state: State) -> None:
        """
        See `state`.
        """
        if not state.tx_restore_points:
            flush(state)
        state.tx_restore_points.append(len(state.journal))

    if has_transient_storage:

        @add_item(patches)
        def begin_transaction(
            state: State,
            transient_storage: Optional[Any] = None,
        ) -> None:
            """
            See `state`.
            """
            _begin_transaction(state)
            if transient_storage is not None:
                transient_storage._snapshots.append(
                    {
                        k: state_mod.copy_trie(t)
                        for (k, t) in transient_storage._tries.items()
                    }
                )

    else:

        @add_item(patches)
        def begin_transaction(state: State) -> None:
            """
            See `state`.
            """
            _begin_transaction(state)

    def _commit_transaction(state: State) -> None:
        """
        See `state`.
        """
        state.tx_restore_points.pop()
        if not state.tx_restore_points:
            state.journal.clear()
            state.created_accounts.clear()
            flush(state)

    if has_transient_storage:

        @add_item(patches)
        def commit_transaction(
            state: State, transient_storage: Optional[Any] = None
        ) -> None:
            """
            See `state`.
            """
            _commit_transaction(state)

            if transient_storage and transient_storage._snapshots:
                transient_storage._snapshots.pop()

    else:

        @add_item(patches)
        def commit_transaction(state: State) -> None:
            """
            See `state`.
            """
            _commit_transaction(state)

    def _rollback_transaction(state: State) -> None:
        """
        See `state`.
        """
        restore_point = state.tx_restore_points.pop()
        while len(state.journal) > restore_point:
            item = state.journal.pop()
            if len(item) == 3:
                # Revert a storage key write
                if item[2] is Unmodified:
                    del state.dirty_storage[item[0]][item[1]]
                else:
                    state.dirty_storage[item[0]][item[1]] = item[2]
            elif type(item[1]) is dict:
                # Restore storage that was destroyed by `destroy_storage()`
                state.destroyed_accounts[item[0]] -= Uint(1)
                if state.destroyed_accounts[item[0]] == 0:
                    state.dirty_storage[item[0]] = item[1]
                    del state.destroyed_accounts[item[0]]
            else:
                # Revert a change to an account
                if item[1] is Unmodified:
                    del state.dirty_accounts[item[0]]
                else:
                    state.dirty_accounts[item[0]] = item[1]

        if not state.tx_restore_points:
            state.created_accounts.clear()

    if has_transient_storage:

        @add_item(patches)
        def rollback_transaction(
            state: State,
            transient_storage: Optional[Any] = None,
        ) -> None:
            """
            See `state`.
            """
            _rollback_transaction(state)

            if transient_storage and transient_storage._snapshots:
                transient_storage._tries = transient_storage._snapshots.pop()

    else:

        @add_item(patches)
        def rollback_transaction(state: State) -> None:
            """
            See `state`.
            """
            _rollback_transaction(state)

    @add_item(patches)
    def get_storage(state: State, address: Address, key: Bytes32) -> U256:
        """
        See `state`.
        """
        if (
            address in state.dirty_storage
            and key in state.dirty_storage[address]
        ):
            return state.dirty_storage[address][key]

        if state.destroyed_accounts[address]:
            return U256(0)
        else:
            return U256(state.db.get_storage(address, key))

    @add_item(patches)
    def get_storage_original(
        state: State, address: Address, key: Bytes32
    ) -> U256:
        """
        See `state`.
        """
        if address in state.created_accounts:
            return U256(0)
        else:
            return U256(state.db.get_storage(address, key))

    @add_item(patches)
    def set_storage(
        state: State, address: Address, key: Bytes32, value: U256
    ) -> None:
        """
        See `state`.
        """
        if address not in state.dirty_accounts:
            state.dirty_accounts[address] = get_account_optional(
                state, address
            )
        if address not in state.dirty_storage:
            state.dirty_storage[address] = {}
        if key not in state.dirty_storage[address]:
            state.journal.append((address, key, Unmodified))
        else:
            state.journal.append(
                (address, key, state.dirty_storage[address][key])
            )
        state.dirty_storage[address][key] = value

    @add_item(patches)
    def get_account_optional(
        state: State, address: Address
    ) -> Optional[Account_]:
        """
        See `state`.
        """
        if address in state.dirty_accounts:
            return state.dirty_accounts[address]
        account = state.db.get_account_optional(address)
        if account is not None:
            return Account(Uint(account[0]), U256(account[1]), account[2])
        else:
            return None

    @add_item(patches)
    def set_account(
        state: State, address: Address, account: Optional[Account_]
    ) -> None:
        """
        See `state`.
        """
        if address not in state.dirty_accounts:
            state.journal.append((address, Unmodified))
        if address in state.dirty_accounts:
            state.journal.append((address, state.dirty_accounts[address]))
        state.dirty_accounts[address] = account

    @add_item(patches)
    def destroy_storage(state: State, address: Address) -> None:
        """
        See `state`.
        """
        state.journal.append((address, state.dirty_storage.pop(address, {})))
        state.destroyed_accounts[address] += Uint(1)
        set_account(state, address, get_account_optional(state, address))

    @add_item(patches)
    def mark_account_created(state: State, address: Address) -> None:
        """
        See `state`.
        """
        state.created_accounts.add(address)

    @add_item(patches)
    def account_has_storage(state: State, address: Address) -> bool:
        """
        See `state`.
        """
        if address in state.dirty_storage:
            for v in state.dirty_storage[address].values():
                if v != U256(0):
                    return True

        if state.destroyed_accounts[address]:
            return False

        return state.db.has_storage(address)

    return patches
