import asyncio
import time
import numpy as np
import os

from virtual_gpu.vram import VRAM
from virtual_gpu.virtual_ram import VirtualRAM
from virtual_gpu.vgpu import VirtualGPU
from virtual_gpu.driver import GPUDriver
from virtual_gpu.render import Renderer
from virtual_gpu.ai import AIAccelerator
from virtual_gpu.shader import ShaderManager
from virtual_gpu.display import DisplayManager, DisplayMode

async def test_virtual_ram_transfer():
    print("Starting Virtual RAM to VRAM Transfer Test...")
    print("=" * 50)

    # 1. Initialize VirtualRAM(128) and VirtualVRAM(500)
    print("Initializing VirtualRAM and VRAM...")
    ram = VirtualRAM(capacity_gb=128)
    vram = VRAM(memory_size_gb=500)
    print("✓ VirtualRAM (128GB) and VRAM (500GB) initialized.")

    # Initialize other vGPU components for a complete environment
    renderer = Renderer(vram)
    ai_accelerator = AIAccelerator(vram)
    vgpu = VirtualGPU(num_sms=800, total_cores=50000)
    vgpu.set_modules(vram, renderer, ai_accelerator, None)
    driver = GPUDriver(vgpu)
    vgpu.driver = driver
    shader_manager = ShaderManager()
    display_manager = DisplayManager(vram)
    file_display = display_manager.add_display("file", DisplayMode.FILE, output_dir="./test_output_ram_transfer")
    print("✓ Other vGPU components initialized.")

    print("\n" + "=" * 50)

    # 2. Allocate a 1MB tensor in RAM.
    tensor_size_mb = 1 # Use 1MB for actual data allocation
    tensor_size_bytes = tensor_size_mb * 1024 * 1024
    
    print(f"Allocating a {tensor_size_mb}MB tensor in VirtualRAM...")
    ram.allocate_block("my_ai_tensor", tensor_size_bytes, store_data=True) # Request real data
    ram.print_info()
    print("✓ 1MB tensor allocated with real data in VirtualRAM.")

    print("\n" + "=" * 50)

    # 3. Transfer that tensor into VRAM as 
    print(f"Transferring \'my_ai_tensor\' from VirtualRAM to VRAM as \'ai_tensor_buffer\'")
    transfer_start_time = time.time()
    vram_tensor_id = ram.transfer_to_vram("my_ai_tensor", vram, "ai_tensor_buffer")
    transfer_end_time = time.time()
    time_taken_to_transfer = transfer_end_time - transfer_start_time
    
    if vram_tensor_id:
        print(f"✓ Tensor transferred to VRAM. VRAM ID: {vram_tensor_id}")
        print(f"Time taken to transfer: {time_taken_to_transfer:.4f} seconds")
    else:
        print("✗ Failed to transfer tensor to VRAM.")
        return

    print("\n" + "=" * 50)

    # 4. Apply a real AI operation (matrix multiplication) on the VRAM buffer.
    print("Applying a real AI operation (matrix multiplication) on the VRAM buffer...")   
    # Retrieve the transferred data from VRAM (now it should contain actual data)
    vram_block_data = vram.get_texture(vram_tensor_id) # This should now return actual data
    if vram_block_data is None:
        print("✗ Could not retrieve actual tensor data from VRAM for AI operation.")
        return

    # Create a dummy matrix B with real data for multiplication
    # Ensure dimensions are compatible for multiplication
    # Let\'s assume the transferred data is a flattened 1MB (1024*1024 bytes) of uint8
    # We\'ll reshape it to (1024, 1024) for a dummy matrix A
    try:
        matrix_a_reshaped = vram_block_data.reshape((1024, 1024)).astype(np.float32)
    except ValueError:
        print("✗ Could not reshape transferred data for matrix A. Data size might not be suitable.")
        return

    # Create a dummy matrix B (e.g., 1024x512) for multiplication
    matrix_b_data = np.random.rand(1024, 512).astype(np.float32)

    # Load these real matrices into AI Accelerator (which will then load them into VRAM)
    ai_accelerator.set_vram(vram) # Ensure AI accelerator has VRAM reference
    matrix_a_id = ai_accelerator.load_matrix(matrix_a_reshaped, "real_ai_matrix_a")
    matrix_b_id = ai_accelerator.load_matrix(matrix_b_data, "real_ai_matrix_b")

    # Perform real matrix multiplication
    result_id = ai_accelerator.matrix_multiply(matrix_a_id, matrix_b_id, "real_ai_op_result")
    
    if result_id:
        result_matrix = ai_accelerator.get_matrix(result_id)
        if result_matrix is not None:
            print(f"✓ Real AI operation (matrix multiplication) completed. Result shape: {result_matrix.shape}")
            # Optional: Verify a small part of the result
            expected_result = np.dot(matrix_a_reshaped, matrix_b_data)
            if np.allclose(result_matrix[:5,:5], expected_result[:5,:5]): # Compare a small section
                print("✓ Result verification (partial) successful.")
            else:
                print("✗ Result verification (partial) failed.")
        else:
            print("✗ Could not retrieve result matrix from VRAM.")
    else:
        print("✗ Real AI operation failed.")

    print("\n" + "=" * 50)

    # 5. Output logs showing:
    # RAM usage
    # VRAM usage
    # Time taken to transfer
    # Operation success

    print("FINAL LOGS AND STATISTICS")
    print("=" * 50)

    print("\nVirtual RAM Usage:")
    ram.print_info()

    print("\nVRAM Usage:")
    vram_stats = vram.get_stats()
    print(f'  Total memory: {vram_stats["total_memory_gb"]:.1f} GB')
    print(f'  Utilization: {vram_stats["utilization_percent"]:.2f}%')
    print(f'  Framebuffers: {vram_stats["framebuffers_count"]}')
    print(f'  Textures/Buffers: {vram_stats["textures_count"]}')
    print(f'  Allocated bytes: {vram_stats["allocated_bytes"]:,} bytes')

    print(f"\nTime taken to transfer tensor from RAM to VRAM: {time_taken_to_transfer:.4f} seconds")
    success_status = "Success" if result_id else "Failed"
    print(f"AI Operation Success: {success_status}")

    print("\n" + "=" * 50)
    print("Virtual RAM Transfer Test Complete!")
    print("Check ./test_output_ram_transfer/ for any saved frame images (if applicable).")
    print("=" * 50)

    # Clean up (optional)
    # ram.release_block("my_ai_tensor")
    # vram.delete_texture(vram_tensor_id)

if __name__ == "__main__":
    # Ensure the output directory exists
    os.makedirs("./test_output_ram_transfer", exist_ok=True)
    asyncio.run(test_virtual_ram_transfer())