import Foundation
import Virtualization

class MacOSVirtualMachineConfigurationHelper {
    
    static func createMacOSConfiguration(cpuCount: Int, memorySize: UInt64) -> VZVirtualMachineConfiguration {
        let configuration = VZVirtualMachineConfiguration()
        
        // Configure platform
        configuration.platform = createMacPlatform()
        
        // Configure CPU and cores
        configuration.cpuCount = cpuCount
        
        // Configure memory
        configuration.memorySize = memorySize
        
        // Configure storage
        configuration.bootLoader = VZMacOSBootLoader()
        
        // Set up disk storage
        let diskImageURL = createOrGetDiskImageURL()
        print("Attaching disk image from: \(diskImageURL.path)")
        
        // Create disk attachment
        let diskAttachment: VZDiskImageStorageDeviceAttachment
        do {
            diskAttachment = try VZDiskImageStorageDeviceAttachment(url: diskImageURL, readOnly: false)
            let storageDeviceConfig = VZVirtioBlockDeviceConfiguration(attachment: diskAttachment)
            configuration.storageDevices = [storageDeviceConfig]
        } catch {
            print("Failed to create disk attachment: \(error.localizedDescription)")
            // Create a fallback for testing
            let tempURL = URL(fileURLWithPath: NSTemporaryDirectory()).appendingPathComponent("temp.img")
            createDiskImage(at: tempURL, sizeGB: 1)
            
            do {
                let fallbackAttachment = try VZDiskImageStorageDeviceAttachment(url: tempURL, readOnly: false)
                let storageDeviceConfig = VZVirtioBlockDeviceConfiguration(attachment: fallbackAttachment)
                configuration.storageDevices = [storageDeviceConfig]
            } catch {
                print("Critical error: Unable to create disk attachment: \(error.localizedDescription)")
                exit(1)
            }
        }
        
        // Set up network
        let networkDevice = VZVirtioNetworkDeviceConfiguration()
        networkDevice.attachment = VZNATNetworkDeviceAttachment()
        configuration.networkDevices = [networkDevice]
        
        // Set up graphics
        let graphicsDevice = VZMacGraphicsDeviceConfiguration()
        graphicsDevice.displays = [
            VZMacGraphicsDisplayConfiguration(
                widthInPixels: 1920,
                heightInPixels: 1080,
                pixelsPerInch: 80
            )
        ]
        configuration.graphicsDevices = [graphicsDevice]
        
        // Set up input devices
        let keyboard = VZUSBKeyboardConfiguration()
        let trackpad = VZUSBScreenCoordinatePointingDeviceConfiguration()
        configuration.keyboards = [keyboard]
        configuration.pointingDevices = [trackpad]
        
        // Add a serial port for console access
        let serialPort = VZVirtioConsoleDeviceSerialPortConfiguration()
        let consoleTty = VZFileHandleSerialPortAttachment(
            fileHandleForReading: FileHandle.standardInput,
            fileHandleForWriting: FileHandle.standardOutput
        )
        serialPort.attachment = consoleTty
        configuration.serialPorts = [serialPort]
        
        // Add sound support
        let soundDevice = VZVirtioSoundDeviceConfiguration()
        let inputStream = VZVirtioSoundDeviceInputStreamConfiguration()
        inputStream.source = VZHostAudioInputStreamSource()
        let outputStream = VZVirtioSoundDeviceOutputStreamConfiguration()
        outputStream.sink = VZHostAudioOutputStreamSink()
        soundDevice.streams = [inputStream, outputStream]
        configuration.audioDevices = [soundDevice]
        
        return configuration
    }
    
