open-standards/wai/wai-sink-swift/Sources/WaiSink/IntDecoder.swift

// The integer decode a native sink runs in its own process (wai/extensions/native-sink.md §13),
// wrapped for Swift: one `WaiIntDecoder` per process finds the device's GPU and checks it once;
// each decode returns the output's bytes and the answer, whose `dispatch` report names the unit
// that ran it (`gpu` or `cpu`) and why any other path was refused.
import Foundation
import WaiFFI

/// One decode's outcome: the bytes, what they are, and the dispatch report.
public struct WaiIntDecoded {
    /// The output: RGB8 rows (a picture, or each frame of a clip in order), int16 PCM
    /// little-endian, or a splat attribute tensor. Empty for a refusal.
    public let bytes: Data
    /// The answer of native-sink §13.3: `capability`, `media`, `refused`, `unit`, `dispatch`,
    /// `micros`.
    public let answer: [String: Any]

    /// The unit that ran the decode: `"gpu"` or `"cpu"`, or nil when nothing ran.
    public var unit: String? { answer["unit"] as? String }
    /// The report `wai-neural-dispatch/1` (neural-dispatch §5); empty for a capability with no
    /// network (`wai.video.int_motion`), which decodes on the CPU.
    public var dispatch: [String: Any] { answer["dispatch"] as? [String: Any] ?? [:] }
    /// The refusal's code, or nil for a decode that gave its output.
    public var refused: String? { (answer["refused"] as? [String: Any])?["code"] as? String }
    /// What the bytes are: `form` (`rgb8`, `rgb8-frames`, `pcm-s16le`, `splat-attributes`) and
    /// its dimensions.
    public var media: [String: Any]? { answer["media"] as? [String: Any] }
    /// For an object of a frame stream: the frame admitted, `kind` (`intra` or `inter`), `index`
    /// and `content_hash`; nil for a refused object and for a whole payload.
    public var frame: [String: Any]? { answer["frame"] as? [String: Any] }
    /// Under a model policy (native-sink §14): `policy`, `set` (`sha256`, `inner`), `at` and
    /// `kept`, the policy the decode was checked against, the set it reads, the decoder's instant
    /// and whether its run was kept; nil without a policy or for a capability that runs no set.
    public var model: [String: Any]? { answer["model"] as? [String: Any] }
}

/// A stream of one integer frame form (`wai.video.int_motion.frame`, `wai.video.int_hyper.frame`,
/// `wai.video.int_cond.frame`; frame-stream §3), one frame an object: it holds the reference,
/// the last frame decoded and the state it left (frame-stream §5), and a copy of its parameter
/// set. Give its objects to `WaiIntDecoder.decode(frame:stream:)` in the order received.
public final class WaiIntFrameStream: @unchecked Sendable {
    fileprivate let handle: OpaquePointer

    /// A stream of `capability` decoded against `set` (empty for `wai.video.int_motion.frame`).
    public init(capability: String, set: [String: Data] = [:], listing: Data? = nil) throws {
        handle = try WaiIntDecoder.withSet(set, listing: listing) { s in
            var out: OpaquePointer?
            let st = capability.withCString { wai_int_frames_open($0, s, &out) }
            guard st == 0, let h = out else { throw WaiSink.error(st) }
            return h
        }
    }

    deinit { wai_int_frames_free(handle) }

    /// Drop the reference: the next frame decoded is an intra frame.
    public func reset() throws {
        let st = wai_int_frames_reset(handle)
        guard st == 0 else { throw WaiSink.error(st) }
    }
}

/// The integer decoder (native-sink §13). Safe to call from any thread; decodes on one device
/// run one at a time.
public final class WaiIntDecoder: @unchecked Sendable {
    /// What the host asks of the GPU.
    public enum GPU: String, Sendable {
        /// Find the device's GPU, run the start-up exactness check, and offer it.
        case auto
        /// Decode on the CPU alone.
        case off
    }

    /// Which unit runs each graph while the GPU is ready (native-sink §13.2).
    public enum UnitRule: String, Sendable {
        /// The unit measured cheaper for the graph's shape on this device: energy where the
        /// platform gives a reading, time otherwise. The default.
        case measured
        /// The GPU for every graph it runs.
        case gpuFirst = "gpu-first"
    }

    private let handle: OpaquePointer

