#!/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"
URL = f"http://{ESP32_IP}/api/mcp"
ARTIFACTS_DIR = "/Users/adolforeyna/.gemini/antigravity/brain/da0e7ff3-d127-4865-bfea-f8ae9198bd40"

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.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 run_stream_test(protocol):
    print(f"\n--- Starting {protocol.upper()} Stream Verification ---")
    
    # 1. Start test_stream.py in background
    cmd = ["python3", "test_stream.py", "--ip", ESP32_IP, "--protocol", protocol, "--source", "animation"]
    proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    print(f"Stream client process started in background (PID: {proc.pid}).")
    
    # 2. Wait for frames to stream and render
    print("Streaming frames... Waiting 3 seconds...")
    time.sleep(3.0)
    
    # 3. Capture screen rendering
    filename = f"{protocol}_stream_success.png"
    capture_screenshot(filename)
    
    # 4. Terminate stream client
    print(f"Terminating stream client process...")
    proc.terminate()
    proc.wait()
    print(f"Stream client stopped.")

def main():
    if not os.path.exists(ARTIFACTS_DIR):
        os.makedirs(ARTIFACTS_DIR)
        
    # Test TCP
    run_stream_test("tcp")
    
    # Test Dashboard Auto-Timeout (stream timeout is 3.0s, let's wait 4.0s)
    print("\n--- Verifying Stream Inactivity Timeout ---")
    print("Waiting 4 seconds for stream to time out on ESP32...")
    time.sleep(4.0)
    capture_screenshot("dashboard_resumed.png")
    
    # Test UDP
    run_stream_test("udp")
    
    # Clean up stream
    print("\nWait 4 seconds to let the screen return to dashboard...")
    time.sleep(4.0)
    
    print("\nVerification Completed successfully.")

if __name__ == "__main__":
    main()