apple--containerization/Sources/Containerization/Image/ImageStore/ImageStore.swift

//===----------------------------------------------------------------------===//
// Copyright © 2025-2026 Apple Inc. and the Containerization project authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//===----------------------------------------------------------------------===//

import ContainerizationError
import ContainerizationExtras
import ContainerizationOCI
import Foundation

/// An ImageStore handles the mappings between an image's
/// reference and the underlying descriptor inside of a content store.
public actor ImageStore: Sendable {
    /// The ImageStore path it was created with.
    public nonisolated let path: URL

    private let referenceManager: ReferenceManager
    internal let contentStore: ContentStore
    internal let lock: AsyncLock = AsyncLock()

    public init(path: URL, contentStore: ContentStore? = nil) throws {
        try FileManager.default.createDirectory(at: path, withIntermediateDirectories: true)

        if let contentStore {
            self.contentStore = contentStore
        } else {
            self.contentStore = try LocalContentStore(path: path.appendingPathComponent("content"))
        }

        self.path = path
        self.referenceManager = try ReferenceManager(path: path)
    }

    /// Return the default image store for the current user.
    public static let `default`: ImageStore = {
        do {
            let root = try defaultRoot()
            return try ImageStore(path: root)
        } catch {
            fatalError("unable to initialize default ImageStore \(error)")
        }
    }()

    private static func defaultRoot() throws -> URL {
        let root = FileManager.default.urls(
            for: .applicationSupportDirectory,
            in: .userDomainMask
        ).first
        guard let root else {
            throw ContainerizationError(.notFound, message: "unable to get Application Support directory for current user")
        }
        return root.appendingPathComponent("com.apple.containerization")

    }
}

extension ImageStore {
    /// Get an image from the `ImageStore`.
    ///
    /// - Parameters:
    ///   - reference: Name of the image.
    ///   - pull: Pull the image if it is not found.
    ///
    /// - Returns: A `Containerization.Image`  object whose `reference` matches the given string.
    ///   This  method throws a `ContainerizationError(code: .notFound)` if the provided reference does not exist in the `ImageStore`.
    public func get(reference: String, pull: Bool = false) async throws -> Image {
        do {
            let desc = try await self.referenceManager.get(reference: reference)
            return Image(description: desc, contentStore: self.contentStore)
        } catch let error as ContainerizationError {
            if error.code == .notFound && pull {
                return try await self.pull(reference: reference)
            }
            throw error
        }
    }

    /// Get a list of all images in the `ImageStore`.
    ///
    /// - Returns: A `[Containerization.Image]` for all the images in the `ImageStore`.
    public func list() async throws -> [Image] {
        try await self.referenceManager.list().map { desc in
            Image(description: desc, contentStore: self.contentStore)
        }
    }

    /// Create a new image in the `ImageStore`.
    ///
    /// - Parameters:
    ///   - description: The underlying `Image.Description` that contains information about the reference and index descriptor for the image to be created.
    ///
    /// - Note: It is assumed that the underlying manifests and blob layers for the image already exists in the `ContentStore` that the `ImageStore` was initialized with. This method is invoked when the `pull(...)` , `load(...)` and `tag(...)` methods are used.
    /// - Returns: A `Containerization.Image`
    @discardableResult
    public func create(description: Image.Description) async throws -> Image {
        try await self.lock.withLock { ctx in
            try await self._create(description: description, lock: ctx)
        }
    }

    @discardableResult
    internal func _create(description: Image.Description, lock: AsyncLock.Context) async throws -> Image {
        try await self.referenceManager.create(description: description)
        return Image(description: description, contentStore: self.contentStore)
    }

    /// Delete an image from the `ImageStore`.
    ///
    /// - Parameters:
    ///   - reference: Name of the image that is to be deleted.
    ///   - performCleanup: Perform a garbage collection on the `ContentStore`, removing all unreferenced image layers and manifests,
    public func delete(reference: String, performCleanup: Bool = false) async throws {
        try await self.lock.withLock { lockCtx in
            try await self.referenceManager.delete(reference: reference)
            if performCleanup {
                try await self._cleanUpOrphanedBlobs(lockCtx)
            }
        }
    }

