from hashlib import sha256

import argparse

from chik_base.bls12_381 import BLSPublicKey, BLSSignature


DEFAULT_PARENT_COIN_ID = sha256(b"fake_id").digest()


def bytes32_fromhex(s: str) -> bytes:
    try:
        b = bytes.fromhex(s)
        if len(b) == 32:
            return b
    except Exception:
        pass
    raise ValueError("not 64 bytes of hex")


DESCRIPTION = "Proof of secret exponent request verifier."
EPILOG = (
    'Note: signature output from hsmk is in "qrint" form. '
    "Use `qrint -H` to convert it to hex."
)


def create_parser():
    parser = argparse.ArgumentParser(description=DESCRIPTION, epilog=EPILOG)
    parser.add_argument(
        "--network-id",
        type=bytes.fromhex,
        default=b"",
        help="hex representing network id",
    )
    parser.add_argument("bech32m_public_key", help="bech32m-encoded public key")
    parser.add_argument("message", help="message to embed in challenge")
    parser.add_argument(
        "challenge_coin_id",
        type=bytes32_fromhex,
        help="64 character hex representing coin id generated by challenge",
    )
    parser.add_argument(
        "signature",
        type=bytes.fromhex,
        help="hex representing signature response to challenge",
    )
    return parser


def main():
    parser = create_parser()
    args = parser.parse_args()

    ONE = bytes([1])
    TWO = bytes([2])

    coin_name = args.challenge_coin_id
    pk = BLSPublicKey.from_bech32m(args.bech32m_public_key)

    try:
        sig = BLSSignature.from_bytes(args.signature)
    except Exception:
        pass

    def s256(b):
        return sha256(b).digest()

    hashed_message = s256(TWO + s256(ONE + ONE) + s256(ONE + args.message.encode()))

    signed_message = hashed_message + coin_name

    r = sig.verify([(pk, signed_message)])

    if r:
        print("message signed correctly")
    else:
        print("BAD SIGNATURE")


if __name__ == "__main__":
    main()
