Graphics (16-bit)

import struct
from PIL import Image
import re
import argparse

# Wicked palette (default)
WICKED_PALETTE = [
    (0, 0, 0),      # 0: #000000
    (51, 119, 119), # 1: #337777
    (119, 153, 119),# 2: #779977
    (119, 187, 153),# 3: #77BB99
    (255, 0, 0),    # 4: #FF0000
    (187, 0, 0),    # 5: #BB0000
    (255, 221, 0),  # 6: #FFDD00
    (221, 153, 0),  # 7: #DD9900
    (255, 187, 0),  # 8: #FFBB00
    (0, 0, 0),      # 9: #000000
    (0, 0, 85),     # 10: #000055
    (85, 85, 119),  # 11: #555577
    (153, 153, 153),# 12: #999999
    (119, 153, 187),# 13: #7799BB
    (187, 187, 255),# 14: #BBBBFF
    (85, 119, 153)  # 15: #557799
]

# UI palette for cursors (indices 88–93, if needed)
WICKED_PALETTE_UI = [
    (0, 0, 0),      # 0: Black #000000
    (255, 255, 255),# 1: White #FFFFFF
    (255, 255, 255),# 2: White #FFFFFF
    (255, 255, 255),# 3: White #FFFFFF
    (255, 255, 255),# 4: White #FFFFFF
    (0, 0, 0),      # 5: Black #000000
    (0, 0, 0),      # 6: Black #000000
    (0, 0, 0),      # 7: Black #000000
    (0, 0, 0),      # 8: Black #000000
    (0, 0, 0),      # 9: Black #000000
    (0, 0, 0),      # 10: Black #000055
    (0, 0, 0),      # 11: Black #555577
    (0, 0, 0),      # 12: Black #999999
    (0, 0, 0),      # 13: Black #7799BB
    (0, 0, 0),      # 14: Black #BBBBFF
    (0, 0, 0)       # 15: Black #557799
]

def parse_assembler_data(data_text):
    """Parse Amiga Assembler data (DC.L, DC.W, DC.B, DS.L) into byte arrays."""
    indices = {}
    current_index = None
    current_data = bytearray()
    width, height, mode = None, None, None
    lines = data_text.splitlines()
    line_idx = 0
    
    while line_idx < len(lines):
        line = lines[line_idx]
        orig_line = line
        line = line.split(';', 1)[0].strip()
        if not line:
            line_idx += 1
            continue
            
        # Check for any label ending with a colon, excluding assembler directives
        if line.endswith(':') and not any(line.startswith(d) for d in ['DC.L', 'DC.W', 'DC.B', 'DS.L']):
            if current_index is not None:
                indices[current_index] = {
                    'data': bytes(current_data),
                    'width_minus_1': width,
                    'height_minus_1': height,
                    'mode': mode
                }
            current_index = line[:-1]  # Remove the colon
            current_data = bytearray()
            # Look for metadata in the next line
            if line_idx + 1 < len(lines):
                next_line = lines[line_idx + 1]
                match = re.search(r'(Width|Depth|W)=([0-9a-fA-F]+)[\s,;:]*(Height|Depth|H)=([0-9a-fA-F]+)[\s,;:]*(Mode|M)=([0-9a-fA-F]+)', next_line, re.IGNORECASE)
                if match:
                    width = int(match.group(2), 16)
                    height = int(match.group(4), 16)
                    mode = int(match.group(6), 16)
                    if mode == 0xFFFF or mode == -1:
                        mode = -1
                    print(f"Parsed metadata for {current_index}: Width={width:04x}, Height={height:04x}, Mode={mode:04x}")
                else:
                    print(f"Warning: Could not parse metadata for {current_index}: {orig_line}")
                    print(f"Next line: {next_line}")
                    width, height, mode = None, None, None
            else:
                print(f"Warning: Could not parse metadata for {current_index}: {orig_line}")
                print("No next line available")
                width, height, mode = None, None, None
        elif 'DC.L' in line:
            values_part = line.split('DC.L')[1].strip()
            if values_part:
                values = [v.strip() for v in values_part.split(',') if v.strip()]
                for v in values:
                    try:
                        val = int(v.replace('$', '0x'), 16)
                        if mode == 6:
                            # Take first 4 bytes for mode 6, skip duplicates
                            current_data.extend(struct.pack('>I', val)[:4])
                        else:
                            current_data.extend(struct.pack('>I', val))
                    except ValueError as e:
                        print(f"Warning: Skipping invalid DC.L value '{v}' in line: {line}")
        elif 'DC.W' in line:
            values_part = line.split('DC.W')[1].strip()
            if values_part:
                values = [v.strip() for v in values_part.split(',') if v.strip()]
                for v in values:
                    try:
                        val = int(v.replace('$', '0x'), 16)
                        current_data.extend(struct.pack('>H', val))
                    except ValueError as e:
                        print(f"Warning: Skipping invalid DC.W value '{v}' in line: {line}")
        elif 'DC.B' in line:
            values_part = line.split('DC.B')[1].strip()
            if values_part:
                values = [v.strip() for v in values_part.split(',') if v.strip()]
                for v in values:
                    try:
                        val = int(v.replace('$', '0x'), 16)
                        current_data.append(val)
                    except ValueError as e:
                        print(f"Warning: Skipping invalid DC.B value '{v}' in line: {line}")
        elif 'DS.L' in line and mode != 6:  # Skip DS.L for mode 6
            values_part = line.split('DS.L')[1].strip()
            if values_part:
                try:
                    count = int(values_part.replace('$', '0x'), 16)
                    current_data.extend([0] * count * 4)
                except ValueError as e:
                    print(f"Warning: Skipping invalid DS.L count '{values_part}' in line: {line}")
        line_idx += 1
    
    if current_index is not None:
        indices[current_index] = {
            'data': bytes(current_data),
            'width_minus_1': width,
            'height_minus_1': height,
            'mode': mode
        }
    
    return indices

