import struct
import zlib
from PIL import Image, ImageDraw, ImageFont
import io
import os

# Create AeThexOS boot logo
print("Creating AeThexOS boot logo...")
img = Image.new('RGBA', (800, 1280), color=(5, 5, 15, 255))
draw = ImageDraw.Draw(img)

# Try to use a nice font
try:
    font = ImageFont.truetype("arial.ttf", 72)
    small_font = ImageFont.truetype("arial.ttf", 24)
except:
    font = ImageFont.load_default()
    small_font = font

# Draw AeThexOS text centered
text = "AeThexOS"
bbox = draw.textbbox((0, 0), text, font=font)
text_width = bbox[2] - bbox[0]
text_height = bbox[3] - bbox[1]
x = (800 - text_width) // 2
y = (1280 - text_height) // 2 - 50

# Glow effect
for offset in range(10, 0, -2):
    alpha = int(255 * (1 - offset/10))
    glow_color = (0, 200, 255, alpha)
    draw.text((x-offset//2, y-offset//2), text, fill=glow_color, font=font)
    draw.text((x+offset//2, y+offset//2), text, fill=glow_color, font=font)

# Main text
draw.text((x, y), text, fill=(0, 220, 255, 255), font=font)

# Subtitle
sub_text = "Starting..."
sub_bbox = draw.textbbox((0, 0), sub_text, font=small_font)
sub_x = (800 - (sub_bbox[2] - sub_bbox[0])) // 2
draw.text((sub_x, y + 100), sub_text, fill=(100, 100, 120, 255), font=small_font)

img.save('aethex_logo.png')
print("Saved aethex_logo.png")

# Convert to BGRA raw bytes
raw_bgra = img.tobytes('raw', 'BGRA')
print(f"Raw BGRA size: {len(raw_bgra)} bytes")

# Compress with zlib
compressed = zlib.compress(raw_bgra, 9)
print(f"Compressed size: {len(compressed)} bytes")

# Now patch the logo.bin
print("\nPatching logo.bin...")
data = bytearray(open('logo.bin', 'rb').read())

# The first logo is at offset 688 (zlib compressed)
# Find the compressed data and its size
original_compressed_start = 688

# Get original block header info from around offset 512
# We need to find where the size is stored

# Let's look at the header structure more carefully
# At offset 512, there seems to be block info
print(f"Bytes at 512: {data[512:560].hex()}")

# Parse potential block header
# Format might be: offset(4) + compressed_size(4) + width(4) + height(4) + ...
block_offset = struct.unpack('<I', data[512:516])[0]
block_size = struct.unpack('<I', data[516:520])[0]
print(f"Block 1: offset={block_offset}, size={block_size}")

# Replace the compressed data if sizes are similar
original_size = struct.unpack('<I', data[516:520])[0]
print(f"Original compressed size at offset 688: ~{original_size}")
print(f"New compressed size: {len(compressed)}")

if len(compressed) <= original_size + 50000:  # Allow some margin
    # Pad new data to match original size
    new_compressed = compressed + b'\x00' * (original_size - len(compressed))
    
    # Replace in data
    data[688:688+original_size] = new_compressed
    
    # Update size in header if needed
    struct.pack_into('<I', data, 516, len(compressed))
    
    # Save patched logo
    open('logo_patched.bin', 'wb').write(data)
    print("Saved logo_patched.bin")
else:
    print(f"New data too large! Need to reconstruct logo partition.")
    print("Creating minimal logo partition...")
    
    # Create a new logo partition with just our image
    # This is risky but might work