#!/usr/bin/env python3
"""
Upload dtu_provision.py to ESP32-S3 via webrepl_cli, then execute it and stream output.
"""
import asyncio
import websockets
import time
import subprocess
import sys

WEBREPL_URI = "ws://192.168.40.90:8266"
PASSWORD = "claw2287"
WEBREPL_CLI = "/tmp/webrepl_cli.py"
LOCAL_SCRIPT = "/home/crogers2287/.openclaw/agents/clinker/dtu_provision.py"
REMOTE_SCRIPT = "/dtu_provision.py"

def upload_file():
    print(f"Uploading {LOCAL_SCRIPT} -> {REMOTE_SCRIPT}")
    result = subprocess.run(
        ["python3", WEBREPL_CLI, "-p", "claw2287", LOCAL_SCRIPT, f"192.168.40.90:{REMOTE_SCRIPT}"],
        capture_output=True, text=True, timeout=60
    )
    print("STDOUT:", result.stdout)
    print("STDERR:", result.stderr)
    return result.returncode == 0

async def run_and_capture():
    async with websockets.connect(WEBREPL_URI, ping_interval=None, ping_timeout=None) as ws:
        # Banner
        try:
            banner = await asyncio.wait_for(ws.recv(), timeout=5)
            print(f"[banner] {banner!r}")
        except:
            print("[no banner]")
        
        # Auth
        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!r}")
        except:
            print("[auth timeout]")
        
        # Interrupt any running code
        for _ in range(3):
            await ws.send(b'\x03')
            await asyncio.sleep(0.3)
            try:
                await asyncio.wait_for(ws.recv(), timeout=0.5)
            except:
                pass
        
        # Clear line
        await ws.send(b'\r\n')
        await asyncio.sleep(0.3)
        try:
            await asyncio.wait_for(ws.recv(), timeout=1)
        except:
            pass
        
        # Execute script
        cmd = f"exec(open('{REMOTE_SCRIPT}').read())\r\n"
        print(f"[exec] {cmd.strip()}")
        await ws.send(cmd.encode())
        
        # Stream output - the script takes ~120s max
        print("\n--- ESP32-S3 OUTPUT ---")
        deadline = time.time() + 150
        output_chunks = []
        last_output = time.time()
        
        while time.time() < deadline:
            try:
                data = await asyncio.wait_for(ws.recv(), timeout=10)
                if isinstance(data, bytes) and data:
                    text = data.decode('utf-8', errors='replace')
                    print(text, end='', flush=True)
                    output_chunks.append(text)
                    last_output = time.time()
            except asyncio.TimeoutError:
                combined = ''.join(output_chunks)
                if 'DONE' in combined:
                    print("\n[Script completed - DONE marker found]")
                    break
                # If no output for 30s after script started producing output
                if output_chunks and (time.time() - last_output > 30):
                    print("\n[No output for 30s, assuming done]")
                    break
        
        return ''.join(output_chunks)

def read_result_file():
    """Read /result.txt from ESP32 using webrepl_cli"""
    import tempfile, os
    tmpfile = "/tmp/dtu_result.txt"
    result = subprocess.run(
        ["python3", WEBREPL_CLI, "-p", "claw2287", f"192.168.40.90:/result.txt", tmpfile],
        capture_output=True, text=True, timeout=30
    )
    if result.returncode == 0 and os.path.exists(tmpfile):
        with open(tmpfile) as f:
            return f.read()
    return None

if __name__ == '__main__':
    print("=== Step 1: Upload provision script ===")
    if not upload_file():
        print("Upload failed! Trying anyway...")
    
    print("\n=== Step 2: Execute provision script ===")
    output = asyncio.run(run_and_capture())
    
    print("\n=== Step 3: Read /result.txt from ESP32 ===")
    time.sleep(2)
    result = read_result_file()
    if result:
        print("\n--- /result.txt contents ---")
        print(result)
    else:
        print("Could not read /result.txt")
    
    print("\n=== Done ===")