import math
from pathlib import Path

import qrcode
from PIL import Image, ImageDraw, ImageFont

def create_cassette_side(qrs, title, output_path):
    """
    Composites exactly 4 square QR images into a 2x2 grid with a large title at the top.
    """
    qr_size = qrs[0].width
    
    # Grid dimensions (2x2)
    grid_w = qr_size * 2
    grid_h = qr_size * 2
    
    # Allocate space for the title
    title_h = 80
    total_w = grid_w
    total_h = grid_h + title_h
    
    img = Image.new("RGB", (total_w, total_h), "white")
    draw = ImageDraw.Draw(img)
    
    # Attempt to load a large, readable font (fallback gracefully for older Pillow versions)
    try:
        font = ImageFont.truetype("arial.ttf", 60)
    except OSError:
        try:
            font = ImageFont.load_default(size=60)
        except TypeError:
            font = ImageFont.load_default()
            
    # Calculate text bounding box to center the title
    try:
        bbox = draw.textbbox((0, 0), title, font=font)
        tw = bbox[2] - bbox[0]
        th = bbox[3] - bbox[1]
    except AttributeError:
        # Fallback for very old PIL versions
        tw, th = 120, 20
        
    draw.text(((total_w - tw) // 2, (title_h - th) // 2), title, fill="black", font=font)
    
    # Paste QRs into the 2x2 layout underneath the title
    img.paste(qrs[0], (0, title_h))
    img.paste(qrs[1], (qr_size, title_h))
    img.paste(qrs[2], (0, title_h + qr_size))
    img.paste(qrs[3], (qr_size, title_h + qr_size))
    
    img.save(output_path)

def main():
    raw_file = Path("music.raw")
    if not raw_file.exists():
        raise FileNotFoundError("Input file 'music.raw' not found in the current directory.")
        
    data = raw_file.read_bytes()
    
    # Calculate chunk size to ensure exactly 8 splits
    chunk_size = math.ceil(len(data) / 8)
    chunks = [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)]
    
    # Pad out the array just in case the file was extremely small or empty
    while len(chunks) < 8:
        chunks.append(b"")
        
    qr_images = []
    
    print("Generating QR Codes...")
    for i in range(8):
        # Prefix the payload with a 1-byte ID (1 through 8)
        payload = bytes([i + 1]) + chunks[i]
        
        qr = qrcode.QRCode(
            version=None,
            error_correction=qrcode.constants.ERROR_CORRECT_L, # Lowest error correction for maximum density
            box_size=8,
            border=3
        )
        
        # optimize=0 guarantees strict byte encoding without accidental alphanumeric mixing
        qr.add_data(payload, optimize=0) 
        qr.make(fit=True)
        
        qr_img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
        qr_images.append(qr_img)
        
    # Standardize image dimensions to prevent grid distortion 
    # (The last QR code might generate a smaller version matrix if the trailing payload is small)
    max_size = max(img.width for img in qr_images)
    
    uniform_qrs = []
    for img in qr_images:
        if img.width != max_size:
            padded = Image.new("RGB", (max_size, max_size), "white")
            # Center the smaller QR inside the max-size bounding box
            padded.paste(img, ((max_size - img.width) // 2, (max_size - img.height) // 2))
            uniform_qrs.append(padded)
        else:
            uniform_qrs.append(img)
            
    print("Compositing cassettes...")
    create_cassette_side(uniform_qrs[0:4], "SIDE A", "cassette_side_a.png")
    create_cassette_side(uniform_qrs[4:8], "SIDE B", "cassette_side_b.png")
    
    print("Done! Saved cassette_side_a.png and cassette_side_b.png.")

if __name__ == "__main__":
    main()