def decode_wicked_bitmap(mode, width_minus_1, height_minus_1, data):
    """Decode bitmap for any mode."""
    cols = width_minus_1 + 1
    rows = height_minus_1 + 1
    U = cols * rows
    cell_size_px = 16 if mode in [0, 2, 3, 5, -1] else 8
    expected_size = 8 * U if mode in [0, 2, 3, 5] else 4 * U
    
    if mode in [3, 5]:
        by_plane = bytearray()
        pos = 0
        if len(data) < 2:
            print(f"Warning: Data too short for mode {mode}: {len(data)} bytes")
            return bytearray(expected_size * 2)
        ctrl_hi = data[0]
        ctrl2 = data[1]
        pos += 2
        
        while pos < len(data):
            b = data[pos]
            pos += 1
            if (b & 0xF0) == (ctrl_hi & 0xFF):
                c = data[pos] if pos < len(data) else 0
                pos += 1
                n = b & 0x0F
                if c == 0:
                    by_plane.append(b)
                elif n == 0:
                    by_plane.extend([0] * c)
                elif n == 1:
                    break
                else:
                    by_plane.extend([c] * n)
            else:
                by_plane.append(b)
        
        if len(by_plane) != expected_size * 2:
            print(f"Warning: Decoded size {len(by_plane)} bytes, expected {expected_size * 2}")
            if len(by_plane) < expected_size * 2:
                by_plane.extend([0] * (expected_size * 2 - len(by_plane)))
            else:
                by_plane = by_plane[:expected_size * 2]
        
        interleaved = bytearray(len(by_plane))
        if ctrl2 != 0:
            plane_size = 2 * U
            for i in range(U):
                base = i * 2
                p0 = by_plane[base:base+2]
                p1 = by_plane[plane_size + base:plane_size + base+2]
                p2 = by_plane[2*plane_size + base:2*plane_size + base+2]
                p3 = by_plane[3*plane_size + base:3*plane_size + base+2]
                interleaved[i*8:i*8+8] = p0 + p1 + p2 + p3
        else:
            interleaved = by_plane
        return interleaved
    
    elif mode in [0, 2]:
        if len(data) < expected_size * 2:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size * 2}")
            data = data + bytearray(expected_size * 2 - len(data))
        elif len(data) > expected_size * 2:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size * 2}")
            data = data[:expected_size * 2]
        return data
    
    elif mode in [1, 6]:
        if len(data) < expected_size:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size}")
            data = data + bytearray(expected_size - len(data))
        elif len(data) > expected_size:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size}")
            data = data[:expected_size]
        if mode == 6:
            print(f"Mode 6 data for index (truncated): {list(data[:56])}")  # Debug first 56 bytes
        return data
    
    elif mode == -1:
        expected_size = 2 * U * 4
        if len(data) < expected_size // 4:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size // 4}")
            data = data + bytearray((expected_size // 4) - len(data))
        elif len(data) > expected_size // 4:
            print(f"Warning: Data size {len(data)} bytes, expected {expected_size // 4}")
            data = data[:expected_size // 4]
        return data
    
    else:
        raise ValueError(f"Unsupported mode: {mode}")

def simulate_masked_render(interleaved, mode, cols, rows, bitplanes):
    """Simulate masked rendering for modes 0, 1, 2, 3, 5, 6."""
    cell_bytes = 8 if mode in [0, 2, 3, 5] else 4
    cell_px = 16 if mode in [0, 2, 3, 5] else 8
    stride = 40
    
    if mode == 6:
        print(f"Mode 6 bitplane input: {list(interleaved[:56])}")  # Debug input
    for row in range(rows):
        line_offset = row * stride
        for col in range(cols):
            cell_data = interleaved[(row * cols + col) * cell_bytes : (row * cols + col + 1) * cell_bytes]
            mask = 0
            for i in range(4):
                try:
                    val = struct.unpack('>H', cell_data[i*2:i*2+2])[0] if mode in [0, 2, 3, 5] else cell_data[i]
                    mask |= val
                except (struct.error, IndexError):
                    print(f"Warning: Failed to unpack cell data at row {row}, col {col}, mode {mode}")
                    val = 0
            mask = ~mask & ((1 << cell_px) - 1)
            
            for p in range(4):
                plane_offset = line_offset + col * (cell_px // 8)
                try:
                    if mode in [0, 2, 3, 5]:
                        dest_val = struct.unpack('>H', bitplanes[p][plane_offset:plane_offset+2])[0]
                        dest_val &= mask
                        new_val = struct.unpack('>H', cell_data[p*2:p*2+2])[0]
                        dest_val |= new_val
                        bitplanes[p][plane_offset:plane_offset+2] = struct.pack('>H', dest_val)
                    else:
                        dest_val = bitplanes[p][plane_offset]
                        dest_val &= mask & 0xFF
                        new_val = cell_data[p]
                        dest_val |= new_val
                        bitplanes[p][plane_offset] = dest_val
                except (struct.error, IndexError):
                    print(f"Warning: Failed to unpack bitplane data at plane {p}, offset {plane_offset}")
                    continue
    return bitplanes

def simulate_direct_render(data, cols, rows, bitplanes):
    """Simulate direct rendering for mode -1."""
    stride = 40
    U = cols * rows
    for p in range(4):
        plane_data = data[p * (U * 2) : (p + 1) * (U * 2)]
        for row in range(rows):
            offset = row * stride
            try:
                word = struct.unpack('>H', plane_data[row * 2 : row * 2 + 2])[0] if len(plane_data[row * 2 : row * 2 + 2]) == 2 else 0
                bitplanes[p][offset:offset+2] = struct.pack('>H', word)
            except struct.error:
                print(f"Warning: Failed to unpack plane {p}, row {row}, expected 2 bytes, got {len(plane_data[row * 2 : row * 2 + 2])}")
                bitplanes[p][offset:offset+2] = struct.pack('>H', 0)
    return bitplanes

def pixels_from_bitplanes(bitplanes, rows, width_px):
    """Extract 2D pixels from 4 bitplanes."""
    stride = 40
    pixels = []
    for row in range(rows):
        row_pixels = []
        line_offset = row * stride
        for px_byte in range(width_px // 8):
            offset = line_offset + px_byte
            try:
                p0_byte = bitplanes[0][offset]
                p1_byte = bitplanes[1][offset]
                p2_byte = bitplanes[2][offset]
                p3_byte = bitplanes[3][offset]
            except IndexError:
                print(f"Warning: Bitplane index out of range at offset {offset}")
                p0_byte = p1_byte = p2_byte = p3_byte = 0
            for bit in range(7, -1, -1):
                index = ((p0_byte >> bit) & 1) + 2*((p1_byte >> bit) & 1) + 4*((p2_byte >> bit) & 1) + 8*((p3_byte >> bit) & 1)
                row_pixels.append(index)
        pixels.append(row_pixels)
        if row < 14:
            print(f"Pixel row {row}: {row_pixels[:8]}")
    return pixels

def save_bitmap_to_png(pixels, filename, palette, index_name, scale_factor=1):
    """Save 2D pixel array as PNG, scaled by scale_factor."""
    height = len(pixels)
    width = len(pixels[0]) if pixels else 1
    img = Image.new('P', (width, height))
    # Use UI palette for cursor indices if needed (adjust if specific names are cursors)
    selected_palette = WICKED_PALETTE_UI if index_name in ['88', '89', '90', '91', '92', '93'] else WICKED_PALETTE
    img.putpalette([c for rgb in selected_palette for c in rgb])
    for y in range(height):
        for x in range(width):
            img.putpixel((x, y), pixels[y][x])
    # Scale image by scale_factor
    scaled_img = img.resize((width * scale_factor, height * scale_factor), Image.NEAREST)
    scaled_img.save(filename, 'PNG')
    return width * scale_factor, height * scale_factor  # Return scaled dimensions

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Render Amiga graphics to PNG")
parser.add_argument('--scale', type=int, choices=[1, 2], default=1, help="Scale factor for output PNG (1 for normal size, 2 for double size)")
args = parser.parse_args()

# Parse assembler data
try:
    with open('graphicsdata.txt', 'r') as f:
        assembler_data = f.read()
except FileNotFoundError:
    print("Error: 'graphicsdata.txt' not found in current directory")
    exit(1)

indices = parse_assembler_data(assembler_data)

# Process each index
for index_name, info in indices.items():
    mode = info['mode']
    width_minus_1 = info['width_minus_1']
    height_minus_1 = info['height_minus_1']
    data = info['data']
    
    if None in (mode, width_minus_1, height_minus_1):
        print(f"Skipping Index {index_name}: Incomplete metadata (mode={mode}, width_minus_1={width_minus_1}, height_minus_1={height_minus_1})")
        continue
    
    cols = width_minus_1 + 1
    rows = height_minus_1 + 1
    width_px = cols * (16 if mode in [0, 2, 3, 5, -1] else 8)
    
    try:
        # Decode
        decoded = decode_wicked_bitmap(mode, width_minus_1, height_minus_1, data)
        
        # Initialize bitplanes
        bitplanes = [bytearray(40 * rows) for _ in range(4)]
        
        # Render
        if mode in [0, 1, 2, 3, 5, 6]:
            bitplanes = simulate_masked_render(decoded, mode, cols, rows, bitplanes)
        else:
            bitplanes = simulate_direct_render(decoded, cols, rows, bitplanes)
        
        # Extract pixels
        pixels = pixels_from_bitplanes(bitplanes, rows, width_px)
        
        # Save PNG
        filename = f"{index_name}_{args.scale}x.png"
        scaled_width, scaled_height = save_bitmap_to_png(pixels, filename, WICKED_PALETTE, index_name, scale_factor=args.scale)
        print(f"Saved {filename} ({scaled_width}x{scaled_height} pixels, mode {mode}, scaled {args.scale}x)")
        
        # Print all rows for mode 6
        if mode == 6:
            for i, row in enumerate(pixels[:14]):
                print(f"Index {index_name}, Row {i}: {row[:8]}")
            
    except ValueError as e:
        print(f"Error processing Index {index_name}: {e}")
    except Exception as e:
        print(f"Unexpected error in Index {index_name}: {e}")