#!/usr/bin/env python3
import argparse
import os.path
import PIL.Image
import PIL.ImageDraw
import PIL.ImageFont
import qrcode
import re
import sys

script_dir = os.path.dirname(os.path.realpath(__file__))

BASE36 = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'

PAYLOAD_VERSION = 0
PAYLOAD_FLAGS = 2  # 2=IP, 4=BLE, 8=IP_WAC

def gen_homekit_setup_uri(category, password, setup_id, version=0, reserved=0, flags=2):
    payload = 0
    payload |= (version & 0x7)

    payload <<= 4
    payload |=(reserved & 0xf)  # reserved bits

    payload <<= 8
    payload |= (category & 0xff)

    payload <<= 4
    payload |= (flags & 0xf)

    payload <<= 27
    payload |= (int(password.replace('-', '')) & 0x7fffffff)

    encodedPayload = ''
    for _ in range(9):
        encodedPayload += BASE36[payload % 36]
        payload //= 36

    return 'X-HM://%s%s' % (''.join(reversed(encodedPayload)), setup_id)


def gen_homekit_qrcode(setup_uri, password, transparent=False):
    back_color = 'transparent' if transparent else 'white'

    # open template
    template = PIL.Image.open(os.path.join(script_dir, 'qrcode.png'))

    img = PIL.Image.new('RGBA', template.size, None)
    draw = PIL.ImageDraw.Draw(img)

    if not transparent:
        draw.rounded_rectangle((5, 5, img.width-5, img.height-5), 60, fill='white')

    img.paste(template, (0, 0), template)

    # add QR code to it
    code = qrcode.QRCode(version=2, border=0, box_size=12.5,
                         error_correction=qrcode.constants.ERROR_CORRECT_Q)
    code.add_data(setup_uri)
    img.paste(
        code.
            make_image(fill_color='black',
                       back_color='transparent' if transparent else 'white').
            get_image().
            resize((320, 320)),
        (40, 180)
    )

    # add password digits
    setup_code = password.replace('-', '')

    font = PIL.ImageFont.truetype(os.path.join(script_dir, 'SF-Mono-Bold.otf'), 69)

    for i in range(4):
        draw.text((173 + i*49, 38-13), setup_code[i], font=font, fill=(0, 0, 0))
        draw.text((173 + i*49, 88), setup_code[i+4], font=font, fill=(0, 0, 0))

    return img


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--transparent', action='store_true')
    parser.add_argument('category', type=int)
    parser.add_argument('password')
    parser.add_argument('setup_id')
    parser.add_argument('output_image')

    args = parser.parse_args()

    if not re.match('\d{3}-\d{2}-\d{3}', args.password):
        raise ValueError('Invalid password')

    if not re.match('[0-9A-Z]{4}', args.setup_id):
        raise ValueError('Invalid setup ID')

    setupURI = gen_homekit_setup_uri(args.category, args.password, args.setup_id)

    qrcodeImage = gen_homekit_qrcode(
        setupURI, args.password, transparent=args.transparent
    )

    qrcodeImage.save(args.output_image)


if __name__ == '__main__':
    main()