#!/usr/bin/env swift
import Foundation
import CoreAudio

// Extensions and Helpers
extension Int32 {
    var fourCC: String {
        let utf16 = [
            UInt16((self >> 24) & 0xFF),
            UInt16((self >> 16) & 0xFF),
            UInt16((self >> 8) & 0xFF),
            UInt16(self & 0xFF)
        ]
        return String(utf16CodeUnits: utf16, count: 4)
    }
}

// Safer Property Getter
func getPropertyData<T>(objectID: AudioObjectID, selector: AudioObjectPropertySelector, initialValue: T) -> T? {
    var address = AudioObjectPropertyAddress(
        mSelector: selector,
        mScope: kAudioObjectPropertyScopeGlobal,
        mElement: kAudioObjectPropertyElementMain
    )
    var size = UInt32(MemoryLayout<T>.size)
    var value = initialValue
    
    let status = AudioObjectGetPropertyData(objectID, &address, 0, nil, &size, &value)
    if status == noErr {
        return value
    }
    return nil
}

// CFString Helper
func getStringProperty(objectID: AudioObjectID, selector: AudioObjectPropertySelector) -> String? {
    var address = AudioObjectPropertyAddress(
        mSelector: selector,
        mScope: kAudioObjectPropertyScopeGlobal,
        mElement: kAudioObjectPropertyElementMain
    )
    // CFStringRef is just a pointer, so size of Optional<Unmanaged<CFString>> is pointer size
    var size = UInt32(MemoryLayout<Unmanaged<CFString>?>.size)
    var value: Unmanaged<CFString>?
    
    let status = AudioObjectGetPropertyData(objectID, &address, 0, nil, &size, &value)
    if status == noErr, let existingValue = value {
        return existingValue.takeRetainedValue() as String
    }
    return nil
}

func findDeviceByName(_ name: String) -> AudioObjectID? {
    // System Object is 1
    let systemID = AudioObjectID(kAudioObjectSystemObject)
    
    // Get all devices
    var address = AudioObjectPropertyAddress(
        mSelector: kAudioHardwarePropertyDevices,
        mScope: kAudioObjectPropertyScopeGlobal,
        mElement: kAudioObjectPropertyElementMain
    )
    var size: UInt32 = 0
    guard AudioObjectGetPropertyDataSize(systemID, &address, 0, nil, &size) == noErr else { return nil }
    
    let count = Int(size) / MemoryLayout<AudioObjectID>.size
    var deviceIDs = [AudioObjectID](repeating: 0, count: count)
    guard AudioObjectGetPropertyData(systemID, &address, 0, nil, &size, &deviceIDs) == noErr else { return nil }
    
    for id in deviceIDs {
        if let devName = getStringProperty(objectID: id, selector: kAudioDevicePropertyDeviceNameCFString) {
            if devName == name {
                return id
            }
        }
    }
    return nil
}

func findDeviceByUID(_ uid: String) -> AudioObjectID? {
    let systemID = AudioObjectID(kAudioObjectSystemObject)
    var address = AudioObjectPropertyAddress(
        mSelector: kAudioHardwarePropertyDevices,
        mScope: kAudioObjectPropertyScopeGlobal,
        mElement: kAudioObjectPropertyElementMain
    )
    var size: UInt32 = 0
    guard AudioObjectGetPropertyDataSize(systemID, &address, 0, nil, &size) == noErr else { return nil }
    
    let count = Int(size) / MemoryLayout<AudioObjectID>.size
    var deviceIDs = [AudioObjectID](repeating: 0, count: count)
    guard AudioObjectGetPropertyData(systemID, &address, 0, nil, &size, &deviceIDs) == noErr else { return nil }
    
    for id in deviceIDs {
        if let devUID = getStringProperty(objectID: id, selector: kAudioDevicePropertyDeviceUID) {
            if devUID == uid {
                return id
            }
        }
    }
    return nil
}

func createAggregateDevice() {
    print("Searching for devices...")
    
    guard let blackHoleID = findDeviceByName("BlackHole 2ch") else {
        print("Error: BlackHole 2ch not found. Please install it first.")
        exit(1)
    }
    print("Found BlackHole 2ch (ID: \(blackHoleID))")
    
    // Default Input
    var defaultInputID: AudioObjectID = 0
    var size = UInt32(MemoryLayout<AudioObjectID>.size)
    var address = AudioObjectPropertyAddress(
        mSelector: kAudioHardwarePropertyDefaultInputDevice,
        mScope: kAudioObjectPropertyScopeGlobal,
        mElement: kAudioObjectPropertyElementMain
    )
    
    if AudioObjectGetPropertyData(AudioObjectID(kAudioObjectSystemObject), &address, 0, nil, &size, &defaultInputID) != noErr {
         print("Error: Could not find default input.")
         exit(1)
    }
    print("Found Default Input (ID: \(defaultInputID))")
    
    // Check for existing "Hearbit Audio" by UID
    let targetUID = "hearbit_audio_aggregate_v1"
    if let existingID = findDeviceByUID(targetUID) {
        print("Found existing Hearbit Audio device (ID: \(existingID)). Destroying to recreate...")
        if AudioHardwareDestroyAggregateDevice(existingID) != noErr {
            print("Warning: Failed to destroy existing device.")
        } else {
            print("Existing device destroyed.")
        }
        Thread.sleep(forTimeInterval: 0.5)
    }
    
    // Build SubDevice List
    guard let bhUID = getStringProperty(objectID: blackHoleID, selector: kAudioDevicePropertyDeviceUID) else {
        print("Error: Could not get BlackHole UID.")
        exit(1)
    }
    guard let micUID = getStringProperty(objectID: defaultInputID, selector: kAudioDevicePropertyDeviceUID) else {
        print("Error: Could not get Default Input UID.")
        exit(1)
    }
    
    // Dedup: if Mic IS BlackHole (user set BlackHole as default), don't duplicate
    var subDevicesUIDs = [bhUID]
    if micUID != bhUID {
        subDevicesUIDs.append(micUID)
    }
    
    let subDevicesArray = subDevicesUIDs.map { 
        [kAudioSubDeviceUIDKey: $0]
    }
    
    let desc: [String: Any] = [
        kAudioAggregateDeviceNameKey: "Hearbit Audio",
        kAudioAggregateDeviceUIDKey: targetUID,
        kAudioAggregateDeviceIsPrivateKey: Int(0),
        kAudioAggregateDeviceIsStackedKey: Int(0),
        kAudioAggregateDeviceSubDeviceListKey: subDevicesArray
    ]
    
    print("Creating Aggregate Device with UIDs: \(subDevicesUIDs)")
    
    var outID: AudioObjectID = 0
    let err = AudioHardwareCreateAggregateDevice(desc as CFDictionary, &outID)
    
    if err == noErr {
        print("Success! Created 'Hearbit Audio' with ID: \(outID)")
        exit(0)
    } else {
        print("Failed to create device. Error code: \(err) (\(err.fourCC))")
        exit(1)
    }
}

createAggregateDevice()