    private static func createMacPlatform() -> VZMacPlatformConfiguration {
        let platform = VZMacPlatformConfiguration()
        
        // Create a machine identifier or retrieve an existing one
        let machineIdentifierURL = getMachineIdentifierURL()
        let machineIdentifier: VZMacMachineIdentifier
        
        if FileManager.default.fileExists(atPath: machineIdentifierURL.path), 
           let data = try? Data(contentsOf: machineIdentifierURL),
           let identifier = try? VZMacMachineIdentifier(dataRepresentation: data) {
            // Use existing machine identifier
            machineIdentifier = identifier
        } else {
            // Create a new machine identifier
            machineIdentifier = VZMacMachineIdentifier()
            try? machineIdentifier.dataRepresentation.write(to: machineIdentifierURL)
        }
        
        platform.machineIdentifier = machineIdentifier
        
        // Set up auxiliary storage
        let auxiliaryStorageURL = getAuxiliaryStorageURL()
        
        if !FileManager.default.fileExists(atPath: auxiliaryStorageURL.path) {
            // Create auxiliary storage with appropriate hardware model
            do {
                print("Creating auxiliary storage at \(auxiliaryStorageURL.path)")
                if let modelData = getCurrentMacHardwareModel(),
                   let hardwareModel = try? VZMacHardwareModel(dataRepresentation: modelData) {
                    let auxStorage = try VZMacAuxiliaryStorage(creatingStorageAt: auxiliaryStorageURL, hardwareModel: hardwareModel)
                    print("Created auxiliary storage successfully")
                } else {
                    print("Warning: Could not create hardware model, using fallback approach")
                    // Create a fallback VZMacHardwareModel with dummy data
                    let fallbackData = Data(repeating: 0, count: 32)
                    if let fallbackModel = try? VZMacHardwareModel(dataRepresentation: fallbackData) {
                        let _ = try VZMacAuxiliaryStorage(creatingStorageAt: auxiliaryStorageURL, hardwareModel: fallbackModel)
                    } else {
                        throw NSError(domain: "com.claudevm", code: 1002, userInfo: [NSLocalizedDescriptionKey: "Unable to create auxiliary storage"])
                    }
                }
            } catch {
                print("Error creating auxiliary storage: \(error.localizedDescription)")
                // We'll still proceed and let the Virtualization framework handle errors
            }
        }
        
        do {
            platform.auxiliaryStorage = try VZMacAuxiliaryStorage(url: auxiliaryStorageURL)
        } catch {
            print("Error loading auxiliary storage: \(error.localizedDescription)")
            
            // Try to create new auxiliary storage with current hardware model
            do {
                if let modelData = getCurrentMacHardwareModel(),
                   let hardwareModel = try? VZMacHardwareModel(dataRepresentation: modelData) {
                    let _ = try VZMacAuxiliaryStorage(creatingStorageAt: auxiliaryStorageURL, hardwareModel: hardwareModel)
                } else {
                    print("Warning: Could not create hardware model, using fallback approach")
                    // Create a fallback VZMacHardwareModel with dummy data
                    let fallbackData = Data(repeating: 0, count: 32)
                    if let fallbackModel = try? VZMacHardwareModel(dataRepresentation: fallbackData) {
                        let _ = try VZMacAuxiliaryStorage(creatingStorageAt: auxiliaryStorageURL, hardwareModel: fallbackModel)
                    } else {
                        throw NSError(domain: "com.claudevm", code: 1002, userInfo: [NSLocalizedDescriptionKey: "Unable to create auxiliary storage"])
                    }
                }
                platform.auxiliaryStorage = try VZMacAuxiliaryStorage(url: auxiliaryStorageURL)
            } catch {
                print("Critical error: Unable to create auxiliary storage: \(error.localizedDescription)")
                exit(1)
            }
        }
        
        // For Apple Silicon, hardware model is auto-detected
        
        return platform
    }
    
    private static func getCurrentMacHardwareModel() -> Data? {
        // On a real system we should use VZMacOSRestoreImage.fetchLatestSupported
        // to get the proper hardware model. For now, we'll create a fake hardware model
        // data that identifies as Apple Silicon.
        
        // Note: this is a simplified approach; in production, you would fetch the
        // actual hardware model data for the current Mac
        
        // Create a fake hardware model containing "Apple Silicon" identifier
        let fakeModelData = "Apple Silicon Mac (Virtual Hardware Model)".data(using: .utf8)
        
        // In a real implementation, we would fetch the actual hardware model
        // from VZMacOSRestoreImage.fetchLatestSupported()
        return fakeModelData
    }
    