    /// Clean up orphaned blobs that are no longer referenced by any image.
    ///
    /// - Returns: Returns a tuple of `(deleted, freed)`.
    ///   `deleted` :  A  list of the names of the content items that were deleted from the `ContentStore`,
    ///   `freed` : The total size of the items that were deleted.
    @discardableResult
    public func cleanUpOrphanedBlobs() async throws -> (deleted: [String], freed: UInt64) {
        try await self.lock.withLock { lockCtx in
            try await self._cleanUpOrphanedBlobs(lockCtx)
        }
    }

    /// Calculate the size of orphaned blobs without deleting them.
    ///
    /// - Returns: The total size in bytes of blobs that are not referenced by any image.
    public func calculateOrphanedBlobsSize() async throws -> UInt64 {
        try await self.lock.withLock { lockCtx in
            try await self._calculateOrphanedBlobsSize(lockCtx)
        }
    }

    @discardableResult
    private func _cleanUpOrphanedBlobs(_ lock: AsyncLock.Context) async throws -> (deleted: [String], freed: UInt64) {
        let images = try await self.list()
        var referenced: [String] = []
        for image in images {
            try await referenced.append(contentsOf: image.referencedDigests().uniqued())
        }
        let (deleted, size) = try await self.contentStore.delete(keeping: referenced)
        return (deleted, size)
    }

    private func _calculateOrphanedBlobsSize(_ lock: AsyncLock.Context) async throws -> UInt64 {
        let images = try await self.list()
        var referenced: [String] = []
        for image in images {
            try await referenced.append(contentsOf: image.referencedDigests().uniqued())
        }

        // Calculate size of blobs not in the referenced list
        let referencedSet = Set(referenced.map { $0.trimmingDigestPrefix })
        let blobsPath = self.path.appendingPathComponent("content/blobs/sha256")

        let fileManager = FileManager.default
        let allBlobs = try fileManager.contentsOfDirectory(
            at: blobsPath,
            includingPropertiesForKeys: [.fileSizeKey],
            options: [.skipsHiddenFiles]
        )

        var orphanedSize: UInt64 = 0
        for blobURL in allBlobs {
            let digest = blobURL.lastPathComponent
            if !referencedSet.contains(digest) {
                if let resourceValues = try? blobURL.resourceValues(forKeys: [.fileSizeKey]),
                    let size = resourceValues.fileSize
                {
                    orphanedSize += UInt64(size)
                }
            }
        }

        return orphanedSize
    }

    /// Tag an existing image such that it can be referenced by another name.
    ///
    /// - Parameters:
    ///   - existing: The reference to an image that already exists in the `ImageStore`.
    ///   - new: The new reference by which the image should also be referenced as.
    /// - Note: The new image created in the `ImageStore` will have the same `Image.Description`
    ///         as that of the image with reference `existing.`
    /// - Returns: A `Containerization.Image` object to the newly created image.
    public func tag(existing: String, new: String) async throws -> Image {
        let old = try await self.get(reference: existing)
        let descriptor = old.descriptor
        do {
            _ = try Reference.parse(new)
        } catch {
            throw ContainerizationError(.invalidArgument, message: "invalid reference \(new), error: \(error)")
        }
        let newDescription = Image.Description(reference: new, descriptor: descriptor)
        return try await self.create(description: newDescription)
    }
}

