from PIL import Image
from pyzbar.pyzbar import decode
import itertools
from tqdm import tqdm
import numpy as np
from multiprocessing import Pool
from math import perm
import sys

BOX_SIZE = 2
TILE = 78
# 7x7 finder, upscaled to 14x14 (black=1, white=0)
TEMPLATE = np.kron(np.array([
    [1,1,1,1,1,1,1],
    [1,0,0,0,0,0,1],
    [1,0,1,1,1,0,1],
    [1,0,1,1,1,0,1],
    [1,0,1,1,1,0,1],
    [1,0,0,0,0,0,1],
    [1,1,1,1,1,1,1],
], dtype=np.uint8), np.ones((BOX_SIZE, BOX_SIZE), np.uint8))
H, W = TEMPLATE.shape  # 14x14


def decodeQRCode(image : Image.Image) -> str | None:
    # 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 cut_pieces(image):
    # Get the size of the image
    width, height = image.size

    if width == height:
        NUM_PIECES_WIDTH = 3
    else:
        NUM_PIECES_WIDTH = 4

    NUM_PIECES_HEIGHT = 3

    pieces = []

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

    return pieces


def to_bw01(img):
    # Force 1-bit, then map black->1, white->0
    a = np.array(img.convert("1"))  # {0,255}
    return (a == 0).astype(np.uint8)


def finder_corner_exact(tile_bw):
    crops = {
        'TL': tile_bw[0:H,            0:W            ],
        'TR': tile_bw[0:H,            TILE-W:TILE    ],
        'BL': tile_bw[TILE-H:TILE,    0:W            ],
    }
    for corner, crop in crops.items():
        if np.array_equal(crop, TEMPLATE):
            return corner
    return None  # no exact finder in this tile


def three_tiles_with_finders(tiles):
    hits = []
    for tile in tiles:
        corner = finder_corner_exact(to_bw01(tile))
        if corner is not None:
            hits.append((tile, corner))
            if len(hits) == 3:
                break
    return hits

def solve_combo(combo):
    """
    Function to be executed by each process in the pool.
    """
    a, b, c, d, e, f = combo
    # Assuming `new_image` is a base image you can copy or an in-memory image object.
    # You may need to pass `new_image` to the function or have it loaded by each process.
    local_image = new_image.copy() # Make a copy to avoid race conditions
    local_image.paste(a, (78, 0))
    local_image.paste(b, (0, 78))
    local_image.paste(c, (78, 78))
    local_image.paste(d, (156, 78))
    local_image.paste(e, (78, 156))
    local_image.paste(f, (156, 156))
    
    output = decodeQRCode(local_image)
    if output:
        return output
    return None


def main(img_path="./output/secret_message.png", verbose=True):
    img = Image.open(img_path)
    global new_image
    new_image = Image.new("RGB", (img.height, img.height))


    pieces = cut_pieces(img)
    corners = three_tiles_with_finders(pieces)

    for tile, corner in corners:
        if corner == 'TL':
            new_image.paste(tile, (0,0))
        elif corner == 'TR':
            new_image.paste(tile, (156,0))
        elif corner == 'BL':
            new_image.paste(tile, (0,156))
        else:
            sys.exit()


    pieces = [piece for piece in pieces if piece not in [tile for tile, _ in corners]]
    possible_combos = itertools.permutations(pieces, 6)

    result = None
    
    # Use `multiprocessing.Pool` with the number of available cores
    with Pool() as pool:
        # `imap_unordered` returns results as they are completed.
        # This is a lazy iterator, so it doesn't create all combos at once.
        # `tqdm` can still be used to show progress.
        for result in tqdm(pool.imap_unordered(solve_combo, possible_combos), total=perm(len(pieces), 6), disable=(not verbose)):
            if result:
                if verbose:
                    print(f"QR Code decoded: {result}")
                pool.terminate() # Stop all worker processes
                break

    return result


if __name__ == "__main__":
    if len(sys.argv) == 2:
        main(sys.argv[1])
    else:
        main()