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

ESP32_IP = "192.168.68.123"
BROADCAST_IP = "192.168.68.255"
URL = f"http://{ESP32_IP}/api/mcp"
ARTIFACTS_DIR = "/Users/adolforeyna/.gemini/antigravity/brain/8b421e8b-10a8-4d06-a7d0-27bd82b42804"

def capture_screenshot(output_filename):
    print(f"Requesting screenshot from ESP32 ({ESP32_IP})...")
    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(
        URL,
        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, output_filename)
                img.save(dest_path)
                os.remove(temp_pbm)
                print(f"Success: Screenshot saved to {dest_path}")
                return True
            else:
                print("Error: Invalid screenshot response format.")
                return False
    except Exception as e:
        print(f"Failed to capture screenshot: {e}")
        return False

def main():
    print(f"--- Starting UDP Broadcast Stream Verification ({BROADCAST_IP}) ---")
    
    # 1. Start test_stream.py targeting broadcast IP
    cmd = ["python3", "test_stream.py", "--ip", BROADCAST_IP, "--protocol", "udp", "--source", "animation"]
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    print(f"Broadcast stream started in background (PID: {proc.pid}).")
    
    # 2. Wait for packets to propagate
    print("Broadcasting UDP packets... Waiting 3 seconds...")
    time.sleep(3.0)
    
    # 3. Capture screen rendering on ESP32
    capture_screenshot("udp_broadcast_success.png")
    
    # 4. Terminate stream client
    print(f"Terminating broadcast client process...")
    proc.terminate()
    proc.wait()
    print(f"Broadcast stream client stopped.")
    
    # Let the screen return to dashboard
    print("Waiting 4 seconds to let the screen return to dashboard...")
    time.sleep(4.0)
    print("Completed.")

if __name__ == "__main__":
    main()