import Foundation
#if arch(arm64)
import Virtualization
#endif

#if arch(arm64)
/// A one-shot resume gate: whichever of the VZ stop delegate or the timeout fires first delivers its
/// value to the awaiting continuation exactly once; the loser is a no-op. Thread-safe (both fire from
/// arbitrary queues). Used to await a self-powering-off base-build boot without blocking a thread.
final class OnceResumer: @unchecked Sendable {
    private let lock = NSLock()
    private var cont: CheckedContinuation<Bool, Never>?
    private var pending: Bool?

    /// Attach the awaiting continuation. If a value already arrived, resume immediately.
    func attach(_ c: CheckedContinuation<Bool, Never>) {
        lock.lock()
        if let v = pending {
            lock.unlock()
            c.resume(returning: v)
        } else {
            cont = c
            lock.unlock()
        }
    }

    /// Deliver `value` once. Resumes an already-attached continuation, else latches for `attach`.
    func fire(_ value: Bool) {
        lock.lock()
        if pending != nil { lock.unlock(); return }
        pending = value
        let c = cont
        cont = nil
        lock.unlock()
        c?.resume(returning: value)
    }
}

extension MacVMEngine {
    // MARK: - Linux base build: the two in-guest boots

    /// **Phase 1** — boot the kernel + bootstrap initramfs with the empty root disk and the provisioning
    /// payload shared in over virtiofs, and wait for its `/init` to format the disk, unpack the rootfs,
    /// bake in the agent, and power off. Headless: no graphics/vsock/network — just block + virtiofs +
    /// console (the console log lands next to the bundle for diagnosis).
    func runLinuxBootstrapBoot(bundle: MacVMBundle, artifacts: LinuxArtifacts) async throws {
        let initramfs = try ensureBootstrapInitramfs()
        let payload = try stageProvisioningPayload(rootfsTarball: artifacts.rootfsTarball)
        defer { try? FileManager.default.removeItem(at: payload) }

        let config = try makeLinuxPhaseConfig(
            kernel: bundle.kernelURL, initrd: initramfs, cmdline: "console=hvc0",
            rootDisk: bundle.diskImageURL, consoleLog: bundle.root.appendingPathComponent("bootstrap.log"),
            withNetwork: false, withVsock: false,
            shares: [(Self.nucleicProvisionTag, payload, true)])

        try await bootAndAwaitPowerOff(
            config: config, label: "linux-bootstrap", timeout: 600,
            onTimeout: "the bootstrap boot didn't finish assembling the root filesystem in time")
    }

    /// **Phase 2** — boot the installed rootfs; its first-boot unit installs the desktop + toolchain. We
    /// wait for the agent (vsock), poll the `provisioned` flag it writes (streaming the guest's own
    /// progress line into the base progress), then power off cleanly to seal the base.
    func runLinuxProvisionBoot(bundle: MacVMBundle) async throws {
        let config = try makeLinuxPhaseConfig(
            kernel: bundle.kernelURL, initrd: bundle.initrdURL, cmdline: Self.defaultLinuxCommandLine,
            rootDisk: bundle.diskImageURL, consoleLog: bundle.root.appendingPathComponent("provision.log"),
            withNetwork: true, withVsock: true, shares: [])

        let instance = MacVMInstance(configuration: config, label: "linux-provision")
        try await instance.start()
        defer { Task { await instance.stop() } }

        // Wait for the agent to answer (it starts at multi-user.target, well before apt finishes).
        try await awaitAgentReadyOnInstance(instance, timeout: 300)

        // Poll the readiness flag the provisioner writes when the desktop + toolchain are in.
        let deadline = Date().addingTimeInterval(3600)  // apt of a desktop can take a while
        while Date() < deadline {
            let flag = await runOnInstance(instance, "test -f /var/lib/nucleic/provisioned && echo READY")
            if flag.stdout.contains("READY") {
                await instance.stop()
                return
            }
            let step = await runOnInstance(instance, "cat /var/lib/nucleic/progress 2>/dev/null | tr -d '\\n'")
            let detail = step.stdout.trimmingCharacters(in: .whitespacesAndNewlines)
            baseProgress = MacVMBaseProgress(
                phase: .provisioning, fraction: nil,
                detail: detail.isEmpty ? "Installing the Linux desktop + toolchain" : detail)
            try? await Task.sleep(nanoseconds: 5_000_000_000)
        }
        throw MacVMError.provisionFailed(
            "Linux base provisioning didn't finish within the deadline (see provision.log in the bundle)")
    }

