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

extension MacVMEngine {
    // MARK: - Linux base bundle resolution

    /// Resolve the golden **Linux** base bundle to clone from, without kicking off the (heavy) build.
    /// Order: a user-supplied prebuilt bundle if configured and complete, else the engine's own built
    /// Linux base. When neither exists, throw `.baseImageMissing` pointing at the explicit build path —
    /// the multi-GB kernel/rootfs download + desktop provisioning is a deliberate one-time action, never
    /// a surprise on an agent's first turn (mirrors ``ensureBaseBundle`` for macOS).
    func ensureLinuxBaseBundle() throws -> MacVMBundle {
        if let prebuilt = MacVMSettings.linuxBasePrebuiltPath {
            let bundle = MacVMBundle(root: URL(fileURLWithPath: prebuilt))
            if bundle.isComplete(for: .linux) { return bundle }
            throw MacVMError.baseImageMissing(
                "the configured prebuilt Linux base at \(prebuilt) is missing its kernel/initrd/rootfs")
        }
        let built = MacVMBundle(root: builtLinuxBaseDir)
        if built.isComplete(for: .linux) { return built }
        throw MacVMError.baseImageMissing(
            "no Linux base image has been built yet. Build it once via Settings ▸ Virtual Machines ▸ "
                + "Build Linux base image.")
    }

    /// The current Linux base bundle's provisioning status (from its `bundle.json`), for the settings
    /// panel. Prefers a configured prebuilt bundle, else the engine-built one; all-false when neither
    /// is complete. `os` is forced to `.linux` so the panel renders the Linux summary.
    public func currentLinuxBaseStatus() -> MacVMBaseStatus {
        let bundle = MacVMSettings.linuxBasePrebuiltPath.map { MacVMBundle(root: URL(fileURLWithPath: $0)) }
            ?? MacVMBundle(root: builtLinuxBaseDir)
        guard bundle.isComplete(for: .linux) else {
            return MacVMBaseStatus(os: .linux)
        }
        var status = readBaseStatus(bundle)
        status.installed = true
        status.os = .linux
        return status
    }

    /// Filesystem path of the current (complete) Linux base bundle, or `nil` when none exists.
    public func linuxBaseBundleRootPath() -> String? {
        let bundle = MacVMSettings.linuxBasePrebuiltPath.map { MacVMBundle(root: URL(fileURLWithPath: $0)) }
            ?? MacVMBundle(root: builtLinuxBaseDir)
        return bundle.isComplete(for: .linux) ? bundle.root.path : nil
    }

    /// Delete the engine-built Linux base so the next ``buildLinuxBaseImage`` rebuilds from scratch.
    /// Refuses while the base is busy. A configured *prebuilt* base is left untouched.
    ///
    /// Deletes the downloaded artifacts too, so "delete and rebuild" really does start over. Keeping
    /// them meant a base broken by a bad or stale download rebuilt from the very same bytes and failed
    /// identically — the cache is exactly what a user reaches for this button to clear.
    /// `force` skips the busy check and force-cancels whatever is holding the base first — see
    /// ``forceCancelBaseWork()``. The settings panel always forces.
    public func deleteLinuxBaseImage(force: Bool = false) async throws {
        if force {
            await forceCancelBaseWork()
        }
        guard force || !baseIsBusy else {
            throw MacVMError.baseBusy("a base image build is in progress — wait for it to finish, then delete.")
        }
        try? FileManager.default.removeItem(at: builtLinuxBaseDir)
        purgeLinuxArtifactCache()
        baseProgress = nil
    }
}

#if arch(arm64)
extension MacVMEngine {
    // MARK: - Linux base build (fully automated, two in-guest phases)

