###############################################################################
#
# The MIT License (MIT)
#
# Copyright (c) typedef int GmbH
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
###############################################################################
import binascii
import configparser
import os
from collections.abc import MutableMapping
from threading import Lock
from typing import Any, Dict, Iterator, List, Optional, Union

import nacl
import txaio
from autobahn.util import parse_keyfile
from autobahn.wamp.cryptosign import CryptosignKey
from autobahn.wamp.interfaces import IEthereumKey, ISecurityModule
from eth_account.account import Account
from eth_account.signers.local import LocalAccount
from py_eth_sig_utils.eip712 import encode_typed_data
from py_eth_sig_utils.signing import signature_to_v_r_s, v_r_s_to_signature
from py_eth_sig_utils.utils import checksum_encode, ecrecover_to_pub, ecsign, sha3

from xbr._mnemonic import mnemonic_to_private_key

__all__ = (
    "EthereumKey",
    "SecurityModuleMemory",
)


class EthereumKey(object):
    """
    Base class to implement :class:`autobahn.wamp.interfaces.IEthereumKey`.
    """

    def __init__(
        self,
        key_or_address: Union[LocalAccount, str, bytes],
        can_sign: bool,
        security_module: Optional[ISecurityModule] = None,
        key_no: Optional[int] = None,
    ) -> None:
        if can_sign:
            # https://eth-account.readthedocs.io/en/latest/eth_account.html#eth_account.account.Account
            assert type(key_or_address) == LocalAccount
            self._key = key_or_address
            self._address = key_or_address.address
        else:
            assert type(key_or_address) in (str, bytes)
            self._key = None
            self._address = key_or_address
        self._can_sign = can_sign
        self._security_module = security_module
        self._key_no = key_no

    @property
    def security_module(self) -> Optional["ISecurityModule"]:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.security_module`.
        """
        return self._security_module

    @property
    def key_no(self) -> Optional[int]:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.key_no`.
        """
        return self._key_no

    @property
    def key_type(self) -> str:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.key_type`.
        """
        return "ethereum"

    def public_key(self, binary: bool = False) -> Union[str, bytes]:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.public_key`.
        """
        raise NotImplementedError()

    @property
    def can_sign(self) -> bool:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.can_sign`.
        """
        return self._can_sign

    def address(self, binary: bool = False) -> Union[str, bytes]:
        """
        Implements :meth:`autobahn.wamp.interfaces.IEthereumKey.address`.
        """
        if binary:
            return binascii.a2b_hex(self._address[2:])
        else:
            return self._address

    def sign(self, data: bytes) -> bytes:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.sign`.
        """
        # FIXME: implement signing of raw data
        raise NotImplementedError()

    def recover(self, data: bytes, signature: bytes) -> bytes:
        """
        Implements :meth:`autobahn.wamp.interfaces.IKey.recover`.
        """
        # FIXME: implement signing address recovery from signature of raw data
        raise NotImplementedError()

    def sign_typed_data(self, data: Dict[str, Any], binary=True) -> bytes:
        """
        Implements :meth:`autobahn.wamp.interfaces.IEthereumKey.sign_typed_data`.
        """
        if self._security_module:
            assert self._security_module.is_open and not self._security_module.is_locked, (
                "security module must be open and unlocked"
            )
        try:
            # encode typed data dict and return message hash
            msg_hash = encode_typed_data(data)

            # ECDSA signatures in Ethereum consist of three parameters: v, r and s.
            # The signature is always 65-bytes in length.
            #     r = first 32 bytes of signature
            #     s = second 32 bytes of signature
            #     v = final 1 byte of signature
            signature_vrs = ecsign(msg_hash, self._key.key)

            # concatenate signature components into byte string
            signature = v_r_s_to_signature(*signature_vrs)
        except Exception as e:
            return txaio.create_future_error(e)
        else:
            if binary:
                return txaio.create_future_success(signature)
            else:
                return txaio.create_future_success(binascii.b2a_hex(signature).decode())

    def verify_typed_data(self, data: Dict[str, Any], signature: bytes) -> bool:
        """
        Implements :meth:`autobahn.wamp.interfaces.IEthereumKey.verify_typed_data`.
        """
        if self._security_module:
            assert self._security_module.is_open and not self._security_module.is_locked, (
                "security module must be open and unlocked"
            )
        try:
            msg_hash = encode_typed_data(data)
            signature_vrs = signature_to_v_r_s(signature)
            public_key = ecrecover_to_pub(msg_hash, *signature_vrs)
            address_bytes = sha3(public_key)[-20:]
            address = checksum_encode(address_bytes)
        except Exception as e:
            return txaio.create_future_error(e)
        else:
            return txaio.create_future_success(address == self._address)

    @classmethod
    def from_address(cls, address: Union[str, bytes]) -> "EthereumKey":
        """
        Create a public key from an address, which can be used to verify signatures.

        :param address: The Ethereum address (20 octets).
        :return: New instance of :class:`EthereumKey`
        """
        return EthereumKey(key_or_address=address, can_sign=False)

    @classmethod
    def from_bytes(cls, key: bytes) -> "EthereumKey":
        """
        Create a private key from seed bytes, which can be used to sign and create signatures.

        :param key: The Ethereum private key seed (32 octets).
        :return: New instance of :class:`EthereumKey`
        """
        if type(key) != bytes:
            raise ValueError("invalid seed type {} (expected binary)".format(type(key)))

        if len(key) != 32:
            raise ValueError("invalid seed length {} (expected 32)".format(len(key)))

        account: LocalAccount = Account.from_key(key)
        return EthereumKey(key_or_address=account, can_sign=True)

    @classmethod
    def from_seedphrase(cls, seedphrase: str, index: int = 0) -> "EthereumKey":
        """
        Create a private key from the given BIP-39 mnemonic seed phrase and index,
        which can be used to sign and create signatures.

        :param seedphrase: The BIP-39 seedphrase ("Mnemonic") from which to derive the account.
        :param index: The account index in account hierarchy defined by the seedphrase.
        :return: New instance of :class:`EthereumKey`
        """
        # Base HD Path:  m/44'/60'/0'/0/{account_index}
        derivation_path = "m/44'/60'/0'/0/{}".format(index)

        key = mnemonic_to_private_key(seedphrase, str_derivation_path=derivation_path)
        assert type(key) == bytes
        assert len(key) == 32

        account: LocalAccount = Account.from_key(key)
        return EthereumKey(key_or_address=account, can_sign=True)

    @classmethod
    def from_keyfile(cls, keyfile: str) -> "EthereumKey":
        """
        Create a public or private key from reading the given public or private key file.

        Here is an example key file that includes an Ethereum private key ``private-key-eth``, which
        is loaded in this function, and other fields, which are ignored by this function:

        .. code-block::

            This is a comment (all lines until the first empty line are comments indeed).

            creator: oberstet@intel-nuci7
            created-at: 2022-07-05T12:29:48.832Z
            user-id: oberstet@intel-nuci7
            public-key-ed25519: 7326d9dc0307681cc6940fde0e60eb31a6e4d642a81e55c434462ce31f95deed
            public-adr-eth: 0x10848feBdf7f200Ba989CDf7E3eEB3EC03ae7768
            private-key-ed25519: f750f42b0430e28a2e272c3cedcae4dcc4a1cf33bc345c35099d3322626ab666
            private-key-eth: 4d787714dcb0ae52e1c5d2144648c255d660b9a55eac9deeb80d9f506f501025

        :param keyfile: Path (relative or absolute) to a public or private keys file.
        :return: New instance of :class:`EthereumKey`
        """
        if not os.path.exists(keyfile) or not os.path.isfile(keyfile):
            raise RuntimeError('keyfile "{}" is not a file'.format(keyfile))

        # now load the private or public key file - this returns a dict which should
        # include (for a private key):
        #
        #   private-key-eth: 6b08b6e186bd2a3b9b2f36e6ece3f8031fe788ab3dc4a1cfd3a489ea387c496b
        #
        # or (for a public key only):
        #
        #   public-adr-eth: 0x10848feBdf7f200Ba989CDf7E3eEB3EC03ae7768
        #
        data = parse_keyfile(keyfile)

        privkey_eth_hex = data.get("private-key-eth", None)
        if privkey_eth_hex is None:
            pub_adr_eth = data.get("public-adr-eth", None)
            if pub_adr_eth is None:
                raise RuntimeError(
                    'neither "private-key-eth" nor "public-adr-eth" found in keyfile {}'.format(keyfile)
                )
            else:
                return EthereumKey.from_address(pub_adr_eth)
        else:
            return EthereumKey.from_bytes(binascii.a2b_hex(privkey_eth_hex))


