#!/usr/bin/env swift

import Foundation
import Metal

print("=== Test Metal Basic ===")
print("Start time: \(Date())")
fflush(stdout)

do {
    // Try to create Metal device
    guard let device = MTLCreateSystemDefaultDevice() else {
        print("ERROR: Cannot create Metal device")
        exit(1)
    }
    print("Metal device: \(device.name)")
    fflush(stdout)
    
    guard let commandQueue = device.makeCommandQueue() else {
        print("ERROR: Cannot create command queue")
        exit(1)
    }
    print("Command queue created")
    fflush(stdout)
    
    // Try a simple compute operation
    let source = """
    #include <metal_stdlib>
    using namespace metal;
    
    kernel void simple_add(device const float* a [[buffer(0)]],
                          device const float* b [[buffer(1)]],
                          device float* result [[buffer(2)]],
                          uint id [[thread_position_in_grid]]) {
        result[id] = a[id] + b[id];
    }
    """
    
    let library = try device.makeLibrary(source: source, options: nil)
    print("Library compiled")
    fflush(stdout)
    
    guard let function = library.makeFunction(name: "simple_add") else {
        print("ERROR: Cannot create function")
        exit(1)
    }
    print("Function created")
    fflush(stdout)
    
    let pipeline = try device.makeComputePipelineState(function: function)
    print("Pipeline created")
    fflush(stdout)
    
    // Create buffers
    let count = 1024
    guard let bufferA = device.makeBuffer(length: count * MemoryLayout<Float>.stride),
          let bufferB = device.makeBuffer(length: count * MemoryLayout<Float>.stride),
          let bufferResult = device.makeBuffer(length: count * MemoryLayout<Float>.stride) else {
        print("ERROR: Cannot create buffers")
        exit(1)
    }
    
    // Fill buffers
    let ptrA = bufferA.contents().bindMemory(to: Float.self, capacity: count)
    let ptrB = bufferB.contents().bindMemory(to: Float.self, capacity: count)
    for i in 0..<count {
        ptrA[i] = Float(i)
        ptrB[i] = Float(i * 2)
    }
    
    // Execute
    guard let commandBuffer = commandQueue.makeCommandBuffer(),
          let encoder = commandBuffer.makeComputeCommandEncoder() else {
        print("ERROR: Cannot create command buffer/encoder")
        exit(1)
    }
    
    encoder.setComputePipelineState(pipeline)
    encoder.setBuffer(bufferA, offset: 0, index: 0)
    encoder.setBuffer(bufferB, offset: 0, index: 1)
    encoder.setBuffer(bufferResult, offset: 0, index: 2)
    
    let gridSize = MTLSize(width: count, height: 1, depth: 1)
    let threadGroupSize = MTLSize(width: min(count, pipeline.maxTotalThreadsPerThreadgroup), height: 1, depth: 1)
    encoder.dispatchThreads(gridSize, threadsPerThreadgroup: threadGroupSize)
    encoder.endEncoding()
    
    print("Command buffer committed")
    fflush(stdout)
    
    commandBuffer.commit()
    commandBuffer.waitUntilCompleted()
    
    print("Command buffer completed")
    fflush(stdout)
    
    // Verify result
    let ptrResult = bufferResult.contents().bindMemory(to: Float.self, capacity: count)
    print("Result[0] = \(ptrResult[0]), expected 0")
    print("Result[1] = \(ptrResult[1]), expected 3")
    print("Result[100] = \(ptrResult[100]), expected 300")
    fflush(stdout)
    
    print("Test completed successfully!")
    
} catch {
    print("ERROR: \(error)")
    exit(1)
}