#!/usr/bin/env python3
"""Run a script on ESP32-S3 via WebREPL and capture output."""
import asyncio
import websockets
import time
import sys

WEBREPL_URI = "ws://192.168.40.90:8266"
PASSWORD = "claw2287"

async def run_script(script_name, timeout=120):
    uri = WEBREPL_URI
    async with websockets.connect(uri, ping_interval=None, ping_timeout=None) as ws:
        # Read banner
        try:
            banner = await asyncio.wait_for(ws.recv(), timeout=5)
            print(f"[banner]: {banner}", flush=True)
        except asyncio.TimeoutError:
            print("[no banner]", flush=True)
        
        # Send password
        await ws.send(b'claw2287\r\n')
        await asyncio.sleep(0.8)
        try:
            resp = await asyncio.wait_for(ws.recv(), timeout=5)
            print(f"[auth]: {resp}", flush=True)
        except asyncio.TimeoutError:
            print("[auth timeout]", flush=True)
        
        # Ctrl-C to interrupt any running code
        await ws.send(b'\x03')
        await asyncio.sleep(0.5)
        try:
            await asyncio.wait_for(ws.recv(), timeout=1)
        except:
            pass
        
        # Ctrl-C again
        await ws.send(b'\x03')
        await asyncio.sleep(0.5)
        try:
            await asyncio.wait_for(ws.recv(), timeout=1)
        except:
            pass
        
        # Run the script
        cmd = f"exec(open('{script_name}').read())\r\n"
        print(f"[sending]: {cmd.strip()}", flush=True)
        await ws.send(cmd.encode())
        
        # Collect output
        deadline = time.time() + timeout
        output = []
        while time.time() < deadline:
            try:
                data = await asyncio.wait_for(ws.recv(), timeout=5)
                if isinstance(data, bytes):
                    text = data.decode('utf-8', errors='replace')
                    print(text, end='', flush=True)
                    output.append(text)
            except asyncio.TimeoutError:
                combined = ''.join(output)
                if 'DONE' in combined or 'ERROR' in combined or 'FAILED' in combined:
                    print("\n[got terminal marker, stopping]", flush=True)
                    break
        
        return ''.join(output)

if __name__ == '__main__':
    script = sys.argv[1] if len(sys.argv) > 1 else '/dtu_hm_provision.py'
    asyncio.run(run_script(script))