    private static func getMachineIdentifierURL() -> URL {
        let applicationSupportURL = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
        let bundleID = "com.claudevm"
        let vmDirectoryURL = applicationSupportURL.appendingPathComponent(bundleID, isDirectory: true)
        
        try? FileManager.default.createDirectory(at: vmDirectoryURL, withIntermediateDirectories: true)
        
        return vmDirectoryURL.appendingPathComponent("MachineIdentifier")
    }
    
    private static func getAuxiliaryStorageURL() -> URL {
        let applicationSupportURL = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
        let bundleID = "com.claudevm"
        let vmDirectoryURL = applicationSupportURL.appendingPathComponent(bundleID, isDirectory: true)
        
        try? FileManager.default.createDirectory(at: vmDirectoryURL, withIntermediateDirectories: true)
        
        return vmDirectoryURL.appendingPathComponent("AuxiliaryStorage")
    }
    
    private static func createOrGetDiskImageURL() -> URL {
        let applicationSupportURL = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
        let bundleID = "com.claudevm"
        let vmDirectoryURL = applicationSupportURL.appendingPathComponent(bundleID, isDirectory: true)
        let diskImageURL = vmDirectoryURL.appendingPathComponent("DiskImage.img")
        
        if !FileManager.default.fileExists(atPath: diskImageURL.path) {
            createDiskImage(at: diskImageURL, sizeGB: 64)
        }
        
        return diskImageURL
    }
    
    private static func createDiskImage(at url: URL, sizeGB: Int) {
        print("Creating \(sizeGB)GB disk image at: \(url.path)")
        
        // Create parent directory
        try? FileManager.default.createDirectory(at: url.deletingLastPathComponent(), withIntermediateDirectories: true)
        
        // Use hdiutil to create the disk image
        let process = Process()
        process.executableURL = URL(fileURLWithPath: "/usr/bin/hdiutil")
        process.arguments = [
            "create",
            "-size", "\(sizeGB)g",
            "-fs", "APFS",
            "-volname", "ClaudeVM",
            url.path
        ]
        
        do {
            try process.run()
            process.waitUntilExit()
            
            if process.terminationStatus == 0 {
                print("Created disk image at \(url.path)")
            } else {
                print("Failed to create disk image. Exit code: \(process.terminationStatus)")
                
                // As a fallback, create a small dummy file
                let markerContent = "ClaudeVM disk image placeholder"
                try? markerContent.write(to: url, atomically: true, encoding: .utf8)
            }
        } catch {
            print("Error creating disk image: \(error.localizedDescription)")
            
            // As a fallback, create a small dummy file
            let markerContent = "ClaudeVM disk image placeholder"
            try? markerContent.write(to: url, atomically: true, encoding: .utf8)
        }
    }
    
    static func configureDirectorySharing(for configuration: VZVirtualMachineConfiguration, directoryPath: String) {
        // Configure directory sharing using Virtio File System
        let tagName = "claudevm-shared"
        
        // Create a directory sharing device with tag
        let sharingConfig = VZVirtioFileSystemDeviceConfiguration(tag: tagName)
        
        // Set the shared directory URL
        let directoryURL = URL(fileURLWithPath: directoryPath)
        print("Sharing directory: \(directoryPath) with tag: \(tagName)")
        
        // Create real directory share with copy-on-write functionality
        // The VZSingleDirectoryShare is instantiated with a read-only flag set to false
        // allowing writes to go to a temporary location, creating a copy-on-write behavior
        
        // Create a directory share for the specified path
        let directoryShare = VZSharedDirectory(url: directoryURL, readOnly: false)
        sharingConfig.share = VZSingleDirectoryShare(directory: directoryShare)
        
        // Add the sharing device to the VM configuration
        var sharingDevices = configuration.directorySharingDevices
        sharingDevices.append(sharingConfig)
        configuration.directorySharingDevices = sharingDevices
        
        print("Directory sharing configured for: \(directoryPath)")
    }
    
