import base58
import ecdsa
from hashlib import sha256


class KeyExtension:
    def get_private_key_from_wif(self, wif: str) -> bytes:
        """
        Converts a WIF private key to a byte array.

        :param wif: WIF private key in string format.
        :return: Private key in byte format (32 bytes).
        :raises ValueError: If the WIF key is incorrect.
        """
        assert not self.is_empty(wif)

        decoded = base58.b58decode_check(wif)
        if len(decoded) != 34 or decoded[0] != 0x80 or decoded[-1] != 0x01:
            raise ValueError("Incorrect WIF private key")

        private_key = decoded[1:-1]
        return private_key

    def get_public_key(self, private_key: bytes) -> bytes:
        """
        Extract public key from Private key

        :param private_key: Private key in byte format (32 bytes).
        :return: compressed public key in byte format (33 bytes).
        :raises ValueError: If the private_key key is empty or null.
        """
        assert not self.is_empty(private_key)

        if len(private_key) != 32:
            raise ValueError(f"Incorrect len of private key, Expected: 32, Actual: {len(private_key)}")

        public_key = ecdsa.SigningKey.from_string(private_key, curve=ecdsa.NIST256p).get_verifying_key()
        compressed_public_key = public_key.to_string("compressed")
        return compressed_public_key

    @staticmethod
    def is_empty(sequence_symbols: bytes | str):
        if len(sequence_symbols) == 0 or sequence_symbols is None:
            raise ValueError(f"Empty sequence symbols of key: {sequence_symbols}")
        return False