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.
    public func deleteLinuxBaseImage() throws {
        guard !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)
        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 (Ubuntu rootfs tarball, the vsock agent, the
    ///      provisioning script + systemd units) shared in over virtiofs. 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 installs a minimal Wayland
    ///      desktop + dev toolchain, 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
        {
            baseProgress = MacVMBaseProgress(phase: .ready, fraction: 1)
            return
        }
        baseBuilding = true
        defer { baseBuilding = false }
        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 Linux 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: "24.04", 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 Linux desktop + toolchain")
            try await runLinuxProvisionBoot(bundle: base)

            status.provisioned = true
            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
        let initrd: URL           // Ubuntu generic initrd (phase-2 boot)
        let rootfsTarball: URL     // Ubuntu root tarball (.tar.xz/.gz), fed to the bootstrap
    }

    /// Download (and cache, version-keyed by filename) the kernel + initrd + rootfs tarball, gunzipping
    /// the kernel if it arrived as a gzip `vmlinuz` (VZLinuxBootLoader needs an uncompressed `Image`).
    func resolveLinuxArtifacts() async throws -> LinuxArtifacts {
        guard let kernelSrc = URL(string: MacVMSettings.linuxKernelURL),
            let initrdSrc = URL(string: MacVMSettings.linuxInitrdURL),
            let rootfsSrc = URL(string: MacVMSettings.linuxRootfsURL)
        else { throw MacVMError.kernelMissing("the configured Linux artifact URLs are invalid") }

        let kernelDL = try await downloadLinuxArtifact(from: kernelSrc, label: "kernel")
        let kernel = try gunzipIfNeeded(kernelDL, to: linuxArtifactsDir.appendingPathComponent("vmlinux"))
        let initrd = try await downloadLinuxArtifact(from: initrdSrc, label: "initrd")
        let rootfs = try await downloadLinuxArtifact(from: rootfsSrc, label: "rootfs")
        return LinuxArtifacts(kernel: kernel, initrd: initrd, rootfsTarball: rootfs)
    }

    /// Download `url` into the version-keyed cache unless present, reporting progress.
    private func downloadLinuxArtifact(from url: URL, label: String) async throws -> URL {
        let fm = FileManager.default
        let dest = linuxArtifactsDir.appendingPathComponent(Self.linuxCacheName(for: url))
        if fm.fileExists(atPath: dest.path) { return dest }
        try await downloadFile(from: url, to: dest) { [weak self] fraction in
            Task {
                await self?.setBaseProgress(.init(
                    phase: .downloadingRestoreImage, fraction: fraction,
                    detail: "Downloading the Linux \(label)"))
            }
        }
        return dest
    }

    /// A filesystem-safe cache filename for a Linux artifact URL — its own leaf (which encodes the
    /// release), so different versions land in different files. 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) }
    }

    /// If `src` is gzip-compressed (magic `1f 8b`), decompress it to `dest`; otherwise hardlink/copy it
    /// through unchanged. Uses `/usr/bin/gunzip` (present on macOS). Cached: skips if `dest` exists.
    func gunzipIfNeeded(_ src: URL, to dest: URL) throws -> URL {
        let fm = FileManager.default
        if fm.fileExists(atPath: dest.path) { return dest }
        let handle = try FileHandle(forReadingFrom: src)
        let magic = try handle.read(upToCount: 2) ?? Data()
        try? handle.close()
        if magic == Data([0x1F, 0x8B]) {
            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 = URL(fileURLWithPath: "/usr/bin/gunzip")
            proc.arguments = ["-c", src.path]
            proc.standardOutput = out
            try proc.run()
            proc.waitUntilExit()
            guard proc.terminationStatus == 0 else {
                throw MacVMError.kernelMissing("could not gunzip the kernel image")
            }
        } else {
            try fm.copyItem(at: src, to: dest)
        }
        return dest
    }
}
#endif