    /// Build the golden **Linux** base bundle end-to-end (docs/LINUX_VM.md §base build). macOS has no
    /// Linux filesystem tooling, so the ext4 root disk is assembled *inside the guest* across two boots:
    ///
    ///   1. **bootstrap** — boot the downloaded kernel + a tiny busybox initramfs with the empty root
    ///      disk attached and the provisioning payload (narOS-vm rootfs tarball, the vsock agent, the
    ///      provisioning script + systemd units) handed in on a virtio-blk disk. Its `/init` formats the
    ///      disk, unpacks the rootfs, bakes in the agent + first-boot unit, and powers off.
    ///   2. **provision** — boot the installed rootfs; the first-boot unit `apt install`s naros-tier-vm
    ///      (the vsock agent + tier), then writes the readiness flag the host waits on before sealing.
    ///
    /// Long-running (multi-GB download, then two boots with an `apt` install); progress rides
    /// ``currentBaseProgress()``. Invoke explicitly (Settings ▸ Build Linux base image / the spike),
    /// never lazily on an agent turn — like ``buildBaseImage``.
    public func buildLinuxBaseImage() async throws {
        guard Self.isSupported else {
            throw MacVMError.unavailable(Self.unsupportedReason ?? "requires Apple silicon")
        }
        guard !baseIsBusy else {
            throw MacVMError.installFailed("a base image build is already in progress")
        }
        // A complete, provisioned base is a no-op (rebuild after an explicit delete).
        let existing = MacVMBundle(root: builtLinuxBaseDir)
        if existing.isComplete(for: .linux), readBaseStatus(existing).provisioned,
            MacVMSettings.linuxBasePrebuiltPath == nil
        {
            // No build runs, so publish no progress: a lingering `.ready` reads as "a build is in
            // flight" to the settings panel and hides the settled row's rebuild/delete menu.
            baseProgress = nil
            return
        }
        baseBuilding = true
        baseWorkCancelled = false  // a new, deliberate pass — re-enable bundle writes
        baseBuildGuest = .linux
        // Hold the anti-nap assertion for the whole build (the assemble/provision VMs never enter
        // `live`), so a Linux base rebuild that fires at launch after an update isn't App-Napped or
        // auto-terminated mid-build. Released in the defer.
        syncBackgroundActivity()
        defer {
            baseBuilding = false
            baseBuildGuest = nil
            clearTerminalBaseProgress()
            syncBackgroundActivity()
        }
        let fm = FileManager.default
        try fm.createDirectory(at: linuxArtifactsDir, withIntermediateDirectories: true)

        // 1. Resolve (download + cache) the boot artifacts and the guest payload.
        baseProgress = MacVMBaseProgress(
            phase: .downloadingRestoreImage, fraction: 0, detail: "Downloading the narOS kernel, initrd, and root filesystem")
        let artifacts = try await resolveLinuxArtifacts()

        // Build into a temp bundle, publish atomically so an interrupted build never looks complete.
        let tmp = MacVMBundle(root: storageRoot.appendingPathComponent("linux-base.building", isDirectory: true))
        try? fm.removeItem(at: tmp.root)
        try fm.createDirectory(at: tmp.root, withIntermediateDirectories: true)

        do {
            // Kernel + initrd + persisted cmdline live in the bundle; the rootfs disk is created empty
            // and filled by the bootstrap boot.
            try fm.copyItem(at: artifacts.kernel, to: tmp.kernelURL)
            try fm.copyItem(at: artifacts.initrd, to: tmp.initrdURL)
            try Self.defaultLinuxCommandLine.write(to: tmp.cmdlineURL, atomically: true, encoding: .utf8)
            try Self.createDiskImage(
                at: tmp.diskImageURL,
                sizeBytes: UInt64(MacVMSettings.baseDiskGiB) * 1024 * 1024 * 1024)

            // 2. Bootstrap boot — assemble the rootfs in-guest.
            baseProgress = MacVMBaseProgress(
                phase: .installing, fraction: nil, detail: "Assembling the root filesystem (bootstrap boot)")
            try await runLinuxBootstrapBoot(bundle: tmp, artifacts: artifacts)

            // Publish the installed-but-unprovisioned base first (so a provisioning failure doesn't
            // discard the multi-GB rootfs), then provision in place.
            var status = MacVMBaseStatus(
                installed: true, provisioned: false, osVersion: MacVMSettings.defaultLinuxOSVersion, os: .linux)
            writeBaseStatus(tmp, status)
            try? fm.removeItem(at: builtLinuxBaseDir)
            try fm.moveItem(at: tmp.root, to: builtLinuxBaseDir)

            // 3. Provision boot — install the desktop + toolchain; wait for the readiness flag.
            let base = MacVMBundle(root: builtLinuxBaseDir)
            baseProgress = MacVMBaseProgress(
                phase: .provisioning, fraction: nil, detail: "Installing the narOS desktop + toolchain")
            try await runLinuxProvisionBoot(bundle: base)

            status.provisioned = true
            // Stamp the recipe this build provisioned under, so a later app update that bumps the Linux
            // recipe can detect this base as stale and force a rebuild (MacVMEngine+Reprovision).
            status.provisioningRecipe = Self.linuxProvisioningRecipe
            writeBaseStatus(base, status)
            baseProgress = MacVMBaseProgress(phase: .ready, fraction: 1)
        } catch {
            try? fm.removeItem(at: tmp.root)
            baseProgress = nil
            if let e = error as? MacVMError { throw e }
            throw MacVMError.installFailed(String(describing: error))
        }
    }

