import asyncio
import websockets
import subprocess
import json

BANNER = """
╔══════════════════════════════════════════╗
║         OnboardAI Local Agent            ║
║  Listening on ws://localhost:7788        ║
║  Waiting for OnboardAI to connect...     ║
╚══════════════════════════════════════════╝
"""

async def handle(websocket):
    print("✓ OnboardAI connected\n")
    try:
        async for raw in websocket:
            data = json.loads(raw)

            if data["type"] == "ping":
                await websocket.send(json.dumps({"type": "pong"}))

            elif data["type"] == "run":
                cmd = data["command"]
                cmd_id = data["id"]

                print(f"\n{'─'*44}")
                print(f"  Command : {cmd}")
                print(f"{'─'*44}")

                proc = await asyncio.create_subprocess_shell(
                    cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    cwd=data.get("cwd")
                )
                stdout, stderr = await proc.communicate()
                output = stdout.decode() or stderr.decode()

                if output:
                    print(output)

                status = "✓ Done" if proc.returncode == 0 else "✗ Failed"
                print(status)

                await websocket.send(json.dumps({
                    "id":      cmd_id,
                    "success": proc.returncode == 0,
                    "output":  output
                }))

    except websockets.exceptions.ConnectionClosed:
        print("\n⚠ OnboardAI disconnected. Waiting for reconnect...\n")

async def start_server():
    print(BANNER)
    async with websockets.serve(handle, "localhost", 7788):
        await asyncio.Future()

def main():
    try:
        asyncio.run(start_server())
    except KeyboardInterrupt:
        print("\nAgent stopped.")

if __name__ == "__main__":
    main()