//
//  CoreCount.swift
//  FlashAttention
//
//  Created by Philip Turner on 6/21/24.
//

#if os(macOS)
import IOKit

/// Finds the core count on macOS devices, using IORegistry.
///
/// Source: [AppleGPUInfo](https://github.com/philipturner/applegpuinfo)
///
/// This code was generated by GPT-4 a few days after launch (early 2023).
/// Since then, it has undergone extensive human review and real-world testing.
/// It proved that proto-AGI could be a practically useful tool, in this case
/// assisting with code creation.
func findCoreCount() -> Int {
  // Create a matching dictionary with "AGXAccelerator" class name
  let matchingDict = IOServiceMatching("AGXAccelerator")
  
  // Get an iterator for matching services
  var iterator: io_iterator_t = 0
  do {
    let io_registry_error =
    IOServiceGetMatchingServices(
      kIOMainPortDefault, matchingDict, &iterator)
    guard io_registry_error == 0 else {
      fatalError(
        "Encountered IORegistry error code \(io_registry_error)")
    }
  }
  
  // Get the first (and only) GPU entry from the iterator
  let gpuEntry = IOIteratorNext(iterator)
  
  // Check if the entry is valid
  if gpuEntry == MACH_PORT_NULL {
    fatalError(
      "Error getting GPU entry at \(#file):\(#line - 5)")
  }
  
  // Release the iterator
  IOObjectRelease(iterator)
  
  // Get the "gpu-core-count" property from gpuEntry
  let key = "gpu-core-count"
  let options: IOOptionBits = 0 // No options needed
  let gpuCoreCount = IORegistryEntrySearchCFProperty(
    gpuEntry, kIOServicePlane, key as CFString, nil, options)
  
  // Check if the property is valid
  if gpuCoreCount == nil {
    fatalError(
      "Error getting gpu-core-count property at \(#file):\(#line - 6)")
  }
  
  // Cast the property to CFNumberRef
  let gpuCoreCountNumber = gpuCoreCount as! CFNumber
  
  // Check if the number type is sInt64
  let type = CFNumberGetType(gpuCoreCountNumber)
  if type != .sInt64Type {
    fatalError(
      "Error: gpu-core-count is not sInt64 at \(#file):\(#line - 3)")
  }
  
  // Get the value of the number as Int64
  var value: Int64 = 0
  let result = CFNumberGetValue(gpuCoreCountNumber, type, &value)
  
  // Check for errors
  if result == false {
    fatalError(
      " Error getting value of gpu-core-count at \(#file):\(#line - 5)")
  }
  
  return Int(value)
}
#endif
