#!/usr/bin/env python3
"""
Picpak BLE Direct Uploader
Allows uploading 400x300 pre-dithered images directly to the Picpak e-ink unit over BLE.
Dependencies: python3 -m pip install bleak pillow numpy

Usage:
  python3 uploader.py --scan
  python3 uploader.py --address "YOUR-DEVICE-UUID-OR-MAC" --image "card.png"
  python3 uploader.py --address "YOUR-DEVICE-UUID-OR-MAC" --folder "/path/to/exported/folder"
"""

import os
import re
import glob
import sys
import argparse
import asyncio
import hashlib
from PIL import Image
import numpy as np

try:
    from bleak import BleakScanner, BleakClient
except ImportError:
    print("❌ Error: 'bleak' library is missing.")
    print("👉 Please install dependencies: python3 -m pip install bleak pillow numpy --break-system-packages")
    sys.exit(1)

# Default Picpak Image Data Write Characteristic UUID (FF01)
DEFAULT_CHAR_UUID = "0000ff01-0000-1000-8000-00805f9b34fb"

# Protocol constants
MAGIC_START = 0xAA
OP_CHUNK_DATA = 0x01
OP_COMMIT_MD5 = 0x04
MAGIC_END = 0xFF
CHUNK_SIZE_LIMIT = 236  # Safe payload chunk size limit per packet

def parse_slot_number(filename, default_slot=1):
    """Tries to extract slot number from common filename formats (e.g. batch_12, slot_5, etc.)"""
    # Look for patterns like '_batch_12_', 'batch_12_', '_12_', etc.
    match = re.search(r'batch_(\d+)', filename, re.IGNORECASE)
    if match:
        return int(match.group(1))
    match = re.search(r'slot_(\d+)', filename, re.IGNORECASE)
    if match:
        return int(match.group(1))
    # General number match if bordered by delimiters
    match = re.search(r'_(\d+)_', filename)
    if match:
        return int(match.group(1))
    return default_slot

def png_to_picpak_bin(png_path):
    """
    Reads a 400x300 PNG image and converts it into Picpak's 2-bit raw format.
    Also applies vertical and horizontal flipping to match the physical orientation of the panel.
    Color indices:
    00: Black
    01: White
    10: Yellow
    11: Red
    Packs 4 pixels per byte. Total binary size must equal 30,000 bytes.
    """
    try:
        img = Image.open(png_path).convert('RGB')
    except Exception as e:
        print(f"❌ Failed to open image: {e}")
        sys.exit(1)

    if img.size != (400, 300):
        print(f"⚠️ Warning: Image size is {img.size[0]}x{img.size[1]}. Picpak requires exactly 400x300.")
        print("👉 Resizing to 400x300...")
        img = img.resize((400, 300), Image.Resampling.NEAREST)

    pixels = np.array(img)
    
    # Flip vertically to match Gu(e, t, n) in the official client protocol
    pixels = np.flipud(pixels)
    
    bin_data = bytearray()
    
    for y in range(300):
        for x in range(0, 400, 4):
            byte_val = 0
            for i in range(4):
                r, g, b = pixels[y, x + i]
                # Determine closest Spectra color code using optimized thresholds
                # 115 is the sweet spot to prevent thin line clipping without causing font bleed
                # 0: Black, 1: White, 2: Yellow, 3: Red
                if r < 115 and g < 115 and b < 115:       # Black
                    color_idx = 0
                elif r > 180 and g > 180 and b < 80:     # Yellow
                    color_idx = 2
                elif r > 180 and g < 80 and b < 80:      # Red
                    color_idx = 3
                else:                                    # White (Default)
                    color_idx = 1
                
                # Pack 2 bits into current byte position
                byte_val |= (color_idx << (6 - i * 2))
            bin_data.append(byte_val)
            
    return bytes(bin_data)

