#!/usr/bin/env python3
import socket
import time
import urllib.request
import json
import base64
import os
from PIL import Image, ImageDraw

ESP32_IP = "192.168.68.123"
BROADCAST_IP = "192.168.68.255"
PORT = 8082
ARTIFACTS_DIR = "/Users/adolforeyna/.gemini/antigravity/brain/8b421e8b-10a8-4d06-a7d0-27bd82b42804"

def main():
    print(f"--- Starting UDP Broadcast Single-Frame Test ({BROADCAST_IP}) ---")
    
    # 1. Generate a test frame with Pillow displaying "BROADCAST SUCCESS!"
    img = Image.new("1", (400, 300), 0)
    draw = ImageDraw.Draw(img)
    draw.text((100, 140), "BROADCAST SUCCESS!", fill=1)
    draw.rectangle([40, 110, 360, 180], outline=1, width=2)
    
    # Map to Waveshare RLCD buffer format
    px = img.load()
    hw_buffer = bytearray(15000)
    for y in range(300):
        for x in range(400):
            if px[x, y]:
                inv_y = 299 - y
                bx = x // 2
                by = inv_y // 4
                idx = bx * 75 + by
                lx, ly = x % 2, inv_y % 4
                bit = 7 - (ly * 2 + lx)
                hw_buffer[idx] |= (1 << bit)
                
    # 2. Open UDP socket with broadcast permission
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    
    # Send all 15 chunks representing the single frame
    print("Broadcasting chunks...")
    frame_id = 99
    for chunk_idx in range(15):
        packet = bytearray(1002)
        packet[0] = frame_id
        packet[1] = chunk_idx
        start = chunk_idx * 1000
        packet[2:1002] = hw_buffer[start : start + 1000]
        sock.sendto(packet, (BROADCAST_IP, PORT))
        time.sleep(0.002) # Small pacing delay
        
    sock.close()
    print("Frame broadcast sent.")
    
    # 3. Wait 1 second for ESP32 to render and display
    time.sleep(1.0)
    
    # 4. Capture screenshot from ESP32
    print(f"Requesting screenshot from ESP32 ({ESP32_IP}) to verify broadcast reception...")
    payload = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "tools/call",
        "params": {
            "name": "get_screenshot",
            "arguments": {}
        }
    }
    
    req_data = json.dumps(payload).encode('utf-8')
    req = urllib.request.Request(
        f"http://{ESP32_IP}/api/mcp",
        data=req_data,
        headers={"Content-Type": "application/json"},
        method="POST"
    )
    
    try:
        with urllib.request.urlopen(req, timeout=10.0) as response:
            resp_body = response.read().decode('utf-8')
            resp_json = json.loads(resp_body)
            content = resp_json.get("result", {}).get("content", [{}])[0].get("text", "")
            
            if content.startswith("__PBM_BASE64__:"):
                pbm_b64 = content.split(":", 1)[1]
                pbm_bytes = base64.b64decode(pbm_b64)
                
                temp_pbm = "temp_capture_bcast.pbm"
                with open(temp_pbm, "wb") as pf:
                    pf.write(pbm_bytes)
                
                # Convert to PNG using Pillow
                img = Image.open(temp_pbm)
                dest_path = os.path.join(ARTIFACTS_DIR, "udp_broadcast_success.png")
                img.save(dest_path)
                os.remove(temp_pbm)
                print(f"Success: Broadcast screenshot saved to {dest_path}")
            else:
                print("Error: Invalid screenshot response format.")
    except Exception as e:
        print(f"Failed to capture screenshot: {e}")

if __name__ == "__main__":
    main()