import qrcode
from PIL import Image
import random
import sys

from templates import payload, decoy_payload

NUM_PIECES = 3

def generate_qrcode(data : str):
    qr = qrcode.QRCode(version=25, box_size=2, error_correction=qrcode.ERROR_CORRECT_L, border=0)
    qr.add_data(data)
    return qr.make_image()


def cut_pieces(image):
    # Get the size of the image
    width, height = image.size
    assert width == height

    # Calculate the size of each grid cell
    cell_width = width // NUM_PIECES 
    cell_height = height // NUM_PIECES 

    pieces = []

    # Split the original image into 9 pieces
    for i in range(NUM_PIECES**2):
        # Calculate the coordinates of the current grid cell
        x1 = (i % NUM_PIECES) * cell_width
        y1 = (i // NUM_PIECES) * cell_height
        x2 = x1 + cell_width
        y2 = y1 + cell_height

        # Crop the corresponding part of the original image
        cropped_piece = image.crop((x1, y1, x2, y2))
        pieces.append(cropped_piece)

    return pieces


def scramble_qr_code(payloadList, decoyList, password):
    height_piece, width_piece = payloadList[0].size

    new_qr_code = Image.new("RGB", (width_piece * (NUM_PIECES+1), height_piece * NUM_PIECES))
    random.seed(password)
    payloadList += random.sample(decoyList, NUM_PIECES)
    assert len(payloadList) == 12
    random.shuffle(payloadList)

    for x,y in [(x,y) for x in range(NUM_PIECES+1) for y in range(NUM_PIECES)]:
        piece = payloadList[x + (NUM_PIECES+1) * y]
        new_qr_code.paste(piece, (x * width_piece, y * height_piece))
    return new_qr_code


def generate_scrambled_qrcode(passphrase):
    qr_code = generate_qrcode(payload)
    decoy_qr_code = generate_qrcode(decoy_payload)
    qr_pieces = cut_pieces(qr_code)
    decoy_pieces = [piece for id, piece in enumerate(cut_pieces(decoy_qr_code)) if id not in [0,2,6]]

    scrambled_qr_code = scramble_qr_code(qr_pieces, decoy_pieces, passphrase)
    scrambled_qr_code.save("scrambled_test_qr_code.png")


def main():
    if len(sys.argv) != 2:
        print("missing password")
        sys.exit()
    generate_scrambled_qrcode(sys.argv[1])

if __name__ == "__main__":
    main()