    // MARK: - Artifact resolution

    struct LinuxArtifacts {
        let kernel: URL           // uncompressed arm64 Image (external Ubuntu kernel)
        let initrd: URL           // external Ubuntu generic initrd (phase-2 boot)
        let rootfsTarball: URL     // narOS-vm root tarball (.tar.gz/.xz), fed to the bootstrap
    }

    /// Download (and cache) the external kernel + initrd and resolve the narOS-vm rootfs tarball. The
    /// kernel is gunzipped/zboot-unwrapped if needed (VZLinuxBootLoader needs an uncompressed `Image`);
    /// the rootfs comes from the narOS OCI artifact (``resolveLinuxRootfs``).
    func resolveLinuxArtifacts() async throws -> LinuxArtifacts {
        guard let kernelSrc = URL(string: MacVMSettings.linuxKernelURL),
            let initrdSrc = URL(string: MacVMSettings.linuxInitrdURL)
        else { throw MacVMError.kernelMissing("the configured Linux kernel/initrd URLs are invalid") }

        // The unwrapped `vmlinux` is *derived* from the downloaded vmlinuz, so it is declared as such:
        // when Ubuntu respins the rolling cloud-image directory the derived Image must be discarded
        // too, or ``prepareBootableKernel``'s "dest exists" fast path would keep booting the previous
        // kernel against the new kernel's modules (MacVMEngine+LinuxArtifactCache).
        let vmlinux = linuxArtifactsDir.appendingPathComponent("vmlinux")
        let kernelDL = try await downloadLinuxArtifact(
            from: kernelSrc, label: "the narOS kernel", derived: [vmlinux])
        let kernel = try prepareBootableKernel(kernelDL, to: vmlinux)
        let initrd = try await downloadLinuxArtifact(from: initrdSrc, label: "the narOS initrd")
        let rootfs = try await resolveLinuxRootfs()
        return LinuxArtifacts(kernel: kernel, initrd: initrd, rootfsTarball: rootfs)
    }

    /// Resolve the payload rootfs tarball. Default: the **narOS-vm** single-blob OCI artifact
    /// (``MacVMSettings/linuxRootfsImage``), cached like the other artifacts. Escape hatch: a non-empty
    /// plain ``MacVMSettings/linuxRootfsURL`` override (e.g. a `file://` local tarball) takes precedence.
    func resolveLinuxRootfs() async throws -> URL {
        let override = MacVMSettings.linuxRootfsURL
        if !override.isEmpty, let url = URL(string: override) {
            return try await downloadLinuxArtifact(from: url, label: "the narOS VM rootfs")
        }
        let ref = MacVMSettings.linuxRootfsImage
        // The OCI blob is a gzip tar; name the cache `.tar.gz` so the payload stager + bootstrap `/init`
        // take the gunzip path (stageProvisioningPayload keys on the file extension). The tag stays in
        // the name so a *tag* bump is visible on disk, but freshness is decided by the layer digest —
        // `:26.07` is republished in place whenever os/ is rebuilt at the same version.
        let cached = linuxArtifactsDir.appendingPathComponent("naros-vm-" + Self.linuxRefSlug(ref) + ".tar.gz")
        do {
            return try await cachedOCIArtifact(ref: ref, cache: cached, label: "the narOS VM rootfs")
        } catch {
            throw MacVMError.installFailed(
                "couldn't download the narOS VM rootfs from \(ref): \(error). Check your network "
                    + "connection, or set \(MacVMSettings.linuxRootfsURLKey) to a local rootfs tarball.")
        }
    }

    /// A filesystem-safe slug for an image ref, so a version/tag bump lands in a fresh cache file.
    static func linuxRefSlug(_ ref: String) -> String {
        ref.map { ($0.isLetter || $0.isNumber || $0 == "." || $0 == "-") ? $0 : "_" }
            .reduce(into: "") { $0.append($1) }
    }

