from typing import Callable, Tuple, Type

__all__ = [
    "Hash",
    "BlockCipher",
    "ECDLP",
    "EllipticCurve",
    "EllipticCurveCipher",
]


def _inv(x: int, p: int):
    r1 = p
    r2 = x
    t1 = 0
    t2 = 1
    while r2 > 0:
        q, r = divmod(r1, r2)
        r1 = r2
        r2 = r
        t = t1 - q * t2
        t1 = t2
        t2 = t
    return t1 % p


class Hash:
    @classmethod
    def max_msg_length(self) -> int:
        """Get maximum message length in bytes."""

        raise NotImplementedError

    @classmethod
    def hash_length(self) -> int:
        """Get output hash value length in bytes."""

        raise NotImplementedError

    def __init__(self) -> None:
        raise NotImplementedError

    def update(self, data: bytes) -> None:
        """Update internal state.

        Args:
            data (bytes): data stream to be updated.
        """

        raise NotImplementedError

    def value(self) -> bytes:
        """Returns current hash value in bytes.

        Returns:
            bytes: Hash value.
        """

        raise NotImplementedError


class BlockCipher:
    @classmethod
    def key_length(self) -> int:
        """Get key length in bytes."""

        raise NotImplementedError

    @classmethod
    def block_length(self) -> int:
        """Get block length in bytes."""

        raise NotImplementedError

    def __init__(self, key: bytes) -> None:
        """Block Cipher

        Args:
            key (bytes): key used in cipher, has a length of `BlockCipher.key_length()`
        """

        raise NotImplementedError

    def encrypt(self, block: bytes) -> bytes:
        """Encrypt"""

        raise NotImplementedError

    def decrypt(self, block: bytes) -> bytes:
        """Decrypt"""

        raise NotImplementedError


class EllipticCurve:
    """Elliptic Curve (Fp)"""

    INF = (-1, -1)

    @staticmethod
    def isinf(x: int, y: int) -> bool:
        return x < 0 or y < 0

    def __init__(self, p: int, a: int, b: int) -> None:
        self.p = p
        self.a = a
        self.b = b

        self.bitlength = self.p.bit_length()
        self.length = (self.bitlength + 7) >> 3

    def isvalid(self, x: int, y: int) -> bool:
        """Verify if a point is on the curve."""

        if x >= self.p or y >= self.p:
            return False

        if (y * y - x * x * x - self.a * x - self.b) % self.p != 0:
            return False

        return True

    def add(self, x1: int, y1: int, x2: int, y2: int) -> Tuple[int, int]:
        """Add two points. Negative numbers means infinite point."""

        a = self.a
        p = self.p

        if self.isinf(x1, y1):
            return x2, y2
        if self.isinf(x2, y2):
            return x1, y1
        if x1 == x2:
            if y1 + y2 == p:
                return -1, -1
            elif y1 == y2:
                lam = (3 * x1 * x1 + a) * _inv(2 * y1, p)
            else:
                raise ValueError(f"{hex(y1)} and {hex(y2)} is neither equal nor opposite.")
        else:
            if x2 > x1:
                lam = (y2 - y1) * _inv(x2 - x1, p)
            else:
                lam = (y1 - y2) * _inv(x1 - x2, p)

        x3 = (lam * lam - x1 - x2) % p
        y3 = (lam * (x1 - x3) - y1) % p

        return x3, y3

    def sub(self, x1: int, y1: int, x2: int, y2: int) -> Tuple[int, int]:
        """Sub two points."""

        return self.add(x1, y1, x2, self.p - y2)

    def mul(self, k: int, x: int, y: int) -> Tuple[int, int]:
        """k-multiply points."""

        xk = -1
        yk = -1

        for i in f"{k:b}":
            xk, yk = self.add(xk, yk, xk, yk)
            if i == "1":
                xk, yk = self.add(xk, yk, x, y)

        return xk, yk

    def itob(self, i: int) -> bytes:
        """Convert domain elements to bytes."""

        return i.to_bytes(self.length, "big")

    def btoi(self, b: bytes) -> int:
        """Convert bytes to domain elements"""

        return int.from_bytes(b, "big")


