#!/Users/cato/Code/Cato447/ctf/insomnihack24/misc/puzzled/.venv/bin/python3

from PIL import Image
import random
import itertools
from pyzbar.pyzbar import decode
import re


def cutPieces(image_path, output_name_scheme):
    # Open the image
    original_image = Image.open(image_path)

    # Get the size of the image
    width, height = original_image.size
    assert width == height

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

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

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

        # Save the cropped piece
        output_path = output_name_scheme.format(i)
        cropped_piece.save(output_path)


def generateQRCodes(topLeft, topRight, bottomLeft, pieces):
    # Create a new image to store rearranged pieces
    new_image = Image.new("RGB", (234, 234))

    # Paste the QR codes onto the new image
    new_image.paste(topLeft, (0, 0))
    new_image.paste(topRight, (156, 0))
    new_image.paste(bottomLeft, (0, 156))

    print(pieces)

    possible_combos = itertools.permutations(pieces, 6)

    i = 0
    for a, b, c, d, e, f in possible_combos:
        # Paste the remaining pieces onto the new image
        new_image.paste(Image.open(a), (78, 0))
        new_image.paste(Image.open(b), (156, 78))
        new_image.paste(Image.open(c), (0, 78))
        new_image.paste(Image.open(d), (78, 78))
        new_image.paste(Image.open(e), (156, 156))
        new_image.paste(Image.open(f), (78, 156))

        # Save the rearranged image
        new_image.save(f"solutions/possible_qr{i}.png")
        i += 1


def decodeQRCode(image_path) -> str | None:
    image = Image.open(image_path)

    # Decode the QR code
    decoded_objects = decode(image)

    # Check if any QR code was found
    if decoded_objects:
        # Iterate through all the detected QR codes and print their data
        for obj in decoded_objects:
            return obj.data.decode("utf-8")
    else:
        return None


def shuffle3x3GridImage(image_path, output_path):
    # Open the image
    original_image = Image.open(image_path)

    # Get the size of the image
    width, height = original_image.size
    assert width == 234 and height == 234

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

    # Create a new image to store rearranged pieces
    new_image = Image.new("RGB", (width, height))

    # Create a list to store shuffled grid positions
    positions = [(x, y) for x in range(3) for y in range(3)]
    random.shuffle(positions)

    # Split the original image into 16 pieces and rearrange them
    for i, pos in enumerate(positions):
        # Calculate the coordinates of the current grid cell
        x1 = pos[0] * cell_width
        y1 = pos[1] * cell_height
        x2 = x1 + cell_width
        y2 = y1 + cell_height

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

        # Calculate the coordinates to paste the cropped piece
        new_x = (i % 3) * cell_width
        new_y = (i // 3) * cell_height

        # Paste the cropped piece onto the new image
        new_image.paste(cropped_piece, (new_x, new_y))

    # Save the rearranged image
    new_image.save(output_path)


# key = "Not the real key!"
# random.seed(key)
# shuffle3x3GridImage("fullVersion25QRcode.png", "puzzledVersion25QRcode.png")

# cutPieces("puzzledVersion25QRcode.png", "puzzledVersion25QRcode_piece{}.png")

flag_pattern = r"INS{.*}"
pattern = re.compile(flag_pattern)

paths = [f"solutions/possible_qr{i}.png" for i in range(720)]
for path in paths:
    data = decodeQRCode(path)
    if data is None:
        continue
    possible_flag = pattern.search(data)
    if possible_flag:
        print(possible_flag.group())
        break