    /// Download `url` into the cache, reporting progress, and reuse the cached copy only while it still
    /// matches upstream (``cachedHTTPArtifact``). `derived` names files computed from this download that
    /// must be discarded with it when upstream moves.
    ///
    /// "Present on disk" is deliberately *not* the freshness test: these artifacts come from Ubuntu's
    /// rolling `releases/<series>/release/` directory, whose filenames are identical across kernel
    /// respins, so a name-keyed cache pins the first download forever (MacVMEngine+LinuxArtifactCache).
    private func downloadLinuxArtifact(
        from url: URL, label: String, derived: [URL] = []
    ) async throws -> URL {
        try await cachedHTTPArtifact(
            from: url,
            to: linuxArtifactsDir.appendingPathComponent(Self.linuxCacheName(for: url)),
            label: label, derived: derived)
    }

    /// A filesystem-safe cache filename for a Linux artifact URL — its own leaf. Note this leaf does
    /// **not** encode the release for Ubuntu's rolling cloud-image URLs, which is why freshness is
    /// decided by the upstream validator rather than the name. Unlike ``restoreCacheName`` it does not
    /// force an `.ipsw` suffix.
    static func linuxCacheName(for url: URL) -> String {
        var leaf = url.lastPathComponent
        if leaf.isEmpty || leaf == "/" { leaf = "artifact" }
        return leaf.map { ($0 == "/" || $0 == ":" || $0 == "?") ? "_" : $0 }.reduce(into: "") { $0.append($1) }
    }

    /// Turn a downloaded `vmlinuz` into a raw arm64 `Image` that `VZLinuxBootLoader` can boot, writing
    /// the result to `dest`. Handles the three shapes Ubuntu has shipped:
    ///   • a **raw arm64 `Image`** (magic `ARM\x64` at offset 0x38) or a pre-unwrapped override → copied through;
    ///   • a **gzip** `vmlinuz` (magic `1f 8b`, ≤24.04) → `gunzip`;
    ///   • an **EFI-zboot** self-decompressing PE (`MZ` + a `zimg` header, ≥24.10) → the embedded payload
    ///     is sliced out and decompressed by its declared compressor (`gzip` via `/usr/bin/gunzip`;
    ///     `zstd` via a discovered `zstd` binary — macOS has no system zstd, so this needs Homebrew/CI
    ///     zstd, else a clear error pointing at a pre-unwrapped `linuxKernelURL`).
    /// Cached: skips if `dest` exists.
    func prepareBootableKernel(_ src: URL, to dest: URL) throws -> URL {
        let fm = FileManager.default
        if fm.fileExists(atPath: dest.path) { return dest }
        let data = try Data(contentsOf: src)

        // Already a raw arm64 Image? (arm64 Linux Image header magic "ARM\x64" at byte 0x38.)
        let armMagic = Data([0x41, 0x52, 0x4D, 0x64])
        if data.count > 0x3C, data.subdata(in: 0x38..<0x3C) == armMagic {
            try fm.copyItem(at: src, to: dest)
            return dest
        }

        // Plain gzip vmlinuz (Ubuntu ≤24.04): gunzip straight to the Image.
        if data.starts(with: [0x1F, 0x8B]) {
            try decompress(payload: src, with: "gzip", to: dest)
            return dest
        }

        // EFI-zboot: MZ PE carrying a compressed Image behind a "zimg" header. Locate the header, read
        // the declared compression + payload size, find the payload by its compressor magic, decompress.
        if data.starts(with: [0x4D, 0x5A]), let z = data.range(of: Data("zimg".utf8)) {
            let hz = z.lowerBound
            let payloadSize = Int(readLE32(data, at: hz + 8))
            let comp = readCString(data, at: hz + 20).lowercased()
            let magics: [String: [UInt8]] = [
                "gzip": [0x1F, 0x8B, 0x08], "zstd": [0x28, 0xB5, 0x2F, 0xFD],
                "lz4": [0x04, 0x22, 0x4D, 0x18],
            ]
            guard let magic = magics[comp],
                let payloadRange = data.range(of: Data(magic), in: hz..<data.endIndex)
            else {
                throw MacVMError.kernelMissing(
                    "the kernel is an EFI-zboot image with unsupported compression '\(comp)' — set "
                        + "linuxKernelURL to a pre-unwrapped raw arm64 Image (scripts/publish-linux-vm-kernel.sh)")
            }
            let start = payloadRange.lowerBound
            let end = payloadSize > 0 ? min(start + payloadSize, data.endIndex) : data.endIndex
            let payload = linuxArtifactsDir.appendingPathComponent("kernel-payload.\(comp)")
            try data.subdata(in: start..<end).write(to: payload)
            defer { try? fm.removeItem(at: payload) }
            try decompress(payload: payload, with: comp, to: dest)
            return dest
        }

        // Unknown shape — copy through and let VZ surface a boot error if it's not a valid Image.
        try fm.copyItem(at: src, to: dest)
        return dest
    }

