# Copyright 2021 Vincent Texier <vit@free.fr>
#
# This software is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# Python Substrate Interface Library
#
# Copyright 2018-2023 Stichting Polkascan (Polkascan Foundation).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import binascii
import hashlib
import re
import secrets
from hashlib import blake2b
from math import ceil
from typing import Optional, Union

import ed25519_zebra
import nacl.bindings
import nacl.public
import sr25519
from bip39 import bip39_generate, bip39_to_mini_secret, bip39_validate
from eth_keys.datatypes import PrivateKey
from nacl.bindings import (
    crypto_sign_seed_keypair,
)
from scalecodec.base import ScaleBytes
from scalecodec.types import Bytes
from scalecodec.utils.ss58 import ss58_decode, ss58_encode

__all__ = ["DeriveJunction", "Keypair", "KeypairType", "MnemonicLanguageCode"]

RE_JUNCTION = r"(\/\/?)([^/]+)"
JUNCTION_ID_LEN = 32
DEV_PHRASE = "bottom drive obey lake curtain smoke basket hold race lonely fit walk"


class DeriveJunction:
    """
    DeriveJunction class
    """

    def __init__(self, chain_code, is_hard=False):
        """
        Init DeriveJunction instance

        :param chain_code:
        :param is_hard:
        """
        self.chain_code = chain_code
        self.is_hard = is_hard

    @classmethod
    def from_derive_path(cls, path: str, is_hard=False):
        """
        Return DeriveJunction instance from derive path

        :param path: Derive path
        :param is_hard: Hard derivation if True
        :return:
        """
        if path.isnumeric():
            byte_length = ceil(int(path).bit_length() / 8)
            chain_code = int(path).to_bytes(byte_length, "little").ljust(32, b"\x00")

        else:
            path_scale = Bytes()
            path_scale.encode(path)

            if len(path_scale.data) > JUNCTION_ID_LEN:
                chain_code = blake2b(path_scale.data.data, digest_size=32).digest()
            else:
                chain_code = bytes(path_scale.data.data.ljust(32, b"\x00"))

        return cls(chain_code=chain_code, is_hard=is_hard)


def extract_derive_path(derive_path: str):
    """
    Extract derivation path

    :param derive_path: Derive path
    :return:
    """
    path_check = ""
    junctions = []
    paths = re.findall(RE_JUNCTION, derive_path)

    if paths:
        path_check = "".join("".join(path) for path in paths)

        for path_separator, path_value in paths:
            junctions.append(
                DeriveJunction.from_derive_path(
                    path=path_value, is_hard=path_separator == "//"
                )
            )

    if path_check != derive_path:
        raise ValueError(
            'Reconstructed path "{}" does not match input'.format(path_check)
        )

    return junctions


def scale_encode_string(string: str) -> bytes:
    """
    Encode a string to SCALE

    :param string: String to encode
    :return:
    """
    string_in_bytes = string.encode("utf-8")
    length = len(string_in_bytes)

    # Compact encoding pour longueur < 64 (mode single-byte)
    if length < 64:
        encoded_len = (length << 2).to_bytes(1, "little")
    else:
        raise NotImplementedError(
            "SCALE compact encoding only implemented for len < 64"
        )

    return encoded_len + string_in_bytes


def derive_ed25519_hard_junction(seed: bytes, cc: bytes) -> bytes:
    """
    Derive a hard junction

    Same as rust function:

        use codec::Encode;
        ("Ed25519HDKD", secret_seed, cc).using_encoded(sp_crypto_hashing::blake2_256)

    :param seed: Mini-secret seed
    :param cc: Chain code of path junction
    :return:
    """
    assert len(seed) == 32
    assert len(cc) == 32

    encoded = scale_encode_string("Ed25519HDKD") + seed + cc
    return hashlib.blake2b(encoded, digest_size=32).digest()


