import Foundation

// Upstream-pinned caching for the Linux base's downloaded artifacts (docs/LINUX_VM.md §artifact cache).
//
// Every artifact the Linux base build consumes — the external Ubuntu kernel + initrd, the narOS-vm
// rootfs, the guest agents — is fetched from a **mutable** upstream: a rolling cloud-image directory
// (`releases/resolute/release/…`, whose filenames never change across respins) or a mutable OCI tag
// (`:latest`, `:26.07`). The original cache keyed on the *name* — the URL's last path component, or a
// slug of the image ref — and treated "file exists" as "file is current". Neither name changes when
// upstream republishes, so a cached artifact was pinned for life.
//
// That is what desynchronized the kernel from its modules: the kernel is cached under a name that
// never changes, while `scripts/fetch-kernel-modules.sh` re-reads the live manifest on every build. An
// Ubuntu respin (7.0.0-27 → -28) left the guest booting the old kernel with only the new kernel's
// modules present, so `/lib/modules/$(uname -r)` was missing, `modprobe vmw_vsock_virtio_transport`
// failed, and the agent could never bind AF_VSOCK — the guest booted to a login prompt and the host
// timed out waiting for a vsock ping.
//
// So cache on **upstream identity**, not name: an OCI blob pins its layer digest, an HTTP artifact pins
// its ETag (else Last-Modified + Content-Length). The identity is recorded in a sidecar next to the
// file; when it moves, the stale copy is deleted and re-downloaded — along with any `derived` files
// computed from it (the unwrapped `vmlinux`), which would otherwise survive as the stale original.
extension MacVMEngine {

    /// The sidecar recording the upstream identity of a cached artifact. Kept beside the file (rather
    /// than encoded into its name) so re-downloads reuse one path instead of accreting a copy per
    /// upstream revision — these are hundreds of megabytes each.
    static func upstreamStampURL(for cache: URL) -> URL {
        cache.appendingPathExtension("upstream")
    }

    /// The recorded upstream identity of a cached artifact, or `nil` when absent/unreadable. A cache
    /// file with no stamp predates this mechanism and is treated as stale (re-downloaded once).
    static func readUpstreamStamp(for cache: URL) -> String? {
        guard let s = try? String(contentsOf: upstreamStampURL(for: cache), encoding: .utf8) else {
            return nil
        }
        let trimmed = s.trimmingCharacters(in: .whitespacesAndNewlines)
        return trimmed.isEmpty ? nil : trimmed
    }

    static func writeUpstreamStamp(_ id: String, for cache: URL) {
        try? id.write(to: upstreamStampURL(for: cache), atomically: true, encoding: .utf8)
    }

    /// Delete a cached artifact, its stamp, and anything derived from it.
    static func purgeCachedArtifact(_ cache: URL, derived: [URL] = []) {
        let fm = FileManager.default
        try? fm.removeItem(at: cache)
        try? fm.removeItem(at: upstreamStampURL(for: cache))
        for d in derived {
            try? fm.removeItem(at: d)
            try? fm.removeItem(at: upstreamStampURL(for: d))
        }
    }

    /// Resolve a **single-blob OCI artifact** into `cache`, pinned to its layer digest. Reuses the
    /// cached copy only when the tag still resolves to the digest it was downloaded at; otherwise the
    /// stale copy (and `derived`) is deleted and the blob re-fetched. When the registry is unreachable
    /// an existing cached copy is used rather than failing the build — offline should not be fatal for
    /// an artifact we already have.
    func cachedOCIArtifact(
        ref: String, cache: URL, label: String, derived: [URL] = []
    ) async throws -> URL {
        let fm = FileManager.default
        let have = fm.fileExists(atPath: cache.path)

        let upstream: String?
        do {
            upstream = try await OCIArtifact.resolveLayerDigest(imageRef: ref)
        } catch {
            guard have else { throw error }
            return cache  // offline (or registry down) with a copy in hand: use it.
        }
        if have, let upstream, Self.readUpstreamStamp(for: cache) == upstream { return cache }

        if have {
            baseProgress = MacVMBaseProgress(
                phase: .downloadingRestoreImage, fraction: nil,
                detail: "\(label) changed upstream — re-downloading")
        }
        Self.purgeCachedArtifact(cache, derived: derived)
        try fm.createDirectory(
            at: cache.deletingLastPathComponent(), withIntermediateDirectories: true)
        baseProgress = MacVMBaseProgress(
            phase: .downloadingRestoreImage, fraction: nil, detail: "Downloading \(label)")
        let digest = try await OCIArtifact.fetchSingleBlob(imageRef: ref, to: cache)
        Self.writeUpstreamStamp(digest, for: cache)
        return cache
    }

