# Copyright © Supra
# Parts of the project are originally copyright © Aptos Foundation
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import cast

from supra_sdk import asymmetric_crypto, ed25519
from supra_sdk.bcs import Deserializer, Serializer


class PublicKey(asymmetric_crypto.PublicKey):
    ED25519: int = 0

    variant: int
    public_key: asymmetric_crypto.PublicKey

    def __init__(self, public_key: asymmetric_crypto.PublicKey):
        if isinstance(public_key, ed25519.PublicKey):
            self.variant = PublicKey.ED25519
        else:
            raise NotImplementedError()
        self.public_key = public_key

    def to_crypto_bytes(self) -> bytes:
        serializer = Serializer()
        self.serialize(serializer)
        return serializer.output()

    def verify(self, data: bytes, signature: asymmetric_crypto.Signature) -> bool:
        # Convert signature to the original signature
        sig = cast(Signature, signature)

        return self.public_key.verify(data, sig.signature)

    @staticmethod
    def deserialize(deserializer: Deserializer) -> PublicKey:
        variant = deserializer.uleb128()

        if variant == PublicKey.ED25519:
            public_key: asymmetric_crypto.PublicKey = ed25519.PublicKey.deserialize(
                deserializer
            )
        else:
            raise Exception(f"Invalid type: {variant}")

        return PublicKey(public_key)

    def serialize(self, serializer: Serializer):
        serializer.uleb128(self.variant)
        serializer.struct(self.public_key)


class Signature(asymmetric_crypto.Signature):
    ED25519: int = 0

    variant: int
    signature: asymmetric_crypto.Signature

    def __init__(self, signature: asymmetric_crypto.Signature):
        if isinstance(signature, ed25519.Signature):
            self.variant = Signature.ED25519
        else:
            raise NotImplementedError()
        self.signature = signature

    def get_null_signature(self) -> Signature:
        if self.variant == Signature.ED25519:
            return Signature(ed25519.Signature.get_null_signature())
        else:
            raise NotImplementedError()

    @staticmethod
    def deserialize(deserializer: Deserializer) -> Signature:
        variant = deserializer.uleb128()

        if variant == Signature.ED25519:
            signature: asymmetric_crypto.Signature = ed25519.Signature.deserialize(
                deserializer
            )
        else:
            raise Exception(f"Invalid type: {variant}")

        return Signature(signature)

    def serialize(self, serializer: Serializer):
        serializer.uleb128(self.variant)
        serializer.struct(self.signature)


class MultiPublicKey(asymmetric_crypto.PublicKey):
    keys: list[PublicKey]
    threshold: int

    MIN_KEYS = 2
    MAX_KEYS = 32
    MIN_THRESHOLD = 1

    def __init__(self, keys: list[asymmetric_crypto.PublicKey], threshold: int):
        assert self.MIN_KEYS <= len(keys) <= self.MAX_KEYS, (
            f"Must have between {self.MIN_KEYS} and {self.MAX_KEYS} keys."
        )
        assert self.MIN_THRESHOLD <= threshold <= len(keys), (
            f"Threshold must be between {self.MIN_THRESHOLD} and {len(keys)}."
        )

        # Ensure keys are wrapped
        self.keys = []
        for key in keys:
            if isinstance(key, PublicKey):
                self.keys.append(key)
            else:
                self.keys.append(PublicKey(key))

        self.threshold = threshold

    def __str__(self) -> str:
        return f"{self.threshold}-of-{len(self.keys)} Multi key"

    def verify(self, data: bytes, signature: asymmetric_crypto.Signature) -> bool:
        try:
            total_sig = cast(MultiSignature, signature)
            assert self.threshold <= len(total_sig.signatures), (
                f"Insufficient signatures, {self.threshold} > {len(total_sig.signatures)}"
            )

            for idx, signature in total_sig.signatures:
                assert len(self.keys) > idx, (
                    f"Signature index exceeds available keys {len(self.keys)} < {idx}"
                )
                assert self.keys[idx].verify(data, signature), (
                    "Unable to verify signature"
                )

        except Exception:
            return False
        return True

    @staticmethod
    def from_crypto_bytes(indata: bytes) -> MultiPublicKey:
        deserializer = Deserializer(indata)
        return deserializer.struct(MultiPublicKey)

    def to_crypto_bytes(self) -> bytes:
        serializer = Serializer()
        serializer.struct(self)
        return serializer.output()

    @staticmethod
    def deserialize(deserializer: Deserializer) -> MultiPublicKey:
        keys = deserializer.sequence(PublicKey.deserialize)
        threshold = deserializer.u8()
        return MultiPublicKey(keys, threshold)

    def serialize(self, serializer: Serializer):
        serializer.sequence(self.keys, Serializer.struct)
        serializer.u8(self.threshold)


class MultiSignature(asymmetric_crypto.Signature):
    signatures: list[tuple[int, Signature]]
    BITMAP_NUM_OF_BYTES: int = 4

    def __init__(self, signatures: list[tuple[int, asymmetric_crypto.Signature]]):
        # Sort first to ensure no issues in order
        # signatures.sort(key=lambda x: x[0])
        self.signatures = []
        for index, signature in signatures:
            assert index < self.BITMAP_NUM_OF_BYTES * 8, (
                "bitmap value exceeds maximum value"
            )
            if isinstance(signature, Signature):
                self.signatures.append((index, signature))
            else:
                self.signatures.append((index, Signature(signature)))

    def __eq__(self, other: object):
        if not isinstance(other, MultiSignature):
            return NotImplemented
        return self.signatures == other.signatures

    def __str__(self) -> str:
        return f"{self.signatures}"

    @staticmethod
    def deserialize(deserializer: Deserializer) -> MultiSignature:
        signatures = deserializer.sequence(Signature.deserialize)
        deserializer.uleb128()
        bitmap = deserializer.u32()
        num_bits = MultiSignature.BITMAP_NUM_OF_BYTES * 8
        sig_index = 0
        indexed_signatures = []

        for i in range(num_bits):
            has_signature = (bitmap & index_to_bitmap_value(i)) != 0
            if has_signature:
                indexed_signatures.append((i, signatures[sig_index]))
                sig_index += 1

        return MultiSignature(signatures)

    def serialize(self, serializer: Serializer):
        actual_sigs = []
        bitmap = 0

        for i, signature in self.signatures:
            bitmap |= index_to_bitmap_value(i)
            actual_sigs.append(signature)

        serializer.sequence(actual_sigs, Serializer.struct)
        serializer.uleb128(self.BITMAP_NUM_OF_BYTES)
        serializer.u32(bitmap)

    def unset_signatures(self):
        for i, (pubkey_index, signature) in enumerate(self.signatures):
            self.signatures[i] = (pubkey_index, signature.get_null_signature())


def index_to_bitmap_value(i: int) -> int:
    bit = i % 8
    byte = i // 8
    return (128 >> bit) << (byte * 8)