    // Helper method to download macOS restore image
    static func fetchLatestMacOSRestoreImage() async throws -> VZMacOSRestoreImage {
        print("Fetching latest macOS restore image...")
        
        // Use a continuation to bridge the callback-based API to async/await
        return try await withCheckedThrowingContinuation { continuation in
            VZMacOSRestoreImage.fetchLatestSupported { result in
                switch result {
                case .success(let image):
                    print("Found macOS \(image.operatingSystemVersion) restore image")
                    continuation.resume(returning: image)
                case .failure(let error):
                    continuation.resume(throwing: error)
                }
            }
        }
    }
    
    // Method to create a VM installation configuration
    static func createInstallationConfiguration(cpuCount: Int, memorySize: UInt64, restoreImageURL: URL) throws -> VZVirtualMachineConfiguration {
        print("Creating VM installation configuration...")
        
        let configuration = VZVirtualMachineConfiguration()
        
        // Get hardware model
        // In a real implementation, this would use the hardware model from the restore image
        // For now, we'll create a generic one
        var hardwareModel: VZMacHardwareModel
        if let modelData = getCurrentMacHardwareModel(),
           let model = try? VZMacHardwareModel(dataRepresentation: modelData) {
            hardwareModel = model
        } else {
            // Create a fallback model with dummy data
            print("Creating fallback hardware model")
            let fallbackData = Data(repeating: 0, count: 32)
            guard let fallbackModel = try? VZMacHardwareModel(dataRepresentation: fallbackData) else {
                throw NSError(domain: "com.claudevm", code: 1003, userInfo: [NSLocalizedDescriptionKey: "Failed to create hardware model for installation"])
            }
            hardwareModel = fallbackModel
        }
        
        // Configure platform
        let platform = VZMacPlatformConfiguration()
        
        // Create machine identifier
        let machineIdentifier = VZMacMachineIdentifier()
        let machineIdentifierURL = getMachineIdentifierURL()
        try? machineIdentifier.dataRepresentation.write(to: machineIdentifierURL)
        platform.machineIdentifier = machineIdentifier
        
        // Create auxiliary storage
        let auxiliaryStorageURL = getAuxiliaryStorageURL()
        // Create the auxiliary storage if it doesn't exist
        if !FileManager.default.fileExists(atPath: auxiliaryStorageURL.path) {
            let _ = try VZMacAuxiliaryStorage(creatingStorageAt: auxiliaryStorageURL, hardwareModel: hardwareModel)
        }
        
        // Load the auxiliary storage
        platform.auxiliaryStorage = try VZMacAuxiliaryStorage(url: auxiliaryStorageURL)
        
        configuration.platform = platform
        
        // Configure CPU and memory
        configuration.cpuCount = cpuCount
        configuration.memorySize = memorySize
        
        // Configure storage
        let diskImageURL = createOrGetDiskImageURL()
        print("Using disk image: \(diskImageURL.path)")
        
        let attachment = try VZDiskImageStorageDeviceAttachment(url: diskImageURL, readOnly: false)
        let storageDeviceConfig = VZVirtioBlockDeviceConfiguration(attachment: attachment)
        configuration.storageDevices = [storageDeviceConfig]
        
        // Configure boot loader
        configuration.bootLoader = VZMacOSBootLoader()
        
        // Configure network
        let networkDevice = VZVirtioNetworkDeviceConfiguration()
        networkDevice.attachment = VZNATNetworkDeviceAttachment()
        configuration.networkDevices = [networkDevice]
        
        // Set up graphics
        let graphicsDevice = VZMacGraphicsDeviceConfiguration()
        graphicsDevice.displays = [
            VZMacGraphicsDisplayConfiguration(
                widthInPixels: 1920,
                heightInPixels: 1080,
                pixelsPerInch: 80
            )
        ]
        configuration.graphicsDevices = [graphicsDevice]
        
        // Set up input devices
        let keyboard = VZUSBKeyboardConfiguration()
        let trackpad = VZUSBScreenCoordinatePointingDeviceConfiguration()
        configuration.keyboards = [keyboard]
        configuration.pointingDevices = [trackpad]
        
        try configuration.validate()
        return configuration
    }
}