    /// Open the decoder: with `.auto`, the device's GPU is found and checked here, once; `unit`
    /// says which unit runs each graph while it is ready.
    public init(gpu: GPU = .auto, unit: UnitRule = .measured) throws {
        var out: OpaquePointer?
        let st = "{\"gpu\": \"\(gpu.rawValue)\", \"unit\": \"\(unit.rawValue)\"}".withCString { wai_int_decoder_open($0, &out) }
        guard st == 0, let h = out else { throw WaiSink.error(st) }
        handle = h
    }

    deinit { wai_int_decoder_free(handle) }

    /// The GPU: `state` (`ready`, `inexact`, `absent`, `failed`, `off`), the `adapter` as the
    /// backend reports it, the `translator` its kernels went through, and `why` it is not used.
    public func gpu() throws -> [String: Any] {
        var out = wai_buffer_t(data: nil, len: 0, cap: 0)
        let st = wai_int_decoder_info(handle, &out)
        defer { wai_buffer_free(out) }
        guard st == 0 else { throw WaiSink.error(st) }
        return try WaiSink.object(WaiSink.bytes(out))["gpu"] as? [String: Any] ?? [:]
    }

    /// Decode `payload` through `capability` against the parameter set `set` (member name to
    /// bytes; empty for a capability that reads none), unpinned. `intent` and `pinTier` are
    /// the envelope's and a pin's where the host has them (neural-dispatch §2).
    public func decode(capability: String, payload: Data, set: [String: Data] = [:], intent: String? = nil, pinTier: String? = nil) throws -> WaiIntDecoded {
        var opts: [String: Any] = [:]
        if let intent { opts["intent"] = intent }
        if let pinTier { opts["pin_tier"] = pinTier }
        let options = try WaiSink.text(opts)
        return try Self.withSet(set, listing: nil) { s in
            try run { bytes, answer in
                capability.withCString { c in
                    options.withCString { o in
                        payload.withUnsafeBytes { p in
                            wai_int_decode(handle, c, p.bindMemory(to: UInt8.self).baseAddress, UInt(payload.count), s, o, &bytes, &answer)
                        }
                    }
                }
            }
        }
    }

    /// Decode a `WAI1` envelope: its capability, intent and pin from its manifest, the pin
    /// verified against `set` (and `listing`, for a pin in the form `wai.prior.set`) before
    /// anything is decoded.
    public func decode(envelope: Data, set: [String: Data] = [:], listing: Data? = nil) throws -> WaiIntDecoded {
        try Self.withSet(set, listing: listing) { s in
            try run { bytes, answer in
                envelope.withUnsafeBytes { e in
                    wai_int_decode_envelope(handle, e.bindMemory(to: UInt8.self).baseAddress, UInt(envelope.count), s, &bytes, &answer)
                }
            }
        }
    }

    /// Decode the next object of `stream`, a `WAI1` envelope, in the order received
    /// (frame-stream §5): the frame's RGB8 rows, and the answer with `frame`. A refused object
    /// leaves the stream with no reference until an intra frame.
    public func decode(frame envelope: Data, stream: WaiIntFrameStream) throws -> WaiIntDecoded {
        try run { bytes, answer in
            envelope.withUnsafeBytes { e in
                wai_int_decode_frame(handle, stream.handle, e.bindMemory(to: UInt8.self).baseAddress, UInt(envelope.count), &bytes, &answer)
            }
        }
    }

    // MARK: the model policy (native-sink §14)

    /// Hold the operator's model policy: `config` the UTF-8 JSON text `{"policies", "key_log",
    /// "class", "device", "attestation"}`. From then on every neural decode is checked against it
    /// before it runs: one it refuses answers the policy's code (`model_not_in_policy`,
    /// `model_revoked`, `policy_expired`, `policy_not_yet_valid`, `model_tier_below_policy`, …)
    /// with no bytes. A later call replaces it, keeping the runs held under a policy the new one
    /// still holds. Throws for one that does not load, the message starting with
    /// model-governance's code.
    public func holdPolicy(_ config: Data) throws {
        let st = String(decoding: config, as: UTF8.self).withCString { wai_int_decoder_policy(handle, $0) }
        guard st == 0 else { throw WaiSink.error(st) }
    }