    /// Resolve a plain-HTTP artifact into `cache`, pinned to its upstream validator. Same contract as
    /// ``cachedOCIArtifact``: reuse only while upstream is byte-for-byte the copy we hold, else purge
    /// (with `derived`) and re-download; tolerate an unreachable server when we already have a copy.
    func cachedHTTPArtifact(
        from url: URL, to cache: URL, label: String, derived: [URL] = []
    ) async throws -> URL {
        let fm = FileManager.default
        let have = fm.fileExists(atPath: cache.path)

        let upstream = await Self.httpValidator(for: url)
        if have, let upstream, Self.readUpstreamStamp(for: cache) == upstream { return cache }
        // No validator at all (HEAD refused / no useful headers) and a copy in hand: keep it. Without a
        // validator a re-download can't be shown to change anything, so it's pure cost.
        if have, upstream == nil { return cache }

        if have {
            baseProgress = MacVMBaseProgress(
                phase: .downloadingRestoreImage, fraction: nil,
                detail: "\(label) changed upstream — re-downloading")
        }
        Self.purgeCachedArtifact(cache, derived: derived)
        try fm.createDirectory(
            at: cache.deletingLastPathComponent(), withIntermediateDirectories: true)
        do {
            try await downloadFile(from: url, to: cache) { [weak self] fraction in
                Task {
                    await self?.setBaseProgress(.init(
                        phase: .downloadingRestoreImage, fraction: fraction,
                        detail: "Downloading \(label)"))
                }
            }
        } catch {
            // The purge already happened, so a failed re-download must not leave a half-file behind
            // that the next build would mistake for a complete artifact.
            Self.purgeCachedArtifact(cache, derived: derived)
            throw error
        }
        if let upstream { Self.writeUpstreamStamp(upstream, for: cache) }
        return cache
    }

    /// The upstream identity of an HTTP artifact: its `ETag` when the server offers one, else
    /// `Last-Modified` + `Content-Length` (what Ubuntu's cloud-image mirror serves). `nil` when the
    /// server is unreachable or offers neither — the caller then keeps whatever it has.
    static func httpValidator(for url: URL) async -> String? {
        var req = URLRequest(url: url)
        req.httpMethod = "HEAD"
        req.timeoutInterval = 20
        guard let (_, resp) = try? await URLSession.shared.data(for: req),
            let http = resp as? HTTPURLResponse, http.statusCode == 200
        else { return nil }
        if let etag = http.value(forHTTPHeaderField: "ETag"), !etag.isEmpty { return "etag:\(etag)" }
        let modified = http.value(forHTTPHeaderField: "Last-Modified") ?? ""
        let length = http.value(forHTTPHeaderField: "Content-Length") ?? ""
        if modified.isEmpty, length.isEmpty { return nil }
        return "mtime:\(modified);len:\(length)"
    }

    /// The kernel release (`uname -r`, e.g. `7.0.0-28-generic`) of an uncompressed arm64 `Image`, read
    /// from the `Linux version …` banner every kernel embeds, or `nil` when it can't be found. Used to
    /// prove the staged modules belong to the kernel the base actually boots — the host has no other
    /// way to ask, since it never runs this kernel itself.
    static func kernelRelease(of image: URL) -> String? {
        guard let data = try? Data(contentsOf: image, options: .mappedIfSafe) else { return nil }
        let needle = Data("Linux version ".utf8)
        guard let hit = data.firstRange(of: needle) else { return nil }
        // The release token runs to the first space; cap the scan so a corrupt image can't run away.
        let tail = data[hit.upperBound..<min(hit.upperBound + 128, data.endIndex)]
        let token = tail.prefix { $0 != 0x20 && $0 != 0x00 && $0 != 0x0A }
        let release = String(decoding: token, as: UTF8.self)
        return release.isEmpty ? nil : release
    }

    /// Delete every cached Linux artifact so the next base build re-downloads all of them from scratch.
    /// The escape hatch behind "Delete Linux base image": deleting the *bundle* alone leaves these
    /// multi-hundred-megabyte downloads in place, so a rebuild reproduces a build that failed because
    /// of a bad or stale artifact. Best-effort — a cache that won't delete is re-fetched anyway.
    public func purgeLinuxArtifactCache() {
        try? FileManager.default.removeItem(at: linuxArtifactsDir)
    }
}
