import Foundation
import SpeechCore

/// Loads decoded model assets (mel filters, projector weights, embedding
/// table, vocab decode table) from a verified AssetBundle.
/// Internal to the SDK; integrators only see ASRTranscriber.
package final class AssetStore {
    package struct Manifest: Codable {
        package struct PromptIDs: Codable {
            package let user: Int
            package let bos_audio: Int
            package let audio: Int
            package let eos_audio: Int
            package let text: [Int]
            package let assistant: Int
        }
        package struct Tokens: Codable {
            package let audio_token_id: Int
            package let pad_token_id: Int
            package let eos_token_ids: [Int]
            package let extra_block_token_ids: [Int]
        }
        package let prompt_ids: PromptIDs
        package let tokens: Tokens
    }

    package let manifest: Manifest
    package let hannWindow: [Float]
    package let melFilters: [Float]
    package let projNormW: [Float]
    package let projNormB: [Float]
    package let projLinW: [Float]   // 512*1024 row-major [out][in]
    package let projLinB: [Float]

    private let embeddingData: Data          // fp16, vocab*512
    private let vocabBytes: Data
    private let vocabOffsets: [UInt32]
    package let embeddingDim = 512

    package init(bundle: AssetBundle) throws {
        func data(_ name: String) throws -> Data {
            try Data(contentsOf: bundle.url(name), options: .mappedIfSafe)
        }
        func floats(_ name: String) throws -> [Float] {
            try data(name).withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }
        }
        do {
            manifest = try JSONDecoder().decode(Manifest.self, from: data("asr_manifest"))
        } catch let e as SpeechError {
            throw e
        } catch {
            throw SpeechError.assetMissing("asr_manifest (decode: \(error.localizedDescription))")
        }
        hannWindow = try floats("hann_window")
        melFilters = try floats("mel_filters")
        projNormW = try floats("projector_norm_w")
        projNormB = try floats("projector_norm_b")
        projLinW = try floats("projector_lin_w")
        projLinB = try floats("projector_lin_b")
        embeddingData = try data("token_embedding")
        vocabBytes = try data("vocab_bytes")
        vocabOffsets = try data("vocab_offsets").withUnsafeBytes {
            Array($0.bindMemory(to: UInt32.self))
        }
    }

    /// fp16 -> fp32 embedding row lookup
    package func embedding(for tokenID: Int) -> [Float] {
        let rowBytes = embeddingDim * MemoryLayout<Float16>.size
        let start = tokenID * rowBytes
        return embeddingData.subdata(in: start..<(start + rowBytes)).withUnsafeBytes {
            $0.bindMemory(to: Float16.self).map { Float($0) }
        }
    }

    package func decode(_ tokenIDs: [Int]) -> String {
        var bytes = Data()
        for id in tokenIDs where id >= 0 && id + 1 < vocabOffsets.count {
            let a = Int(vocabOffsets[id]), b = Int(vocabOffsets[id + 1])
            bytes.append(vocabBytes.subdata(in: a..<b))
        }
        let raw = String(decoding: bytes, as: UTF8.self)
        return raw.trimmingCharacters(in: .whitespacesAndNewlines)
            .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression)
    }
}