import Containerization
import ContainerizationExtras
import Foundation
import Virtualization

/// Attaches a virtio **memory balloon** to a container's VM and lets the engine retarget it at
/// runtime, so unused guest memory is returned to the host instead of being held until teardown.
///
/// This is a `VZInstanceExtension`: `configureVZ` adds the balloon device while the VZ config is
/// built, and `didCreate` captures the live device once the VM exists. Upstream
/// `apple/containerization` already plumbs `VMConfiguration.extensions` into the VZ config, but
/// `LinuxContainer` — the only entry point the engine uses — didn't forward them; our vendored patch
/// adds `LinuxContainer.Configuration.vmExtensions` to bridge that gap (see
/// `third_party/containerization/PATCHES.md`).
///
/// `ContainerEngine`'s autoballoon loop drives `setTarget` from the guest's live working set; the
/// `MemoryManagementLevel` / ``BalloonPolicy`` picks how aggressively.
///
/// Thread-safety: `configureVZ`/`didCreate` run on the framework's VM-setup path while `setTarget`
/// is called from the engine actor's loop, so a lock guards the stored handles. The balloon device
/// is non-`Sendable` and must only be touched on the VM's own dispatch queue, so every mutation is
/// dispatched onto `queue`.
public final class MemoryBalloon: VZInstanceExtension, @unchecked Sendable {
    private let lock = NSLock()
    private var device: VZVirtioTraditionalMemoryBalloonDevice?
    private var queue: DispatchQueue?
    private var lastTarget: UInt64 = 0

    public init() {}

    /// Add a single traditional memory-balloon device to the VM configuration.
    public func configureVZ(
        _ config: inout VZVirtualMachineConfiguration,
        allocator: any AddressAllocator<Character>,
        storageDeviceCount: Int,
        mountsByID: [String: [Mount]]
    ) throws {
        config.memoryBalloonDevices = [VZVirtioTraditionalMemoryBalloonDeviceConfiguration()]
    }

    /// Capture the live balloon device + the VM's queue once the machine has been created.
    public func didCreate(_ instance: VZVirtualMachineInstance) throws {
        let vm = instance.vzVirtualMachine
        let q = instance.vmQueue
        // VZVirtualMachine properties must be read on its own queue.
        let dev = q.sync { vm.memoryBalloonDevices.first as? VZVirtioTraditionalMemoryBalloonDevice }
        lock.withLock {
            self.device = dev
            self.queue = q
        }
    }

    /// Set the balloon's target guest-memory size (bytes). No-op until the device is captured, or
    /// when the change is smaller than `hysteresis` (avoids churning the balloon on tiny deltas).
    /// Returns the target actually applied, or `nil` if it was skipped.
    @discardableResult
    public func setTarget(_ bytes: UInt64, hysteresis: UInt64) -> UInt64? {
        let (dev, q, apply): (VZVirtioTraditionalMemoryBalloonDevice?, DispatchQueue?, Bool) =
            lock.withLock {
                guard let dev = device, let q = queue else { return (nil, nil, false) }
                let delta = bytes > lastTarget ? bytes - lastTarget : lastTarget - bytes
                if lastTarget != 0 && delta < hysteresis { return (dev, q, false) }
                lastTarget = bytes
                return (dev, q, true)
            }
        guard apply, let dev, let q else { return nil }
        // VZ requires queue-affinity, so the device is only ever touched inside this queue hop;
        // `nonisolated(unsafe)` acknowledges that to the Sendable-capture check.
        nonisolated(unsafe) let device = dev
        q.async { device.targetVirtualMachineMemorySize = bytes }
        return bytes
    }
}
