#!/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))")
    
    // --- PART 1: Hearbit Audio (Input: Mic + BlackHole) ---
    print("\n--- Creating 'Hearbit Audio' (Input) ---")
    
    // 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"
    let inputUID = "hearbit_audio_aggregate_v1"
    if let existingID = findDeviceByUID(inputUID) {
        print("Found existing Hearbit Audio (ID: \(existingID)). Destroying...")
        AudioHardwareDestroyAggregateDevice(existingID)
        Thread.sleep(forTimeInterval: 0.5)
    }
    
    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)
    }
    
    var subDevicesUIDs = [bhUID]
    if micUID != bhUID {
        subDevicesUIDs.append(micUID)
    }
    
    let subDevicesArray = subDevicesUIDs.map { [kAudioSubDeviceUIDKey: $0] }
    let inputDesc: [String: Any] = [
        kAudioAggregateDeviceNameKey: "Hearbit Audio",
        kAudioAggregateDeviceUIDKey: inputUID,
        kAudioAggregateDeviceIsPrivateKey: Int(0),
        kAudioAggregateDeviceIsStackedKey: Int(0),
        kAudioAggregateDeviceSubDeviceListKey: subDevicesArray
    ]
    
    var outInputID: AudioObjectID = 0
    let errIn = AudioHardwareCreateAggregateDevice(inputDesc as CFDictionary, &outInputID)
    if errIn == noErr {
        print("Success! Created 'Hearbit Audio' with ID: \(outInputID)")
    } else {
        print("Failed to create 'Hearbit Audio'. Error: \(errIn)")
    }


    // --- PART 2: Cleanup Unstable "Hearbit Speakers" ---
    // The previous "Hearbit Speakers" device caused MS Teams to crash.
    // We strictly remove it here to restore stability.
    print("\n--- Cleaning up Unstable Devices ---")
    let stopOutputUID = "hearbit_speakers_aggregate_v1"
    if let existingOutID = findDeviceByUID(stopOutputUID) {
        print("Found unstable 'Hearbit Speakers' (ID: \(existingOutID)). Removing to fix Teams crash...")
        let errDist = AudioHardwareDestroyAggregateDevice(existingOutID)
        if errDist == noErr {
            print("Successfully removed unstable device.")
        } else {
             print("Warning: Failed to remove device. Error: \(errDist)")
        }
    } else {
        print("No unstable 'Hearbit Speakers' found. System is clean.")
    }
    
    exit(0)
}

createAggregateDevice()