def make_data_packet(slot, chunk_num, is_last, chunk_data):
    """
    Builds the protocol data packet (dataType = 0x01)
    Format: 0xAA 0x01 <slot_lo> <slot_hi> <packetNum> <isLast> <len_lo> <len_hi> <payload bytes> 0xFF
    """
    packet = bytearray()
    packet.append(MAGIC_START)
    packet.append(OP_CHUNK_DATA)
    # slot as uint16 little-endian
    packet.extend(slot.to_bytes(2, byteorder='little'))
    packet.append(chunk_num)
    packet.append(1 if is_last else 0)
    # chunk length as uint16 little-endian
    packet.extend(len(chunk_data).to_bytes(2, byteorder='little'))
    packet.extend(chunk_data)
    packet.append(MAGIC_END)
    return bytes(packet)

def make_commit_packet(slot, raw_payload):
    """
    Builds the MD5 commit packet (dataType = 0x04)
    Format: 0xAA 0x04 <slot_lo> <slot_hi> <flag> <16 bytes MD5> 0xFF
    """
    # MD5 computed over the packed 30,000-byte raw image payload
    md5_hash = hashlib.md5(raw_payload).digest()
    
    packet = bytearray()
    packet.append(MAGIC_START)
    packet.append(OP_COMMIT_MD5)
    # slot as uint16 little-endian
    packet.extend(slot.to_bytes(2, byteorder='little'))
    packet.append(0x00) # Flag byte (always 0x00)
    packet.extend(md5_hash)
    packet.append(MAGIC_END)
    return bytes(packet)

async def scan_devices():
    """Scans for BLE advertising devices, highlighting anything matching Picpak."""
    print("🔍 Scanning for Bluetooth devices... (Ensure your Picpak light is SOLID)")
    devices = await BleakScanner.discover(timeout=5.0)
    
    found_any = False
    print("\nAvailable BLE Devices:")
    print("-" * 60)
    for d in devices:
        name = d.name or "Unknown"
        is_picpak = "picpak" in name.lower()
        star = "🌟 [Picpak Match]" if is_picpak else ""
        print(f"{d.address:<40} {name:<15} {star}")
        if is_picpak:
            found_any = True
            
    if not found_any:
        print("\n💡 Tip: If you don't see your device, hold the front button for 3 seconds until the light turns solid.")
    print("-" * 60)