    private func readLE32(_ d: Data, at offset: Int) -> UInt32 {
        guard offset + 4 <= d.count else { return 0 }
        let b = d[d.startIndex + offset ..< d.startIndex + offset + 4]
        return b.enumerated().reduce(UInt32(0)) { $0 | (UInt32($1.element) << (8 * $1.offset)) }
    }

    private func readCString(_ d: Data, at offset: Int) -> String {
        var bytes: [UInt8] = []
        var i = d.startIndex + offset
        while i < d.endIndex, d[i] != 0, bytes.count < 32 { bytes.append(d[i]); i = d.index(after: i) }
        return String(decoding: bytes, as: UTF8.self)
    }

    /// Decompress one raw compressed stream file to `dest` using the named compressor.
    private func decompress(payload: URL, with comp: String, to dest: URL) throws {
        let fm = FileManager.default
        let tool: URL
        var args: [String]
        switch comp {
        case "gzip":
            tool = URL(fileURLWithPath: "/usr/bin/gunzip")
            args = ["-c", payload.path]  // stdout → dest
        case "zstd":
            guard let zstd = Self.findExecutable("zstd") else {
                throw MacVMError.kernelMissing(
                    "the kernel payload is zstd-compressed but no `zstd` binary was found (macOS has no "
                        + "system zstd — `brew install zstd`, or set linuxKernelURL to a pre-unwrapped raw "
                        + "arm64 Image via scripts/publish-linux-vm-kernel.sh)")
            }
            tool = zstd
            args = ["-d", "-c", payload.path]  // stdout → dest
        default:
            throw MacVMError.kernelMissing("unsupported kernel compression '\(comp)'")
        }
        guard fm.createFile(atPath: dest.path, contents: nil) else {
            throw MacVMError.kernelMissing("could not create \(dest.lastPathComponent)")
        }
        let out = try FileHandle(forWritingTo: dest)
        defer { try? out.close() }
        let proc = Process()
        proc.executableURL = tool
        proc.arguments = args
        proc.standardOutput = out
        try proc.run()
        proc.waitUntilExit()
        // gunzip/zstd exit non-zero on trailing garbage after a valid stream (harmless for us — the
        // Image decodes fully first), so accept a non-empty output over a hard status check.
        let size = ((try? fm.attributesOfItem(atPath: dest.path))?[.size] as? Int) ?? 0
        if proc.terminationStatus != 0 && size == 0 {
            throw MacVMError.kernelMissing("could not \(comp)-decompress the kernel image")
        }
    }

    /// Resolve an executable by name: common Homebrew/system locations, then `PATH` via `/usr/bin/env`.
    static func findExecutable(_ name: String) -> URL? {
        let fm = FileManager.default
        let candidates = ["/opt/homebrew/bin/\(name)", "/usr/local/bin/\(name)", "/usr/bin/\(name)"]
        for c in candidates where fm.isExecutableFile(atPath: c) { return URL(fileURLWithPath: c) }
        let which = Process()
        which.executableURL = URL(fileURLWithPath: "/usr/bin/env")
        which.arguments = ["which", name]
        let pipe = Pipe()
        which.standardOutput = pipe
        which.standardError = FileHandle.nullDevice
        guard (try? which.run()) != nil else { return nil }
        which.waitUntilExit()
        let out = (try? pipe.fileHandleForReading.readToEnd()).flatMap { $0 } ?? Data()
        let path = String(decoding: out, as: UTF8.self).trimmingCharacters(in: .whitespacesAndNewlines)
        return (which.terminationStatus == 0 && !path.isEmpty) ? URL(fileURLWithPath: path) : nil
    }
}
#endif