extension ImageStore {
    /// Pull an image and its associated manifest and blob layers from a remote registry.
    ///
    /// - Parameters:
    ///   - reference: A string that references an image in a remote registry of the form `<host>[:<port>]/repository:<tag>`
    ///                For example: "docker.io/library/alpine:latest".
    ///   - platform: An optional parameter to indicate the platform to be pulled for the image.
    ///               Defaults to `nil` signifying that layers for all supported platforms by the image will be pulled.
    ///   - insecure: A boolean indicating if the connection to the remote registry should be made via plain-text http or not.
    ///               Defaults to false, meaning the connection to the registry will be over https.
    ///   - auth: An object that implements the `Authentication` protocol,
    ///           used to add any credentials to the HTTP requests that are made to the registry.
    ///           Defaults to `nil` meaning no additional credentials are added to any HTTP requests made to the registry.
    ///   - progress: An optional handler over which progress update events about the pull operation can be received.
    ///
    /// - Returns: A `Containerization.Image` object to the newly pulled image.
    public func pull(
        reference: String, platform: Platform? = nil, insecure: Bool = false,
        auth: Authentication? = nil, progress: ProgressHandler? = nil, maxConcurrentDownloads: Int = 3
    ) async throws -> Image {

        let matcher = createPlatformMatcher(for: platform)
        let client = try RegistryClient(reference: reference, insecure: insecure, auth: auth, tlsConfiguration: TLSUtils.makeEnvironmentAwareTLSConfiguration())

        let ref = try Reference.parse(reference)
        let name = ref.path
        guard let tag = ref.tag ?? ref.digest else {
            throw ContainerizationError(.invalidArgument, message: "invalid tag/digest for image reference \(reference)")
        }

        let rootDescriptor = try await client.resolve(name: name, tag: tag)
        let (id, tempDir) = try await self.contentStore.newIngestSession()
        let operation = ImportOperation(
            name: name, contentStore: self.contentStore, client: client, ingestDir: tempDir, progress: progress, maxConcurrentDownloads: maxConcurrentDownloads)
        do {
            let index = try await operation.import(root: rootDescriptor, matcher: matcher)
            return try await self.lock.withLock { lock in
                try await self.contentStore.completeIngestSession(id)
                let description = Image.Description(reference: reference, descriptor: index)
                let image = try await self._create(description: description, lock: lock)
                return image
            }
        } catch {
            try? await self.contentStore.cancelIngestSession(id)
            throw error
        }
    }

    /// Push an image and its associated manifest and blob layers to a remote registry.
    ///
    /// - Parameters:
    ///   - reference: A string that references an image in the `ImageStore`.  It must be of the form `<host>[:<port>]/repository:<tag>`
    ///                For example: "ghcr.io/foo-bar-baz/image:v1".
    ///   - platform: An optional parameter to indicate the platform to be pushed for the image.
    ///               Defaults to `nil` signifying that layers for all supported platforms by the image will be pushed to the remote registry.
    ///   - insecure: A boolean indicating if the connection to the remote registry should be made via plain-text http or not.
    ///               Defaults to false, meaning the connection to the registry will be over https.
    ///   - auth: An object that implements the `Authentication` protocol,
    ///           used to add any credentials to the HTTP requests that are made to the registry.
    ///           Defaults to `nil` meaning no additional credentials are added to any HTTP requests made to the registry.
    ///   - progress: An optional handler over which progress update events about the push operation can be received.
    ///
    public func push(reference: String, platform: Platform? = nil, insecure: Bool = false, auth: Authentication? = nil, progress: ProgressHandler? = nil) async throws {
        let matcher = createPlatformMatcher(for: platform)
        let client = try RegistryClient(reference: reference, insecure: insecure, auth: auth, tlsConfiguration: TLSUtils.makeEnvironmentAwareTLSConfiguration())
        try await self.pushSingle(reference: reference, client: client, matcher: matcher, progress: progress)
    }