class ECDLP(EllipticCurve):
    """Elliptic Curve Discrete Logarithm Problem"""

    def __init__(self, p: int, a: int, b: int, xG: int, yG: int, n: int, h: int = 1) -> None:
        super().__init__(p, a, b)

        self.xG = xG
        self.yG = yG
        self.n = n
        self.h = h

    def kG(self, k: int) -> Tuple[int, int]:
        return self.mul(k, self.xG, self.yG)


class EllipticCurveCipher:
    """Elliptic Curve Cipher"""

    def __init__(self, ecdlp: ECDLP, hash_cls: Type[Hash], rnd_fn: Callable[[int], int]) -> None:
        """ECC

        Args:
            ecdlp (ECDLP): ECDLP used in cipher.
            hash_fn (Hash): hash function used in cipher.
            rnd_fn ((int) -> int): random function used to generate k-bit random number.
        """

        self.ecdlp = ecdlp
        self._hash_cls = hash_cls
        self._rnd_fn = rnd_fn

    def _hash_fn(self, data: bytes) -> bytes:
        hash_obj = self._hash_cls()
        hash_obj.update(data)
        return hash_obj.value()

    def _randint(self, a: int, b: int) -> int:
        bitlength = b.bit_length()
        while True:
            n = self._rnd_fn(bitlength)
            if n < a or n > b:
                continue
            return n

    def generate_keypair(self) -> Tuple[int, Tuple[int, int]]:
        """Generate key pair.

        Returns:
            int: secret key
            (int, int): public key, (xP, yP)
        """

        d = self._randint(1, self.ecdlp.n - 2)
        return d, self.ecdlp.kG(d)

    def get_pubkey(self, d: int) -> Tuple[int, int]:
        return self.ecdlp.kG(d)

    def verify_pubkey(self, x: int, y: int) -> bool:
        """Verify if a public key is valid.

        Args:
            x (int): x
            y (int): y

        Returns:
            (bool): Whether valid.
        """

        if self.ecdlp.isinf(x, y):
            return False

        if not self.ecdlp.isvalid(x, y):
            return False

        if not self.ecdlp.isinf(self.ecdlp.mul(self.ecdlp.n, x, y)):
            return False

        return True

    def _Z(self, id_: bytes, xP: int, yP: int) -> bytes:
        ENTL = len(id_) << 3
        if ENTL.bit_length() > 16:
            raise ValueError("ID bit length more than 2 bytes.")

        itob = self.ecdlp.itob

        Z = bytearray()
        Z.extend(ENTL.to_bytes(2, "big"))
        Z.extend(id_)
        Z.extend(itob(self.ecdlp.a))
        Z.extend(itob(self.ecdlp.b))
        Z.extend(itob(self.ecdlp.xG))
        Z.extend(itob(self.ecdlp.yG))
        Z.extend(itob(xP))
        Z.extend(itob(yP))

        return self._hash_fn(Z)

    def sign(self, message: bytes, d: int, id_: bytes, xP: int = None, yP: int = None) -> Tuple[int, int]:
        """Sign.

        Args:
            message (bytes): message to be signed.
            d (int): secret key.
            id_ (bytes): user id.
            xP (int): x of public key
            yP (int): y of public key

        Returns:
            (int, int): (r, s)
        """

        if xP is None or yP is None:
            xP, yP = self.get_pubkey(d)

        e = int.from_bytes(self._hash_fn(self._Z(id_, xP, yP) + message), "big")

        ecdlp = self.ecdlp
        n = self.ecdlp.n
        while True:
            k = self._randint(1, n - 1)
            x, _ = ecdlp.kG(k)

            r = (e + x) % n
            if r == 0 or (r + k == n):
                continue

            s = (_inv(1 + d, n) * (k - r * d)) % n
            if s == 0:
                continue

            return r, s

    def verify(self, message: bytes, r: int, s: int, id_: bytes, xP: int, yP: int) -> bool:
        """Verify.

        Args:
            message (bytes): Message to be verified.
            r (int): r
            s (int): s
            id_ (bytes): user id.
            xP (int): x of public key.
            yP (int): y of public key.

        Returns:
            bool: Whether OK.
        """

        ecdlp = self.ecdlp
        n = self.ecdlp.n

        if r < 1 or r > n - 1:
            return False

        if s < 1 or s > n - 1:
            return False

        t = (r + s) % n
        if t == 0:
            return False

        e = int.from_bytes(self._hash_fn(self._Z(id_, xP, yP) + message), "big")

        x, _ = ecdlp.add(*ecdlp.kG(s), *ecdlp.mul(t, xP, yP))
        if (e + x) % n != r:
            return False

        return True

    def _KDF(self, Z: bytes, klen: int) -> bytes:
        """
        Args:
            Z (bytes): secret byets.
            klen (int): key byte length
        """

        hash_fn = self._hash_fn
        v = self._hash_cls.hash_length()

        count, tail = divmod(klen, v)
        if count + (tail > 0) >= 0xffffffff:
            raise OverflowError("klen is too big.")

        K = bytearray()
        for ct in range(1, count + 1):
            K.extend(hash_fn(Z + ct.to_bytes(4, "big")))

        if tail > 0:
            K.extend(hash_fn(Z + (count + 1).to_bytes(4, "big"))[:tail])

        return bytes(K)

    def encrypt(self, plain: bytes, xP: int, yP: int) -> Tuple[Tuple[int, int], bytes, bytes]:
        """Encrypt

        Args:
            data (bytes): plain text to be encrypted.
            xP (int): x of public key.
            yP (int): y of public key.

        Returns:
            (int, int): C1, kG point
            bytes: C2, cipher
            bytes: C3, hash value

        Raises:
            ValueError: Infinite point encountered.

        The return order is `C1, C2, C3`, **NOT** `C1, C3, C2`.
        """

        while True:
            k = self._randint(1, self.ecdlp.n - 1)
            x1, y1 = self.ecdlp.kG(k)  # C1

            if self.ecdlp.isinf(*self.ecdlp.mul(self.ecdlp.h, xP, yP)):
                raise ValueError("Infinite point encountered.")

            x2, y2 = self.ecdlp.mul(k, xP, yP)
            x2 = self.ecdlp.itob(x2)
            y2 = self.ecdlp.itob(y2)

            t = self._KDF(x2 + y2, len(plain))
            if not any(t):
                continue

            C2 = bytes(map(lambda b1, b2: b1 ^ b2, plain, t))
            C3 = self._hash_fn(x2 + plain + y2)

            return (x1, y1), C2, C3

    def decrypt(self, x1: int, y1: int, C2: bytes, C3: bytes, d: int) -> bytes:
        """Decrypt.

        Args:
            x1 (int): x of C1 (kG point).
            y1 (int): y of C1 (kG point).
            C1 (bytes, bytes): kG point
            C2 (bytes): cipher
            C3 (bytes): hash value
            d (int): secret key.

        Returns:
            bytes: plain text.

        Raises:
            ValueError: Invalid C1 point, not on curve.
            ValueError: Infinite point encountered.
            ValueError: Invalid key stream.
            ValueError: Incorrect hash value.
        """

        if not self.ecdlp.isvalid(x1, y1):
            raise ValueError("Invalid C1 point, not on curve.")

        if self.ecdlp.isinf(*self.ecdlp.mul(self.ecdlp.h, x1, y1)):
            raise ValueError("Infinite point encountered.")

        x2, y2 = self.ecdlp.mul(d, x1, y1)
        x2 = self.ecdlp.itob(x2)
        y2 = self.ecdlp.itob(y2)

        t = self._KDF(x2 + y2, len(C2))
        if not any(t):
            raise ValueError("Invalid key stream.")

        M = bytes(map(lambda b1, b2: b1 ^ b2, C2, t))

        if self._hash_fn(x2 + M + y2) != C3:
            raise ValueError("Incorrect hash value.")

        return M