IEthereumKey.register(EthereumKey)


class SecurityModuleMemory(MutableMapping):
    """
    A transient, memory-based implementation of :class:`ISecurityModule`.
    """

    def __init__(self, keys: Optional[List[Union[CryptosignKey, EthereumKey]]] = None):
        self._mutex = Lock()
        self._is_open = False
        self._is_locked = True
        self._keys: Dict[int, Union[CryptosignKey, EthereumKey]] = {}
        self._counters: Dict[int, int] = {}
        if keys:
            for i, key in enumerate(keys):
                self._keys[i] = key

    def __len__(self) -> int:
        """
        Implements :meth:`ISecurityModule.__len__`
        """
        assert self._is_open, "security module not open"

        return len(self._keys)

    def __contains__(self, key_no: int) -> bool:
        assert self._is_open, "security module not open"

        return key_no in self._keys

    def __iter__(self) -> Iterator[int]:
        """
        Implements :meth:`ISecurityModule.__iter__`
        """
        assert self._is_open, "security module not open"

        yield from self._keys

    def __getitem__(self, key_no: int) -> Union[CryptosignKey, EthereumKey]:
        """
        Implements :meth:`ISecurityModule.__getitem__`
        """
        assert self._is_open, "security module not open"

        if key_no in self._keys:
            return self._keys[key_no]
        else:
            raise IndexError("key_no {} not found".format(key_no))

    def __setitem__(self, key_no: int, key: Union[CryptosignKey, EthereumKey]) -> None:
        assert self._is_open, "security module not open"

        assert key_no >= 0
        if key_no in self._keys:
            # FIXME
            pass
        self._keys[key_no] = key

    def __delitem__(self, key_no: int) -> None:
        assert self._is_open, "security module not open"

        if key_no in self._keys:
            del self._keys[key_no]
        else:
            raise IndexError()

    def open(self):
        """
        Implements :meth:`ISecurityModule.open`
        """
        assert not self._is_open, "security module already open"

        self._is_open = True
        return txaio.create_future_success(None)

    def close(self):
        """
        Implements :meth:`ISecurityModule.close`
        """
        assert self._is_open, "security module not open"

        self._is_open = False
        self._is_locked = True
        return txaio.create_future_success(None)

    @property
    def is_open(self) -> bool:
        """
        Implements :meth:`ISecurityModule.is_open`
        """
        return self._is_open

    @property
    def can_lock(self) -> bool:
        """
        Implements :meth:`ISecurityModule.can_lock`
        """
        return True

    @property
    def is_locked(self) -> bool:
        """
        Implements :meth:`ISecurityModule.is_locked`
        """
        return self._is_locked

    def lock(self):
        """
        Implements :meth:`ISecurityModule.lock`
        """
        assert self._is_open, "security module not open"
        assert not self._is_locked

        self._is_locked = True
        return txaio.create_future_success(None)

    def unlock(self):
        """
        Implements :meth:`ISecurityModule.unlock`
        """
        assert self._is_open, "security module not open"
        assert self._is_locked

        self._is_locked = False
        return txaio.create_future_success(None)

    def create_key(self, key_type: str) -> int:
        assert self._is_open, "security module not open"

        key_no = len(self._keys)
        if key_type == "cryptosign":
            key = CryptosignKey(
                key=nacl.signing.SigningKey(os.urandom(32)), can_sign=True, security_module=self, key_no=key_no
            )
        elif key_type == "ethereum":
            key = EthereumKey(
                key_or_address=Account.from_key(os.urandom(32)), can_sign=True, security_module=self, key_no=key_no
            )
        else:
            raise ValueError('invalid key_type "{}"'.format(key_type))
        self._keys[key_no] = key
        return txaio.create_future_success(key_no)

    def delete_key(self, key_no: int):
        assert self._is_open, "security module not open"

        if key_no in self._keys:
            del self._keys[key_no]
            return txaio.create_future_success(key_no)
        else:
            return txaio.create_future_success(None)

    def get_random(self, octets: int) -> bytes:
        """
        Implements :meth:`ISecurityModule.get_random`
        """
        assert self._is_open, "security module not open"

        data = os.urandom(octets)
        return txaio.create_future_success(data)

    def get_counter(self, counter_no: int) -> int:
        """
        Implements :meth:`ISecurityModule.get_counter`
        """
        assert self._is_open, "security module not open"

        self._mutex.acquire()
        res = self._counters.get(counter_no, 0)
        self._mutex.release()
        return txaio.create_future_success(res)

    def increment_counter(self, counter_no: int) -> int:
        """
        Implements :meth:`ISecurityModule.increment_counter`
        """
        assert self._is_open, "security module not open"

        self._mutex.acquire()
        if counter_no not in self._counters:
            self._counters[counter_no] = 0
        self._counters[counter_no] += 1
        res = self._counters[counter_no]
        self._mutex.release()
        return txaio.create_future_success(res)

    @classmethod
    def from_seedphrase(cls, seedphrase: str, num_eth_keys: int = 1, num_cs_keys: int = 1) -> "SecurityModuleMemory":
        """
        Create a new memory-backed security module with

        1. ``num_eth_keys`` keys of type :class:`EthereumKey`, followed by
        2. ``num_cs_keys`` keys of type :class:`CryptosignKey`

        computed from a (common) BIP44 seedphrase.

        :param seedphrase: BIP44 seedphrase to use.
        :param num_eth_keys: Number of Ethereum keys to derive.
        :param num_cs_keys: Number of Cryptosign keys to derive.
        :return: New memory-backed security module instance.
        """
        keys: List[Union[EthereumKey, CryptosignKey]] = []

        # first, add num_eth_keys EthereumKey(s), numbering starting at 0
        for i in range(num_eth_keys):
            key = EthereumKey.from_seedphrase(seedphrase, i)
            keys.append(key)

        # second, add num_cs_keys CryptosignKey(s), numbering starting at num_eth_keys (!)
        for i in range(num_cs_keys):
            key = CryptosignKey.from_seedphrase(seedphrase, i + num_eth_keys)
            keys.append(key)

        # initialize security module from collected keys
        sm = SecurityModuleMemory(keys=keys)
        return sm

    @classmethod
    def from_config(cls, config: str, profile: str = "default") -> "SecurityModuleMemory":
        """
        Create a new memory-backed security module with keys referred from a profile in
        the given configuration file.

        :param config: Path (relative or absolute) to an INI configuration file.
        :param profile: Name of the profile within the given INI configuration file.
        :return: New memory-backed security module instance.
        """
        keys: List[Union[EthereumKey, CryptosignKey]] = []

        cfg = configparser.ConfigParser()
        cfg.read(config)

        if not cfg.has_section(profile):
            raise RuntimeError('profile "{}" not found in configuration file "{}"'.format(profile, config))

        if not cfg.has_option(profile, "privkey"):
            raise RuntimeError(
                'missing option "privkey" in profile "{}" of configuration file "{}"'.format(profile, config)
            )

        privkey = os.path.join(os.path.dirname(config), cfg.get(profile, "privkey"))
        if not os.path.exists(privkey) or not os.path.isfile(privkey):
            raise RuntimeError(
                'privkey "{}" is not a file in profile "{}" of configuration file "{}"'.format(
                    privkey, profile, config
                )
            )

        # now load the private key file - this returns a dict which should include:
        # private-key-eth: 6b08b6e186bd2a3b9b2f36e6ece3f8031fe788ab3dc4a1cfd3a489ea387c496b
        # private-key-ed25519: 20e8c05d0ede9506462bb049c4843032b18e8e75b314583d0c8d8a4942f9be40
        data = parse_keyfile(privkey)

        # first, add Ethereum key
        privkey_eth_hex = data.get("private-key-eth", None)
        keys.append(EthereumKey.from_bytes(binascii.a2b_hex(privkey_eth_hex)))

        # second, add Cryptosign key
        privkey_ed25519_hex = data.get("private-key-ed25519", None)
        keys.append(CryptosignKey.from_bytes(binascii.a2b_hex(privkey_ed25519_hex)))

        # initialize security module from collected keys
        sm = SecurityModuleMemory(keys=keys)
        return sm

    @classmethod
    def from_keyfile(cls, keyfile: str) -> "SecurityModuleMemory":
        """
        Create a new memory-backed security module with keys referred from a profile in
        the given configuration file.

        :param keyfile: Path (relative or absolute) to a private keys file.
        :return: New memory-backed security module instance.
        """
        keys: List[Union[EthereumKey, CryptosignKey]] = []

        if not os.path.exists(keyfile) or not os.path.isfile(keyfile):
            raise RuntimeError('keyfile "{}" is not a file'.format(keyfile))

        # now load the private key file - this returns a dict which should include:
        # private-key-eth: 6b08b6e186bd2a3b9b2f36e6ece3f8031fe788ab3dc4a1cfd3a489ea387c496b
        # private-key-ed25519: 20e8c05d0ede9506462bb049c4843032b18e8e75b314583d0c8d8a4942f9be40
        data = parse_keyfile(keyfile)

        # first, add Ethereum key
        privkey_eth_hex = data.get("private-key-eth", None)
        if privkey_eth_hex is None:
            raise RuntimeError('"private-key-eth" not found in keyfile {}'.format(keyfile))
        keys.append(EthereumKey.from_bytes(binascii.a2b_hex(privkey_eth_hex)))

        # second, add Cryptosign key
        privkey_ed25519_hex = data.get("private-key-ed25519", None)
        if privkey_ed25519_hex is None:
            raise RuntimeError('"private-key-ed25519" not found in keyfile {}'.format(keyfile))
        keys.append(CryptosignKey.from_bytes(binascii.a2b_hex(privkey_ed25519_hex)))

        # initialize security module from collected keys
        sm = SecurityModuleMemory(keys=keys)
        return sm


ISecurityModule.register(SecurityModuleMemory)
