import argparse
import sys

from ecies import encrypt, decrypt
from ecies.utils import generate_eth_key
from ecies.__version__ import __description__


def readablize(b: bytes) -> str:
    try:
        return b.decode()
    except ValueError:
        return b.hex()


def main():
    parser = argparse.ArgumentParser(description=__description__)

    parser.add_argument('-e', '--encrypt', action='store_true', help='encrypt with public key, not compatible with -d')
    parser.add_argument('-d', '--decrypt', action='store_true', help='decrypt with private key, not compatible with -e')
    parser.add_argument('-g', '--generate', action='store_true', help='generate ethereum key pair')
    parser.add_argument(
        '-k', '--key', type=argparse.FileType('r'), help='public or private key file'
    )

    parser.add_argument(
        '-D', '--data', nargs='?', type=argparse.FileType('rb'), default=sys.stdin,
        help='file to encrypt or decrypt, if not specified, it will read from stdin'
    )

    parser.add_argument(
        '-O', '--out', nargs='?', type=argparse.FileType('wb'), default=sys.stdout,
        help='encrypted or decrypted file, if not specified, it will write to stdout'
    )

    args = parser.parse_args()
    if args.generate:
        k = generate_eth_key()
        prv, pub, addr = k.to_hex(), k.public_key.to_hex(), k.public_key.to_checksum_address()
        print('Private: {}\nPublic: {}\nAddress: {}'.format(prv, pub, addr))
        return

    if args.encrypt == args.decrypt:
        parser.print_help()
        return

    if not args.key:
        parser.print_help()
        return

    key = args.key.read().strip()
    if args.encrypt:
        plaintext = args.data.read()
        if isinstance(plaintext, str):
            plaintext = plaintext.encode()
        data = encrypt(key, plaintext)
        if args.out == sys.stdout:
            data = data.hex()
    else:
        ciphertext = args.data.read()
        if isinstance(ciphertext, str):
            # if not bytes, suppose hex string
            ciphertext = bytes.fromhex(ciphertext.strip())
        data = decrypt(key, ciphertext)
        if args.out == sys.stdout:
            # if binary data, print hex; if not, print utf8
            data = readablize(data)

    args.out.write(data)


if __name__ == '__main__':
    main()
