# pqbit/pqbit/dilithium.py

import logging
from pqc.sign.dilithium5 import keypair as _generate_keypair, sign as _sign, verify as _verify

logger = logging.getLogger("pqbit.dilithium")

def dilithium_keypair():
    """
    Generates a key pair for Dilithium5.

    Returns:
        tuple: (public_key: bytes, secret_key: bytes)
    """
    pk, sk = _generate_keypair()
    logger.info("Dilithium5 keypair generated.")
    return pk, sk

def dilithium_sign(message, sk):
    """
    Signs a message using the secret key.

    Args:
        message (bytes): Message to be signed
        sk (bytes): Secret key generated by dilithium_keypair()

    Returns:
        bytes: Digital signature

    Raises:
        TypeError: If message or sk are not of type bytes
    """
    if not isinstance(message, bytes):
        raise TypeError("dilithium_sign expects message of type bytes")
    if not isinstance(sk, bytes):
        raise TypeError("dilithium_sign expects sk of type bytes")

    signature = _sign(message, sk)
    logger.info("Message signed with Dilithium5.")
    return signature

def dilithium_verify(message, signature, pk):
    """
    Verifies the signature of a message using the public key.

    Args:
        message (bytes): Original message
        signature (bytes): Signature generated by dilithium_sign()
        pk (bytes): Public key generated by dilithium_keypair()

    Returns:
        bool: True if the signature is valid, False otherwise

    Raises:
        TypeError: If any argument is not of type bytes
    """
    if not isinstance(message, bytes):
        raise TypeError("dilithium_verify expects message of type bytes")
    if not isinstance(signature, bytes):
        raise TypeError("dilithium_verify expects signature of type bytes")
    if not isinstance(pk, bytes):
        raise TypeError("dilithium_verify expects pk of type bytes")

    try:
        _verify(signature, message, pk)
        logger.info("Signature verification result: valid")
        return True
    except Exception:
        logger.info("Signature verification result: invalid")
        return False

