import asyncio
import numpy as np
from ai import AIAccelerator

async def main():
    print("\n--- Testing AIAccelerator with WebSocket Storage ---")
    try:
        accelerator = AIAccelerator()
        await accelerator.connect_to_storage()
        print("AIAccelerator initialized and connected successfully.")

        # Test model upload
        dummy_model_info = {"layers": 5, "neurons": 100, "type": "CNN"}
        model_name = "test_cnn_model"
        print(f"Attempting to store model: {model_name}")
        if await accelerator.store_model_state(model_name, dummy_model_info):
            print(f"Model \'{model_name}\' stored successfully.")
        else:
            print(f"Failed to store model \'{model_name}\'")

        # Test tensor core initialization (requires VRAM connection)
        print("Attempting to initialize tensor cores...")
        if accelerator.initialize_tensor_cores():
            print("Tensor cores initialized successfully.")
        else:
            print("Failed to initialize tensor cores.")

    except Exception as e:
        print(f"An error occurred during AIAccelerator testing: {e}")

if __name__ == "__main__":
    asyncio.run(main())