SpeechKit/Sources/SpeechCore/AssetBundle.swift

import Foundation
import CryptoKit

/// Verified model-asset bundle. Owns path resolution, schema-version gating,
/// and SHA-256 integrity checks for every asset listed in the manifest.
///
/// Layout on disk (a directory, conventionally named `ASRModels.bundle`):
///   manifest.json          - schema version, asset table with hashes
///   *.mlmodelc / *.onnx / *.bin / *.data
public struct AssetBundle: Sendable {
    /// Manifest schema versions this SpeechKit build can read.
    public static let supportedSchemas = 2...2

    public struct AssetEntry: Codable, Sendable {
        public let file: String
        public let sha256: String
        public let sizeBytes: Int64
    }

    public struct BundleManifest: Codable, Sendable {
        public let schemaVersion: Int
        public let bundleVersion: String
        public let assets: [String: AssetEntry]   // logical name -> entry
        public let extra: [String: String]?
    }

    public let root: URL
    public let manifest: BundleManifest

    /// Opens and validates a bundle directory.
    /// - Parameters:
    ///   - url: bundle directory URL
    ///   - verifyHashes: full SHA-256 verification of every asset (~1-2 s for
    ///     400 MB). Recommended once per install, not per launch; pass
    ///     `false` after a successful first-run check.
    public init(at url: URL, verifyHashes: Bool = false) throws {
        root = url
        let manifestURL = url.appendingPathComponent("manifest.json")
        guard FileManager.default.fileExists(atPath: manifestURL.path) else {
            throw SpeechError.assetBundleNotFound(url.path)
        }
        do {
            manifest = try JSONDecoder().decode(
                BundleManifest.self, from: Data(contentsOf: manifestURL))
        } catch {
            throw SpeechError.assetMissing("manifest.json (unreadable: \(error.localizedDescription))")
        }
        guard Self.supportedSchemas.contains(manifest.schemaVersion) else {
            throw SpeechError.assetSchemaUnsupported(
                found: manifest.schemaVersion, supported: Self.supportedSchemas)
        }
        for (name, entry) in manifest.assets {
            let assetURL = root.appendingPathComponent(entry.file)
            guard FileManager.default.fileExists(atPath: assetURL.path) else {
                throw SpeechError.assetMissing(name)
            }
            if verifyHashes {
                let actual = try Self.sha256(of: assetURL)
                guard actual == entry.sha256 else {
                    throw SpeechError.assetCorrupted(
                        name: name, expectedSHA256: entry.sha256, actualSHA256: actual)
                }
            }
        }
    }

    /// Resolves a logical asset name (e.g. "tower", "lm_prefill") to a URL.
    public func url(_ logicalName: String) throws -> URL {
        guard let entry = manifest.assets[logicalName] else {
            throw SpeechError.assetMissing(logicalName)
        }
        return root.appendingPathComponent(entry.file)
    }

    /// Convenience: locates `<name>.bundle` inside a host bundle's resources.
    public static func named(_ name: String, in hostBundle: Bundle = .main,
                             verifyHashes: Bool = false) throws -> AssetBundle {
        guard let url = hostBundle.url(forResource: name, withExtension: "bundle") else {
            throw SpeechError.assetBundleNotFound("\(name).bundle in \(hostBundle.bundlePath)")
        }
        return try AssetBundle(at: url, verifyHashes: verifyHashes)
    }

    /// Hash of a file, or of a directory (sorted relative paths + contents,
    /// matching Scripts/make_asset_bundle.py).
    static func sha256(of url: URL) throws -> String {
        var isDir: ObjCBool = false
        FileManager.default.fileExists(atPath: url.path, isDirectory: &isDir)
        var hasher = SHA256()
        if isDir.boolValue {
            let files = (FileManager.default.subpaths(atPath: url.path) ?? [])
                .sorted()
                .map { url.appendingPathComponent($0) }
                .filter { (try? $0.resourceValues(forKeys: [.isRegularFileKey]))?.isRegularFile == true }
            for f in files {
                let rel = f.path.replacingOccurrences(of: url.path + "/", with: "")
                hasher.update(data: Data(rel.utf8))
                try hashFileContents(f, into: &hasher)
            }
        } else {
            try hashFileContents(url, into: &hasher)
        }
        return hasher.finalize().map { String(format: "%02x", $0) }.joined()
    }

    private static func hashFileContents(_ url: URL, into hasher: inout SHA256) throws {
        let handle = try FileHandle(forReadingFrom: url)
        defer { try? handle.close() }
        while autoreleasepool(invoking: {
            let chunk = handle.readData(ofLength: 8 << 20)
            if chunk.isEmpty { return false }
            hasher.update(data: chunk)
            return true
        }) {}
    }
}