    /// Push multiple image references to a remote registry, sharing a single ``RegistryClient``.
    ///
    /// All references must resolve to the same registry host. Passing references that target
    /// different hosts throws a ``ContainerizationError`` with code ``invalidArgument``.
    ///
    /// - Parameters:
    ///   - references: An array of fully qualified image reference strings to push.
    ///                  Each must include a host (e.g., `"ghcr.io/myrepo/myimage:v1"`).
    ///   - platform: An optional parameter to indicate the platform to be pushed for each image.
    ///               Defaults to `nil` signifying that layers for all supported platforms will be pushed.
    ///   - insecure: A boolean indicating if the connection to the remote registry should be made via plain-text http or not.
    ///               Defaults to false, meaning the connection to the registry will be over https.
    ///   - auth: An object that implements the `Authentication` protocol,
    ///           used to add any credentials to the HTTP requests that are made to the registry.
    ///           Defaults to `nil` meaning no additional credentials are added to any HTTP requests made to the registry.
    ///   - maxConcurrentUploads: Maximum number of concurrent tag pushes. Defaults to 3.
    ///   - progress: An optional handler over which progress update events about the push operations can be received.
    ///
    public func push(
        references: [String], platform: Platform? = nil, insecure: Bool = false,
        auth: Authentication? = nil, maxConcurrentUploads: Int = 3, progress: ProgressHandler? = nil
    ) async throws {
        guard let firstReference = references.first else {
            return
        }

        // Parse all references upfront: validate hosts and avoid re-parsing inside tasks.
        let parsed = try references.map { ref in try Reference.parse(ref) }
        let hosts = parsed.compactMap { $0.resolvedDomain }
        guard hosts.count == references.count else {
            throw ContainerizationError(.invalidArgument, message: "all references must include a host")
        }
        let uniqueHosts = Set(hosts)
        guard uniqueHosts.count == 1 else {
            throw ContainerizationError(
                .invalidArgument,
                message: "all references must target the same registry host, got: \(uniqueHosts.sorted().joined(separator: ", "))")
        }

        let matcher = createPlatformMatcher(for: platform)
        let client = try RegistryClient(
            reference: firstReference, insecure: insecure, auth: auth,
            tlsConfiguration: TLSUtils.makeEnvironmentAwareTLSConfiguration())

        let pushOne: @Sendable (String) async -> (String, String?) = { reference in
            do {
                try await self.pushSingle(reference: reference, client: client, matcher: matcher, progress: progress)
                return (reference, nil)
            } catch {
                return (reference, String(describing: error))
            }
        }

        var iterator = references.makeIterator()
        var failures: [(reference: String, message: String)] = []

        await withTaskGroup(of: (String, String?).self) { group in
            for _ in 0..<maxConcurrentUploads {
                guard let reference = iterator.next() else { break }
                group.addTask { await pushOne(reference) }
            }
            for await (ref, error) in group {
                if let error {
                    failures.append((ref, error))
                }
                if let reference = iterator.next() {
                    group.addTask { await pushOne(reference) }
                }
            }
        }

        if !failures.isEmpty {
            let details = failures.map { "\($0.reference): \($0.message)" }.joined(separator: "\n")
            throw ContainerizationError(.internalError, message: "failed to push one or more images:\n\(details)")
        }
    }

    private func pushSingle(
        reference: String, client: ContentClient, matcher: @Sendable (Platform) -> Bool, progress: ProgressHandler?
    ) async throws {
        let allowedMediaTypes = [MediaTypes.dockerManifestList, MediaTypes.index]
        let img = try await self.get(reference: reference)
        guard allowedMediaTypes.contains(img.mediaType) else {
            throw ContainerizationError(.internalError, message: "cannot push image \(reference): unsupported media type \(img.mediaType), expected an index or manifest list")
        }
        let ref = try Reference.parse(reference)
        guard let tag = ref.tag ?? ref.digest else {
            throw ContainerizationError(.invalidArgument, message: "invalid tag/digest for image reference \(reference)")
        }
        let operation = ExportOperation(name: ref.path, tag: tag, contentStore: self.contentStore, client: client, progress: progress)
        try await operation.export(index: img.descriptor, platforms: matcher)
    }
}

extension ImageStore {
    /// Get the image for the init block from the image store.
    /// If the image does not exist locally, pull the image.
    public func getInitImage(reference: String, auth: Authentication? = nil, progress: ProgressHandler? = nil) async throws -> InitImage {
        do {
            let image = try await self.get(reference: reference)
            return InitImage(image: image)
        } catch let error as ContainerizationError {
            if error.code == .notFound {
                let image = try await self.pull(reference: reference, auth: auth, progress: progress)
                return InitImage(image: image)
            }
            throw error
        }
    }
}