import base64
import nacl.signing
import nacl.encoding
import base58

class KeyPair:
    def __init__(self, signing_key: nacl.signing.SigningKey):
        self.signing_key = signing_key
        self.verify_key = signing_key.verify_key
        self._did = self._generate_did()

    @classmethod
    def generate(cls):
        """Generate a new random Ed25519 keypair."""
        return cls(nacl.signing.SigningKey.generate())

    @classmethod
    def from_private_key(cls, private_key_b64: str):
        """Load from a base64 encoded private key."""
        private_key_bytes = base64.b64decode(private_key_b64)
        # nacl.signing.SigningKey expects the seed (32 bytes) or full key (64 bytes)
        # tweetnacl usually returns the full 64 bytes as secretKey
        if len(private_key_bytes) == 64:
            # Extract seed (first 32 bytes)
            seed = private_key_bytes[:32]
            return cls(nacl.signing.SigningKey(seed))
        elif len(private_key_bytes) == 32:
            return cls(nacl.signing.SigningKey(private_key_bytes))
        else:
            raise ValueError("Invalid private key length")

    def _generate_did(self) -> str:
        """Generate did:key identifier."""
        pub_key_bytes = self.verify_key.encode(encoder=nacl.encoding.RawEncoder)
        
        # Ed25519 multicodec prefix is 0xed 0x01 (varint for 0xed)
        prefix = b'\xed\x01'
        prefixed_key = prefix + pub_key_bytes
        
        # Base58-BTC encode (z-prefix is handled by multiformats, but here we just encode)
        # Wait, multiformats base58btc usually adds a 'z' prefix to the string.
        # The TS code uses `base58btc.encode(multicodecPubkey)`.
        # Let's check if `base58` library adds 'z'. It usually doesn't.
        encoded = base58.b58encode(prefixed_key).decode('utf-8')
        return f"did:key:z{encoded}"

    @property
    def did(self) -> str:
        return self._did

    @property
    def public_key_b64(self) -> str:
        return base64.b64encode(self.verify_key.encode()).decode('utf-8')

    @property
    def private_key_b64(self) -> str:
        # Return full 64-byte key to match tweetnacl behavior if possible,
        # but pynacl usually just gives the seed. 
        # TweetNaCl secretKey is seed (32) + pubKey (32).
        seed = self.signing_key.encode()
        pub = self.verify_key.encode()
        return base64.b64encode(seed + pub).decode('utf-8')

    def sign(self, message: bytes) -> str:
        """Sign message and return base64 signature."""
        signed = self.signing_key.sign(message)
        # signed.signature is the raw bytes
        return base64.b64encode(signed.signature).decode('utf-8')