    /// Boot `config` on a background queue and resume once the guest powers itself off (its `/init`
    /// calls `poweroff`). Force-stops and throws on the timeout. The VZ stop delegate and the timeout
    /// race to a one-shot resumer, so no cooperative thread is ever blocked on a semaphore.
    private func bootAndAwaitPowerOff(
        config: VZVirtualMachineConfiguration, label: String, timeout: TimeInterval, onTimeout: String
    ) async throws {
        let instance = MacVMInstance(configuration: config, label: label)
        let resumer = OnceResumer()
        // Register the stop hook BEFORE start so a fast poweroff can't be missed.
        instance.setOnStop { resumer.fire(true) }
        try await instance.start()

        let timeoutTask = Task {
            try? await Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
            resumer.fire(false)
        }
        let poweredOff = await withCheckedContinuation { (c: CheckedContinuation<Bool, Never>) in
            resumer.attach(c)
        }
        timeoutTask.cancel()
        if !poweredOff {
            await instance.stop()
            throw MacVMError.installFailed(onTimeout)
        }
    }

    // MARK: - Agent readiness / exec against an out-of-registry instance

    /// Wait until the agent on `instance` answers a vsock `ping` (used during the base build, where the
    /// instance isn't in the `live` registry).
    private func awaitAgentReadyOnInstance(_ instance: MacVMInstance, timeout: TimeInterval) async throws {
        let deadline = Date().addingTimeInterval(timeout)
        while Date() < deadline {
            if let box = try? await withAgentTimeout(seconds: 5, {
                try await instance.connectAgent(port: MacVMAgentWire.port)
            }) {
                let client = MacVMAgentClient(connection: box)
                if let pong = try? await client.request(op: "ping", timeout: 5),
                    pong["version"]?.intValue == MacVMAgentWire.version
                {
                    await client.close()
                    return
                }
                await client.close()
            }
            try? await Task.sleep(nanoseconds: 2_000_000_000)
        }
        throw MacVMError.agentUnavailable("the Linux guest agent never answered within the boot deadline")
    }

    /// Run one command in `instance` over the vsock exec channel and collect its stdout + exit code.
    private func runOnInstance(_ instance: MacVMInstance, _ command: String) async -> (exitCode: Int32, stdout: String) {
        do {
            let channel = try await openExecChannel(
                instance: instance, label: "linux-provision", workdir: nil, env: [:], remoteBody: command)
            var out = ""
            do {
                for try await line in channel.stdoutLines { out += String(decoding: line, as: UTF8.self) + "\n" }
            } catch {}
            let code = await channel.wait()
            return (code, out)
        } catch {
            return (-1, "")
        }
    }

    // MARK: - Provisioning payload + bootstrap initramfs

    /// Stage the phase-1 provisioning payload into a temp directory the bootstrap mounts over virtiofs:
    /// the Ubuntu rootfs tarball, the vsock agent binary, and the provisioning script + systemd units.
    func stageProvisioningPayload(rootfsTarball: URL) throws -> URL {
        let fm = FileManager.default
        let dir = storageRoot.appendingPathComponent("linux-prov-payload", isDirectory: true)
        try? fm.removeItem(at: dir)
        try fm.createDirectory(at: dir, withIntermediateDirectories: true)

        // Name the rootfs so /init picks the right decompressor.
        let rootfsName = rootfsTarball.pathExtension == "gz" ? "rootfs.tar.gz" : "rootfs.tar.xz"
        try fm.copyItem(at: rootfsTarball, to: dir.appendingPathComponent(rootfsName))

        try fm.copyItem(at: linuxAgentBinaryURL(), to: dir.appendingPathComponent("nucleic-linux-agent"))
        try fm.copyItem(
            at: linuxAsset("guest/nucleic-linux-agent/systemd/nucleic-linux-agent.service"),
            to: dir.appendingPathComponent("nucleic-linux-agent.service"))
        try fm.copyItem(
            at: linuxAsset("guest/linux-base/provision/provision-linux-guest.sh"),
            to: dir.appendingPathComponent("provision-linux-guest.sh"))
        try fm.copyItem(
            at: linuxAsset("guest/linux-base/provision/nucleic-firstboot.service"),
            to: dir.appendingPathComponent("nucleic-firstboot.service"))
        return dir
    }

    /// Build (once, cached) the bootstrap initramfs via `scripts/build-linux-bootstrap.sh` — portable
    /// (busybox download + `ar`/`tar`/`gzip`/`python3`, all present on macOS). Returns its path.
    func ensureBootstrapInitramfs() throws -> URL {
        let fm = FileManager.default
        let out = linuxArtifactsDir.appendingPathComponent("nucleic-bootstrap.cpio.gz")
        if fm.fileExists(atPath: out.path) { return out }
        try fm.createDirectory(at: linuxArtifactsDir, withIntermediateDirectories: true)
        let script = linuxAsset("scripts/build-linux-bootstrap.sh")
        let proc = Process()
        proc.executableURL = URL(fileURLWithPath: "/bin/bash")
        proc.arguments = [script.path]
        var env = ProcessInfo.processInfo.environment
        env["OUT"] = out.path
        proc.environment = env
        try proc.run()
        proc.waitUntilExit()
        guard proc.terminationStatus == 0, fm.fileExists(atPath: out.path) else {
            throw MacVMError.installFailed(
                "could not build the Linux bootstrap initramfs (scripts/build-linux-bootstrap.sh exited \(proc.terminationStatus))")
        }
        return out
    }