async def upload_single_payload(client, char_uuid, slot, payload):
    """Helper that uploads a single 30,000 bytes packed payload using the correct BLE protocol."""
    total_bytes = len(payload)
    total_chunks = (total_bytes + CHUNK_SIZE_LIMIT - 1) // CHUNK_SIZE_LIMIT
    
    print(f"🚀 Streaming {total_chunks} wrapped packets to Slot {slot}...")
    
    for u in range(total_chunks):
        c = u * CHUNK_SIZE_LIMIT
        h = min(c + CHUNK_SIZE_LIMIT, total_bytes)
        is_last = (u == total_chunks - 1)
        chunk_data = payload[c:h]
        
        # Build protocol wrapped packet
        packet = make_data_packet(slot, u, is_last, chunk_data)
        
        # Write wrapped chunk to characteristic
        await client.write_gatt_char(char_uuid, packet, response=True)
        
        # Progress status printout
        percent = int(((u + 1) / total_chunks) * 100)
        bar = "█" * (percent // 5) + "░" * (20 - (percent // 5))
        sys.stdout.write(f"\rProgress: |{bar}| {percent}% ({u + 1}/{total_chunks} chunks)")
        sys.stdout.flush()
        
        # Small delay to prevent packet congestion
        await asyncio.sleep(0.04)
    
    print(f"\n📦 Sending MD5 checksum commit packet for Slot {slot}...")
    commit_packet = make_commit_packet(slot, payload)
    await client.write_gatt_char(char_uuid, commit_packet, response=True)

async def upload_image(address, image_path, char_uuid, slot_override=None):
    """Connects to Picpak and uploads a single image."""
    filename = os.path.basename(image_path)
    slot = slot_override if slot_override is not None else parse_slot_number(filename, 1)
    
    print(f"🖼️ Target: {filename} (Parsed Slot: {slot})")
    print(f"📦 Processing to 2-bit Spectra binary...")
    payload = png_to_picpak_bin(image_path)
    
    print(f"⚡ Connecting to BLE device '{address}'...")
    print(f"📝 Target write characteristic: {char_uuid}")
    
    try:
        async with BleakClient(address, timeout=20.0) as client:
            print(f"🔗 Connected: {client.is_connected}")
            await upload_single_payload(client, char_uuid, slot, payload)
            print("\n🎉 Image upload completed successfully!")
            print("⏳ Wait a few seconds for the Picpak screen to refresh...")
    except Exception as e:
        print(f"\n❌ Connection failed: {e}")
        print("👉 Verify your address, make sure the Picpak is close by and has pairing mode enabled.")

async def upload_folder(address, folder_path, char_uuid, delay=10.0, slot_override=None):
    """Finds all PNG files in folder and uploads them sequentially over a single BLE session."""
    png_pattern = os.path.join(folder_path, "*.png")
    png_files = sorted(glob.glob(png_pattern))
    
    if not png_files:
        print(f"❌ No PNG files found in directory: {folder_path}")
        return
        
    print(f"📁 Found {len(png_files)} PNG files to upload in sequence.")
    print(f"⚡ Connecting to BLE device '{address}'...")
    print(f"📝 Target write characteristic: {char_uuid}")
    
    try:
        async with BleakClient(address, timeout=25.0) as client:
            print(f"🔗 Connected: {client.is_connected}")
            
            for idx, img_path in enumerate(png_files):
                filename = os.path.basename(img_path)
                slot = slot_override if slot_override is not None else parse_slot_number(filename, idx + 1)
                
                print(f"\n────────────────────────────────────────────────────────────────────────────────")
                print(f"🖼️  [{idx + 1}/{len(png_files)}] Uploading: {filename} (Slot: {slot})")
                
                payload = png_to_picpak_bin(img_path)
                await upload_single_payload(client, char_uuid, slot, payload)
                
                print(f"✅ Uploaded successfully.")
                
                # Wait for device to process before sending next frame
                if idx < len(png_files) - 1:
                    print(f"⏳ Waiting {delay} seconds for device screen refresh...")
                    await asyncio.sleep(delay)
            
            print("\n🎉 All batch uploads completed successfully!")
    except Exception as e:
        print(f"\n❌ Connection lost or failed: {e}")

def main():
    parser = argparse.ArgumentParser(description="Picpak Direct BLE Uploader")
    parser.add_argument("--scan", action="store_true", help="Scan for nearby BLE devices")
    parser.add_argument("--address", type=str, help="Target Picpak BLE hardware address (MAC or UUID)")
    parser.add_argument("--image", type=str, help="Path to a single dithered 400x300 PNG image")
    parser.add_argument("--folder", type=str, help="Path to a folder of 400x300 PNG images to upload in batch")
    parser.add_argument("--slot", type=int, help="Override destination slot (1-500). If omitted, parsed from filename.")
    parser.add_argument("--char", type=str, default=DEFAULT_CHAR_UUID, help="GATT Write Characteristic UUID")
    parser.add_argument("--delay", type=float, default=10.0, help="Pause delay in seconds between batch uploads")
    
    args = parser.parse_args()
    
    if args.scan:
        asyncio.run(scan_devices())
    elif args.address and args.image:
        asyncio.run(upload_image(args.address, args.image, args.char, args.slot))
    elif args.address and args.folder:
        asyncio.run(upload_folder(args.address, args.folder, args.char, args.delay, args.slot))
    else:
        parser.print_help()

if __name__ == "__main__":
    main()