from logging import disable
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="student_fair/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
    
    MAX_WORKER = 6

    # Use `multiprocessing.Pool` with the number of available cores
    with Pool(MAX_WORKER) 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__":
    main()