    /// The instant the decoder checks and keeps decodes at: a log date-time
    /// (`YYYY-MM-DDTHH:MM:SSZ`), the host's trusted clock, or nil for the system's clock.
    public func setClock(_ at: String?) throws {
        let st: Int32
        if let at {
            st = at.withCString { wai_int_decoder_clock(handle, $0) }
        } else {
            st = wai_int_decoder_clock(handle, nil)
        }
        guard st == 0 else { throw WaiSink.error(st) }
    }

    /// What the decoder holds under its policy: `policy` (`class`, `device`, `operator`,
    /// `policies`) and `runs` (`held`, `reported_to`, `decodes`), both null without a policy.
    public func model() throws -> [String: Any] {
        var out = wai_buffer_t(data: nil, len: 0, cap: 0)
        let st = wai_int_decoder_model(handle, &out)
        defer { wai_buffer_free(out) }
        guard st == 0 else { throw WaiSink.error(st) }
        return try WaiSink.object(WaiSink.bytes(out))
    }

    /// Sign the run report (model-governance §6.1) for the period `from` to `to`, two log
    /// date-times, with the device key whose 32-byte seed is `deviceSeed`, at `issuedAt` (the
    /// decoder's instant where nil): the signed `wai-model-run` record's JSON text, whose runs then
    /// leave the decoder. Throws for a period it cannot sign, the message starting with the
    /// policy's code (`record_overlap`, `outside_period`, `policy_expired`, `bad_key`, …).
    public func runReport(from: String, to: String, issuedAt: String? = nil, deviceSeed: Data) throws -> Data {
        var req: [String: Any] = ["from": from, "to": to]
        if let issuedAt { req["issued_at"] = issuedAt }
        let request = try WaiSink.text(req)
        var out = wai_buffer_t(data: nil, len: 0, cap: 0)
        let st = request.withCString { r in
            deviceSeed.withUnsafeBytes { s in
                wai_int_run_report(handle, r, s.bindMemory(to: UInt8.self).baseAddress, UInt(deviceSeed.count), &out)
            }
        }
        defer { wai_buffer_free(out) }
        guard st == 0 else { throw WaiSink.error(st) }
        return WaiSink.bytes(out)
    }

    /// A model-governance bundle (`*.wai-models.json`, model-governance §8) verified as `wai model
    /// verify` verifies it: `outcome` `accept` with the statement, or `refuse` with `reason` and
    /// `at`.
    public static func verifyModelBundle(_ bundle: Data) throws -> [String: Any] {
        var out = wai_buffer_t(data: nil, len: 0, cap: 0)
        let st = bundle.withUnsafeBytes { b in wai_model_verify(b.bindMemory(to: UInt8.self).baseAddress, UInt(bundle.count), &out) }
        defer { wai_buffer_free(out) }
        guard st == 0 else { throw WaiSink.error(st) }
        return try WaiSink.object(WaiSink.bytes(out))
    }

    // MARK: plumbing

    fileprivate static func withSet<T>(_ files: [String: Data], listing: Data?, _ f: (OpaquePointer?) throws -> T) throws -> T {
        var s: OpaquePointer?
        let st = wai_int_set_new(&s)
        guard st == 0, let set = s else { throw WaiSink.error(st) }
        defer { wai_int_set_free(set) }
        for (name, bytes) in files {
            let st = name.withCString { n in
                bytes.withUnsafeBytes { b in wai_int_set_add(set, n, b.bindMemory(to: UInt8.self).baseAddress, UInt(bytes.count)) }
            }
            guard st == 0 else { throw WaiSink.error(st) }
        }
        if let listing {
            let st = listing.withUnsafeBytes { b in wai_int_set_listing(set, b.bindMemory(to: UInt8.self).baseAddress, UInt(listing.count)) }
            guard st == 0 else { throw WaiSink.error(st) }
        }
        return try f(set)
    }

    private func run(_ f: (inout wai_buffer_t, inout wai_buffer_t) -> Int32) throws -> WaiIntDecoded {
        var bytes = wai_buffer_t(data: nil, len: 0, cap: 0)
        var answer = wai_buffer_t(data: nil, len: 0, cap: 0)
        let st = f(&bytes, &answer)
        defer {
            wai_buffer_free(bytes)
            wai_buffer_free(answer)
        }
        guard st == 0 else { throw WaiSink.error(st) }
        return WaiIntDecoded(bytes: WaiSink.bytes(bytes), answer: try WaiSink.object(WaiSink.bytes(answer)))
    }
}