class KeypairType:
    """
    Type of cryptography, used in `Keypair` instance to encrypt and sign data

    * ED25519 = 0
    * SR25519 = 1
    * ECDSA = 2

    """

    ED25519 = 0
    SR25519 = 1
    ECDSA = 2


class MnemonicLanguageCode:
    """
    Available language codes to generate mnemonics

    * ENGLISH = 'en'
    * CHINESE_SIMPLIFIED = 'zh-hans'
    * CHINESE_TRADITIONAL = 'zh-hant'
    * FRENCH = 'fr'
    * ITALIAN = 'it'
    * JAPANESE = 'ja'
    * KOREAN = 'ko'
    * SPANISH = 'es'

    """

    ENGLISH = "en"
    CHINESE_SIMPLIFIED = "zh-hans"
    CHINESE_TRADITIONAL = "zh-hant"
    FRENCH = "fr"
    ITALIAN = "it"
    JAPANESE = "ja"
    KOREAN = "ko"
    SPANISH = "es"


class Keypair:
    def __init__(
        self,
        ss58_address: Optional[str] = None,
        public_key: Optional[Union[bytes, str]] = None,
        private_key: Optional[Union[bytes, str]] = None,
        ss58_format: Optional[int] = None,
        seed_hex: Optional[Union[str, bytes]] = None,
        crypto_type: int = KeypairType.SR25519,
    ):
        """
        Allows generation of Keypairs from a variety of input combination, such as a public/private key combination,
        mnemonic or URI containing soft and hard derivation paths. With these Keypairs data can be signed and verified

        Parameters
        ----------
        ss58_address: Substrate address
        public_key: hex string or bytes of public_key key
        private_key: hex string or bytes of private key
        ss58_format: Substrate address format, default to 42 when omitted
        seed_hex: hex string of seed
        crypto_type: Use KeypairType.SR25519 or KeypairType.ED25519 cryptography for generating the Keypair
        """

        self.crypto_type = crypto_type
        self.seed_hex = seed_hex
        self.derive_path = None

        if crypto_type != KeypairType.ECDSA and ss58_address and not public_key:
            public_key = ss58_decode(ss58_address, valid_ss58_format=ss58_format)

        if private_key:

            if type(private_key) is str:
                private_key = bytes.fromhex(private_key.replace("0x", ""))  # type: ignore

            if self.crypto_type == KeypairType.SR25519:
                if len(private_key) != 64:
                    raise ValueError("Secret key should be 64 bytes long")
                if not public_key:
                    public_key = sr25519.public_from_secret_key(private_key)

            if self.crypto_type == KeypairType.ECDSA:
                private_key_obj = PrivateKey(private_key)  # type: ignore
                public_key = private_key_obj.public_key.to_address()
                ss58_address = private_key_obj.public_key.to_checksum_address()

        if not public_key:
            raise ValueError("No SS58 formatted address or public key provided")

        if type(public_key) is str:
            public_key = bytes.fromhex(public_key.replace("0x", ""))  # type: ignore

        if crypto_type == KeypairType.ECDSA:
            if len(public_key) != 20:
                raise ValueError("Public key should be 20 bytes long")
        else:
            if len(public_key) != 32:
                raise ValueError("Public key should be 32 bytes long")

            if not ss58_address:
                ss58_address = ss58_encode(public_key, ss58_format=ss58_format)

        self.ss58_format: int = ss58_format  # type: ignore

        self.public_key: bytes = public_key  # type: ignore

        self.ss58_address: str = ss58_address  # type: ignore

        self.private_key: bytes = private_key  # type: ignore

        self.mnemonic = None

    @classmethod
    def generate_mnemonic(
        cls, words: int = 12, language_code: str = MnemonicLanguageCode.ENGLISH
    ) -> str:
        """
        Generates a new seed phrase with given amount of words (default 12)

        Parameters
        ----------
        words: The amount of words to generate, valid values are 12, 15, 18, 21 and 24
        language_code: The language to use, valid values are: 'en', 'zh-hans', 'zh-hant', 'fr', 'it', 'ja', 'ko', 'es'.
                       Defaults to `MnemonicLanguageCode.ENGLISH`

        Returns
        -------
        str: Seed phrase
        """
        return bip39_generate(words, language_code)

    @classmethod
    def validate_mnemonic(
        cls, mnemonic: str, language_code: str = MnemonicLanguageCode.ENGLISH
    ) -> bool:
        """
        Verify if specified mnemonic is valid

        Parameters
        ----------
        mnemonic: Seed phrase
        language_code: The language to use, valid values are: 'en', 'zh-hans', 'zh-hant', 'fr', 'it', 'ja', 'ko', 'es'.
         Defaults to `MnemonicLanguageCode.ENGLISH`

        Returns
        -------
        bool
        """
        return bip39_validate(mnemonic, language_code)

    @classmethod
    def create_from_mnemonic(
        cls,
        mnemonic: str,
        ss58_format: int = 42,
        crypto_type: int = KeypairType.SR25519,
        language_code: str = MnemonicLanguageCode.ENGLISH,
    ) -> "Keypair":
        """
        Create a Keypair for given memonic

        Parameters
        ----------
        mnemonic: Seed phrase
        ss58_format: Substrate address format
        crypto_type: Use `KeypairType.SR25519` or `KeypairType.ED25519` cryptography for generating the Keypair
        language_code: The language to use, valid values are: 'en', 'zh-hans', 'zh-hant', 'fr', 'it', 'ja', 'ko', 'es'.
         Defaults to `MnemonicLanguageCode.ENGLISH`

        Returns
        -------
        Keypair
        """

        if crypto_type == KeypairType.ECDSA:
            raise ValueError("ECDSA not supported")
            # if language_code != MnemonicLanguageCode.ENGLISH:
            #     raise ValueError("ECDSA mnemonic only supports english")
            #
            # private_key = mnemonic_to_ecdsa_private_key(mnemonic)
            # keypair = cls.create_from_private_key(private_key, ss58_format=ss58_format, crypto_type=crypto_type)

        else:
            seed_array = bip39_to_mini_secret(mnemonic, "", language_code)

            keypair = cls.create_from_seed(
                seed_hex=binascii.hexlify(bytearray(seed_array)).decode("ascii"),
                ss58_format=ss58_format,
                crypto_type=crypto_type,
            )

        keypair.mnemonic = mnemonic  # type: ignore

        return keypair

    @classmethod
    def create_from_seed(
        cls,
        seed_hex: Union[bytes, str],
        ss58_format: int = 42,
        crypto_type: int = KeypairType.SR25519,
    ) -> "Keypair":
        """
        Create a Keypair for given seed

        Parameters
        ----------
        seed_hex: hex string of seed
        ss58_format: Substrate address format
        crypto_type: Use KeypairType.SR25519 or KeypairType.ED25519 cryptography for generating the Keypair

        Returns
        -------
        Keypair
        """

        if type(seed_hex) is str:
            seed_hex = bytes.fromhex(seed_hex.replace("0x", ""))  # type: ignore

        if crypto_type == KeypairType.SR25519:
            public_key, private_key = sr25519.pair_from_seed(seed_hex)
        elif crypto_type == KeypairType.ED25519:
            private_key, public_key = ed25519_zebra.ed_from_seed(seed_hex)
        else:
            raise ValueError('crypto_type "{}" not supported'.format(crypto_type))

        ss58_address = ss58_encode(public_key, ss58_format)

        return cls(
            ss58_address=ss58_address,
            public_key=public_key,
            private_key=private_key,
            ss58_format=ss58_format,
            crypto_type=crypto_type,
            seed_hex=seed_hex,
        )

    @classmethod
    def create_from_uri(
        cls,
        suri: str,
        ss58_format: int = 42,
        crypto_type: int = KeypairType.SR25519,
        language_code: str = MnemonicLanguageCode.ENGLISH,
    ) -> "Keypair":
        """
        Creates Keypair for specified suri in format: `[mnemonic]/[soft-path]//[hard-path]`

        Parameters
        ----------
        suri: str
            Secret URI including mnemonic and derivation path.
        ss58_format: int
            Substrate address format
        crypto_type: int
            Use KeypairType.SR25519 or KeypairType.ED25519
        language_code: MnemonicLanguageCode
            Language of the mnemonic phrase

        Returns
        -------
        Keypair
        """

        if suri and suri.startswith("/"):
            suri = DEV_PHRASE + suri

        suri_regex = re.match(
            r"^(?P<phrase>.[^/]+( .[^/]+)*)(?P<path>(//?[^/]+)*)(///(?P<password>.*))?$",
            suri,
        )

        suri_parts = suri_regex.groupdict()  # type: ignore

        if crypto_type == KeypairType.ECDSA:
            raise ValueError("ECDSA not supported")

        if suri_parts["password"]:
            raise NotImplementedError(
                f"Passwords in suri not supported for crypto_type '{crypto_type}'"
            )

        derived_keypair = cls.create_from_mnemonic(
            suri_parts["phrase"],
            ss58_format=ss58_format,
            crypto_type=crypto_type,
            language_code=language_code,
        )

        if suri_parts["path"] != "":
            derived_keypair.derive_path = suri_parts["path"]  # type: ignore
            derive_junctions = extract_derive_path(suri_parts["path"])

            child_pubkey = derived_keypair.public_key
            child_privkey = derived_keypair.private_key
            seed = derived_keypair.seed_hex
            assert isinstance(seed, bytes)

            for junction in derive_junctions:
                if crypto_type == KeypairType.SR25519:
                    if junction.is_hard:
                        _, child_pubkey, child_privkey = sr25519.hard_derive_keypair(
                            (junction.chain_code, child_pubkey, child_privkey), b""
                        )
                    else:
                        _, child_pubkey, child_privkey = sr25519.derive_keypair(
                            (junction.chain_code, child_pubkey, child_privkey), b""
                        )

                elif crypto_type == KeypairType.ED25519:
                    if not junction.is_hard:
                        raise NotImplementedError(
                            "Soft derivation not supported for ED25519"
                        )
                    # get derived seed from parent seed and chaincode
                    seed = derive_ed25519_hard_junction(seed, junction.chain_code)

                    # generate keypair
                    child_pubkey, child_privkey = crypto_sign_seed_keypair(seed)

                else:
                    raise NotImplementedError(
                        f"Derivation for crypto type {crypto_type} not supported"
                    )

            derived_keypair = Keypair(
                public_key=child_pubkey,
                private_key=child_privkey,
                ss58_format=ss58_format,
                crypto_type=crypto_type,
            )

        return derived_keypair

    @classmethod
    def create_from_private_key(
        cls,
        private_key: Union[bytes, str],
        public_key: Optional[Union[bytes, str]] = None,
        ss58_address: Optional[str] = None,
        ss58_format: Optional[int] = None,
        crypto_type: int = KeypairType.SR25519,
    ) -> "Keypair":
        """
        Creates Keypair for specified public/private keys
        Parameters
        ----------
        private_key: hex string or bytes of private key
        public_key: hex string or bytes of public key
        ss58_address: Substrate address
        ss58_format: Substrate address format, default = 42
        crypto_type: Use KeypairType.[SR25519|ED25519|ECDSA] cryptography for generating the Keypair

        Returns
        -------
        Keypair
        """

        return cls(
            ss58_address=ss58_address,
            public_key=public_key,
            private_key=private_key,
            ss58_format=ss58_format,
            crypto_type=crypto_type,
        )

    # @classmethod
    # def create_from_encrypted_json(cls, json_data: Union[str, dict], passphrase: str,
    #                                ss58_format: int = None) -> 'Keypair':
    #     """
    #     Create a Keypair from a PolkadotJS format encrypted JSON file
    #
    #     Parameters
    #     ----------
    #     json_data: Dict or JSON string containing PolkadotJS export format
    #     passphrase: Used to encrypt the keypair
    #     ss58_format: Which network ID to use to format the SS58 address (42 for testnet)
    #
    #     Returns
    #     -------
    #     Keypair
    #     """
    #
    #     if type(json_data) is str:
    #         json_data = json.loads(json_data)
    #
    #     private_key, public_key = decode_pair_from_encrypted_json(json_data, passphrase)
    #
    #     if 'sr25519' in json_data['encoding']['content']:
    #         crypto_type = KeypairType.SR25519
    #     elif 'ed25519' in json_data['encoding']['content']:
    #         crypto_type = KeypairType.ED25519
    #         # Strip the nonce part of the private key
    #         private_key = private_key[0:32]
    #     else:
    #         raise NotImplementedError("Unknown KeypairType found in JSON")
    #
    #     if ss58_format is None and 'address' in json_data:
    #         ss58_format = get_ss58_format(json_data['address'])
    #
    #     return cls.create_from_private_key(private_key, public_key, ss58_format=ss58_format, crypto_type=crypto_type)
    #
    # def export_to_encrypted_json(self, passphrase: str, name: str = None) -> dict:
    #     """
    #     Export Keypair to PolkadotJS format encrypted JSON file
    #
    #     Parameters
    #     ----------
    #     passphrase: Used to encrypt the keypair
    #     name: Display name of Keypair used
    #
    #     Returns
    #     -------
    #     dict
    #     """
    #     if not name:
    #         name = self.ss58_address
    #
    #     if self.crypto_type != KeypairType.SR25519:
    #         raise NotImplementedError(f"Cannot create JSON for crypto_type '{self.crypto_type}'")
    #
    #     # Secret key from PolkadotJS is an Ed25519 expanded secret key, so has to be converted
    #     # https://github.com/polkadot-js/wasm/blob/master/packages/wasm-crypto/src/rs/sr25519.rs#L125
    #     converted_private_key = sr25519.convert_secret_key_to_ed25519(self.private_key)
    #
    #     encoded = encode_pair(self.public_key, converted_private_key, passphrase)
    #
    #     json_data = {
    #         "encoded": b64encode(encoded).decode(),
    #         "encoding": {"content": ["pkcs8", "sr25519"], "type": ["scrypt", "xsalsa20-poly1305"], "version": "3"},
    #         "address": self.ss58_address,
    #         "meta": {
    #             "name": name, "tags": [], "whenCreated": int(time.time())
    #         }
    #     }
    #
    #     return json_data

    def sign(self, data: Union[ScaleBytes, bytes, str]) -> bytes:
        """
        Creates a signature for given data

        Parameters
        ----------
        data: data to sign in `Scalebytes`, bytes or hex string format

        Returns
        -------
        signature in bytes

        """
        if type(data) is ScaleBytes:
            data = bytes(data.data)  # type: ignore
        elif data[0:2] == "0x":
            data = bytes.fromhex(data[2:])  # type: ignore
        elif type(data) is str:
            data = data.encode()  # type: ignore

        if not self.private_key:
            raise ValueError("No private key set to create signatures")

        if self.crypto_type == KeypairType.SR25519:
            signature = sr25519.sign((self.public_key, self.private_key), data)

        elif self.crypto_type == KeypairType.ED25519:
            signature = ed25519_zebra.ed_sign(self.private_key, data)

        elif self.crypto_type == KeypairType.ECDSA:
            raise ValueError("Crypto type not supported")
            # signature = ecdsa_sign(self.private_key, data)

        else:
            raise ValueError("Crypto type not supported")

        return signature

    def verify(
        self, data: Union[ScaleBytes, bytes, str], signature: Union[bytes, str]
    ) -> bool:
        """
        Verifies data with specified signature

        Parameters
        ----------
        data: data to be verified in `Scalebytes`, bytes or hex string format
        signature: signature in bytes or hex string format

        Returns
        -------
        True if data is signed with this Keypair, otherwise False
        """

        if type(data) is ScaleBytes:
            data = bytes(data.data)  # type: ignore
        elif data[0:2] == "0x":
            data = bytes.fromhex(data[2:])  # type: ignore
        elif type(data) is str:
            data = data.encode()  # type: ignore

        if type(signature) is str and signature[0:2] == "0x":
            signature = bytes.fromhex(signature[2:])  # type: ignore

        if type(signature) is not bytes:
            raise TypeError("Signature should be of type bytes or a hex-string")

        if self.crypto_type == KeypairType.SR25519:
            crypto_verify_fn = sr25519.verify
        elif self.crypto_type == KeypairType.ED25519:
            crypto_verify_fn = ed25519_zebra.ed_verify
        elif self.crypto_type == KeypairType.ECDSA:
            raise ValueError("Crypto type not supported")
            # crypto_verify_fn = ecdsa_verify
        else:
            raise ValueError("Crypto type not supported")

        verified = crypto_verify_fn(signature, data, self.public_key)

        if not verified:
            # Another attempt with the data wrapped, as discussed in https://github.com/polkadot-js/extension/pull/743
            # Note: As Python apps are trusted sources on its own, no need to wrap data when signing from this lib
            verified = crypto_verify_fn(
                signature, b"<Bytes>" + data + b"</Bytes>", self.public_key  # type: ignore
            )

        return verified

    def encrypt_message(
        self,
        message: Union[bytes, str],
        recipient_public_key: bytes,
        nonce: bytes = secrets.token_bytes(24),
    ) -> bytes:
        """
        Encrypts message with for specified recipient

        Parameters
        ----------
        message: message to be encrypted, bytes or string
        recipient_public_key: recipient's public key
        nonce: the nonce to use in the encryption

        Returns
        -------
        Encrypted message
        """

        if not self.private_key:
            raise ValueError("No private key set to encrypt")
        if self.crypto_type != KeypairType.ED25519:
            raise ValueError("Only ed25519 keypair type supported")
        curve25519_public_key = nacl.bindings.crypto_sign_ed25519_pk_to_curve25519(
            recipient_public_key
        )
        recipient = nacl.public.PublicKey(curve25519_public_key)
        private_key = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(
            self.private_key + self.public_key
        )
        sender = nacl.public.PrivateKey(private_key)
        box = nacl.public.Box(sender, recipient)
        return box.encrypt(
            message if isinstance(message, bytes) else message.encode("utf-8"), nonce
        )

    def decrypt_message(
        self, encrypted_message_with_nonce: bytes, sender_public_key: bytes
    ) -> bytes:
        """
        Decrypts message from a specified sender

        Parameters
        ----------
        encrypted_message_with_nonce: message to be decrypted
        sender_public_key: sender's public key

        Returns
        -------
        Decrypted message
        """

        if not self.private_key:
            raise ValueError("No private key set to decrypt")
        if self.crypto_type != KeypairType.ED25519:
            raise ValueError("Only ed25519 keypair type supported")
        private_key = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(
            self.private_key + self.public_key
        )
        recipient = nacl.public.PrivateKey(private_key)
        curve25519_public_key = nacl.bindings.crypto_sign_ed25519_pk_to_curve25519(
            sender_public_key
        )
        sender = nacl.public.PublicKey(curve25519_public_key)
        return nacl.public.Box(recipient, sender).decrypt(encrypted_message_with_nonce)

    def __repr__(self):
        if self.ss58_address:
            return "<Keypair (address={})>".format(self.ss58_address)
        else:
            return "<Keypair (public_key=0x{})>".format(self.public_key.hex())
