import asyncio
import websockets
import json
import numpy as np

async def test_websocket():
    uri = "wss://8765-int35aifjdyxr46oxipjg-cf14c883.manusvm.computer"
    async with websockets.connect(uri) as websocket:
        print("Connected to websocket.")

        # Test VRAM write operation
        vram_data = np.random.rand(10, 10).tolist()
        vram_write_message = {
            "operation": "vram",
            "type": "write",
            "block_id": "test_block_1",
            "data": vram_data
        }
        await websocket.send(json.dumps(vram_write_message))
        response = await websocket.recv()
        print(f"VRAM write response: {response}")

        # Test VRAM read operation
        vram_read_message = {
            "operation": "vram",
            "type": "read",
            "block_id": "test_block_1"
        }
        await websocket.send(json.dumps(vram_read_message))
        response = await websocket.recv()
        print(f"VRAM read response: {response}")

        # Test state save operation
        state_data = {"gpu_temp": 65, "fan_speed": 2000}
        state_save_message = {
            "operation": "state",
            "type": "save",
            "component": "gpu_core",
            "state_id": "current_state",
            "data": state_data
        }
        await websocket.send(json.dumps(state_save_message))
        response = await websocket.recv()
        print(f"State save response: {response}")

        # Test state load operation
        state_load_message = {
            "operation": "state",
            "type": "load",
            "component": "gpu_core",
            "state_id": "current_state"
        }
        await websocket.send(json.dumps(state_load_message))
        response = await websocket.recv()
        print(f"State load response: {response}")

        # Test cache set operation
        cache_data = {"model_version": "v1.2", "accuracy": 0.98}
        cache_set_message = {
            "operation": "cache",
            "type": "set",
            "key": "model_info",
            "data": cache_data
        }
        await websocket.send(json.dumps(cache_set_message))
        response = await websocket.recv()
        print(f"Cache set response: {response}")

        # Test cache get operation
        cache_get_message = {
            "operation": "cache",
            "type": "get",
            "key": "model_info"
        }
        await websocket.send(json.dumps(cache_get_message))
        response = await websocket.recv()
        print(f"Cache get response: {response}")

asyncio.run(test_websocket())