import Foundation

/// Old→new token-id remap for the EN+ZH pruned models.
///
/// The `*-en_zh-*` Core ML models were built on a sliced word-embedding matrix
/// (only ASCII/CJK tokens + specials kept). The app must therefore translate the
/// full-vocabulary XLM-R ids produced by the tokenizer into the pruned id space
/// before feeding the model; dropped tokens (`-1`) become `<unk>`.
///
/// Special ids 0/1/2/3 (`<s>`/`<pad>`/`</s>`/`<unk>`) map to themselves, so the
/// same table can be applied uniformly to every id in a window, including the
/// padding and sentinel tokens.
struct VocabRemap {
    private let table: [Int32]   // table[oldId] = newId, or -1 if dropped
    let unkNewId: Int32

    /// Layout (little-endian): magic "WTPR", u32 count, then `count` × i32.
    init(url: URL) throws {
        let data = try Data(contentsOf: url, options: .mappedIfSafe)
        var table = [Int32]()
        try data.withUnsafeBytes { (raw: UnsafeRawBufferPointer) in
            guard raw.count >= 8 else { throw KitError.corruptResource("en_zh_remap.bin too small") }
            let magic = raw.loadUnaligned(fromByteOffset: 0, as: UInt32.self)
            guard magic == 0x52505457 else { throw KitError.corruptResource("bad remap magic") } // "WTPR"
            let count = Int(UInt32(littleEndian: raw.loadUnaligned(fromByteOffset: 4, as: UInt32.self)))
            guard raw.count >= 8 + count * 4 else { throw KitError.corruptResource("remap truncated") }
            table = [Int32](repeating: 0, count: count)
            for k in 0..<count {
                let v = raw.loadUnaligned(fromByteOffset: 8 + k * 4, as: Int32.self)
                table[k] = Int32(littleEndian: v)
            }
        }
        self.table = table
        // <unk> old id is 3; its mapped value is the pruned unk id.
        self.unkNewId = (3 < table.count && table[3] >= 0) ? table[3] : 3
    }

    @inline(__always)
    func map(_ oldId: Int32) -> Int32 {
        guard oldId >= 0 && Int(oldId) < table.count else { return unkNewId }
        let v = table[Int(oldId)]
        return v >= 0 ? v : unkNewId
    }
}