import argparse
from typing import List
import qrcode
from PIL import Image
from pathlib import Path
import random

NUM_PIECES = 3
TILE = 78

def generate_qrcode(data: str) -> Image:
    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: Image) -> List[Image]:
    pieces = []

    for x, y in [(x, y) for x in range(NUM_PIECES) for y in range(NUM_PIECES)]:
        pieces.append(image.crop((x * TILE, y * TILE, x * TILE + TILE, y * TILE + TILE)))

    return pieces


def scramble_qr_code(payloadList: List[Image], seed: str) -> Image:
    height_piece, width_piece = payloadList[0].size
    random.seed(seed)
    scrambled_width = 3
    random.shuffle(payloadList)

    new_qr_code = Image.new("RGB", (width_piece * scrambled_width, height_piece * NUM_PIECES))
    for x, y in [(x, y) for x in range(scrambled_width) for y in range(NUM_PIECES)]:
        piece = payloadList[x + scrambled_width * y]
        new_qr_code.paste(piece, (x * width_piece, y * height_piece))

    return new_qr_code


def generate_scrambled_qrcode(seed: str, output_dir: Path, name: str):
    with open("./secret_message.txt", "r") as secret:
        message = secret.read()
     
    qr_code = generate_qrcode(message + seed)
    
    qr_pieces = cut_pieces(qr_code)
    scrambled_qr_code = scramble_qr_code(qr_pieces, seed)

    output_dir.mkdir(parents=True, exist_ok=True)
    scrambled_qr_code.save(output_dir / name)


def main():
    parser = argparse.ArgumentParser(description="Generates a scrambled QR code with optional decoy pieces.")
    parser.add_argument("seed", type=str, help="The password to scramble the QR code.")
    parser.add_argument("output_dir", type=Path, help="Directory where the output will be stored")
    parser.add_argument("-o", "--output", type=str, default="secret_message.png", help="The name of the output image file.")
    
    args = parser.parse_args()

    generate_scrambled_qrcode(
        seed=args.seed, 
        output_dir=args.output_dir,
        name=args.output
    )


if __name__ == "__main__":
    main()