    /// Resolve the prebuilt Linux **agent binary**. The Mac host can't compile a Linux binary, so this
    /// is a build/CI artifact: an explicit path, a bundled resource, or the repo build output. Throws a
    /// clear, actionable error when absent.
    func linuxAgentBinaryURL() throws -> URL {
        let fm = FileManager.default
        let candidates = [
            linuxAsset("guest/nucleic-linux-agent/build/nucleic-linux-agent"),
        ]
        for c in candidates where fm.fileExists(atPath: c.path) { return c }
        throw MacVMError.provisionFailed(
            "the Linux guest agent binary is missing. Build it once (in the Linux build container / CI): "
                + "`guest/nucleic-linux-agent/build.sh`, then rebuild the base.")
    }

    /// Resolve a repo-relative Linux asset. Order: `NUCLEIC_LINUX_ASSETS` override, a bundled
    /// `LinuxAssets` resource dir, else the repo root inferred from this source file (dev / spike /
    /// tests). Returns the path even if absent so callers surface a precise "missing X" error.
    func linuxAsset(_ relativePath: String) -> URL {
        let fm = FileManager.default
        if let override = ProcessInfo.processInfo.environment["NUCLEIC_LINUX_ASSETS"] {
            let u = URL(fileURLWithPath: override).appendingPathComponent(relativePath)
            if fm.fileExists(atPath: u.path) { return u }
        }
        if let res = Bundle.main.resourceURL {
            let u = res.appendingPathComponent("LinuxAssets").appendingPathComponent(relativePath)
            if fm.fileExists(atPath: u.path) { return u }
        }
        // Dev fallback: …/Sources/NucleicCore/MacVM/<thisfile> → repo root is four levels up.
        let repoRoot = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent().deletingLastPathComponent()
            .deletingLastPathComponent().deletingLastPathComponent()
        return repoRoot.appendingPathComponent(relativePath)
    }

    // MARK: - Phase config builder

    /// Assemble a VZ config for one base-build boot: an external kernel + initrd, a single writable root
    /// disk, a serial console → `consoleLog`, and (optionally) NAT networking, the vsock agent device,
    /// and virtiofs shares. Deliberately minimal — no graphics/HID (the base build is headless).
    func makeLinuxPhaseConfig(
        kernel: URL, initrd: URL, cmdline: String, rootDisk: URL, consoleLog: URL,
        withNetwork: Bool, withVsock: Bool, shares: [(tag: String, dir: URL, readOnly: Bool)]
    ) throws -> VZVirtualMachineConfiguration {
        let config = VZVirtualMachineConfiguration()
        config.platform = VZGenericPlatformConfiguration()
        let bootLoader = VZLinuxBootLoader(kernelURL: kernel)
        bootLoader.initialRamdiskURL = initrd
        bootLoader.commandLine = cmdline
        config.bootLoader = bootLoader

        config.cpuCount = Self.clampCPU(MacVMSettings.defaultVMCPUs)
        config.memorySize = Self.clampMemory(UInt64(MacVMSettings.defaultVMMemoryGiB) * 1024 * 1024 * 1024)

        let attachment = try VZDiskImageStorageDeviceAttachment(url: rootDisk, readOnly: false)
        config.storageDevices = [VZVirtioBlockDeviceConfiguration(attachment: attachment)]

        if withNetwork {
            let net = VZVirtioNetworkDeviceConfiguration()
            net.macAddress = VZMACAddress.randomLocallyAdministered()
            net.attachment = VZNATNetworkDeviceAttachment()
            config.networkDevices = [net]
        }
        if withVsock {
            config.socketDevices = [VZVirtioSocketDeviceConfiguration()]
        }

        let console = VZVirtioConsoleDeviceSerialPortConfiguration()
        console.attachment = try VZFileSerialPortAttachment(url: consoleLog, append: false)
        config.serialPorts = [console]

        config.entropyDevices = [VZVirtioEntropyDeviceConfiguration()]
        config.memoryBalloonDevices = [VZVirtioTraditionalMemoryBalloonDeviceConfiguration()]

        if !shares.isEmpty {
            config.directorySharingDevices = shares.map { share in
                let fs = VZVirtioFileSystemDeviceConfiguration(tag: share.tag)
                fs.share = VZSingleDirectoryShare(
                    directory: VZSharedDirectory(url: share.dir, readOnly: share.readOnly))
                return fs
            }
        }

        try config.validate()
        return config
    }
}
#endif
