import CoreML
import Foundation
import SpeechCore

/// The primary entry point for on-device speech recognition.
///
/// ```swift
/// let transcriber = try ASRTranscriber()          // ASRModels.bundle in main bundle
/// let text = try transcriber.transcribe(samples)  // 16 kHz mono Float32
/// ```
///
/// ## Threading
/// All public methods are thread-safe. Inference is serialized on an internal
/// queue: concurrent `transcribe` calls run one at a time in call order.
/// Synchronous methods must not be called from the main thread (they block);
/// use the `async` variants from UI code.
///
/// ## Lifecycle
/// Initialization compiles/loads Core ML models on first use per install
/// (tens of seconds on-device); subsequent launches hit the system cache.
/// Call `warmUp()` after init to move that cost off the first transcription.
public final class ASRTranscriber {
    public struct Options: Sendable {
        /// Cap on generated text tokens (clamped internally to decoder budget).
        public var maxNewTokens = 128
        /// Compute units for the audio tower. `.cpuAndNeuralEngine` is the
        /// validated production path.
        public var computeUnits: MLComputeUnits = .cpuAndNeuralEngine
        /// Verify SHA-256 of every asset at load (~1-2 s). Enable for the
        /// first launch after install/update.
        public var verifyAssets = false
        public init() {}
    }

    public struct Timings: Sendable {
        public var mel: Double = 0
        public var tower: Double = 0
        public var decode: Double = 0
        public var total: Double = 0
    }

    public struct Result: Sendable {
        public let text: String
        public let tokenIDs: [Int]
        public let hitStop: Bool
        public let timings: Timings
    }

    /// Cancellation handle for an in-flight transcription.
    public final class CancellationToken: @unchecked Sendable {
        private let lock = NSLock()
        private var _cancelled = false
        public var isCancelled: Bool {
            lock.lock(); defer { lock.unlock() }
            return _cancelled
        }
        public func cancel() {
            lock.lock(); _cancelled = true; lock.unlock()
        }
        public init() {}
    }

    package let store: AssetStore
    package let mel: MelExtractor
    package let tower: AudioTower
    package let decoder: LMDecoder
    private let queue = DispatchQueue(label: "com.speechkit.asr.inference", qos: .userInitiated)

    /// Minimum audio length accepted (seconds).
    public static let minimumAudioSeconds = 0.25

    // MARK: - Init

    /// Opens `ASRModels.bundle` from the host bundle's resources.
    public convenience init(options: Options = Options()) throws {
        let bundle = try AssetBundle.named("ASRModels", verifyHashes: options.verifyAssets)
        try self.init(assetBundle: bundle, options: options)
    }

    /// Opens a bundle at an explicit directory URL (e.g. downloaded models).
    public convenience init(bundleURL: URL, options: Options = Options()) throws {
        let bundle = try AssetBundle(at: bundleURL, verifyHashes: options.verifyAssets)
        try self.init(assetBundle: bundle, options: options)
    }

    public init(assetBundle: AssetBundle, options: Options = Options()) throws {
        self.options = options
        store = try AssetStore(bundle: assetBundle)
        mel = MelExtractor(hannWindow: store.hannWindow, melFilters: store.melFilters)
        tower = try AudioTower(unifiedModelURL: assetBundle.url("tower"),
                               store: store, computeUnits: options.computeUnits)
        decoder = try LMDecoder(prefillURL: assetBundle.url("lm_prefill"),
                                decodeURL: assetBundle.url("lm_decode"),
                                maskGenURL: assetBundle.url("mask_gen"),
                                store: store)
    }

    private let options: Options

    /// Pre-loads ANE programs so the first transcription is fast.
    /// Synchronous; call off the main thread.
    public func warmUp() {
        queue.sync { tower.warmUp() }
    }

    // MARK: - Transcription

    /// Transcribes 16 kHz mono Float32 samples. Blocks until complete.
    public func transcribe(_ samples: [Float],
                           cancellation: CancellationToken? = nil) throws -> Result {
        try queue.sync {
            try autoreleasepool {
                try transcribeLocked(samples, cancellation: cancellation)
            }
        }
    }

    /// Async variant for Swift Concurrency callers.
    public func transcribe(_ samples: [Float],
                           cancellation: CancellationToken? = nil) async throws -> Result {
        try await withCheckedThrowingContinuation { cont in
            queue.async {
                do {
                    let r = try autoreleasepool {
                        try self.transcribeLocked(samples, cancellation: cancellation)
                    }
                    cont.resume(returning: r)
                } catch {
                    cont.resume(throwing: error)
                }
            }
        }
    }

    /// Transcribes an audio file (any format/rate AVFoundation can read).
    public func transcribeFile(_ url: URL,
                               cancellation: CancellationToken? = nil) throws -> Result {
        let samples = try AudioResampler.loadFile(url)
        return try transcribe(samples, cancellation: cancellation)
    }

    // MARK: - Internal

    private func transcribeLocked(_ samples: [Float],
                                  cancellation: CancellationToken?) throws -> Result {
        guard Double(samples.count) / 16_000 >= Self.minimumAudioSeconds else {
            throw SpeechError.audioTooShort(minimumSeconds: Self.minimumAudioSeconds)
        }
        func checkCancel() throws {
            if cancellation?.isCancelled == true { throw SpeechError.cancelled }
        }
        var t = Timings()
        let tAll = Date()

        var t0 = Date()
        let (melFeature, encLen, bucketFrames) = mel.extract(samples)
        t.mel = -t0.timeIntervalSinceNow
        try checkCancel()

        t0 = Date()
        let masks: (attn: [Float], valid: [Bool])
        let audioEmbeds: [[Float]]
        do {
            masks = try decoder.generateMasks(encLen: encLen)
            audioEmbeds = try tower.embed(
                mel: melFeature, bucketFrames: bucketFrames,
                attnMask390: masks.attn, validMask390: masks.valid,
                sampleCount: min(samples.count, MelExtractor.maxSamples))
        } catch let e as SpeechError {
            throw e
        } catch {
            throw SpeechError.inferenceFailed(stage: "audio-tower", underlying: error.localizedDescription)
        }
        t.tower = -t0.timeIntervalSinceNow
        try checkCancel()

        let p = store.manifest.prompt_ids
        var promptIDs: [Int] = [p.user, p.bos_audio]
        promptIDs.append(contentsOf: Array(repeating: p.audio, count: audioEmbeds.count))
        promptIDs.append(p.eos_audio)
        promptIDs.append(contentsOf: p.text)
        promptIDs.append(p.assistant)

        var embeds: [[Float]] = []
        var cursor = 0
        for id in promptIDs {
            if id == p.audio {
                embeds.append(audioEmbeds[cursor]); cursor += 1
            } else {
                embeds.append(store.embedding(for: id))
            }
        }

        t0 = Date()
        let generation: LMDecoder.GenerationResult
        do {
            generation = try decoder.generate(
                promptEmbeds: embeds, promptIDs: promptIDs,
                maxNewTokens: options.maxNewTokens,
                isCancelled: { cancellation?.isCancelled == true })
        } catch let e as SpeechError {
            throw e
        } catch {
            throw SpeechError.inferenceFailed(stage: "decoder", underlying: error.localizedDescription)
        }
        t.decode = -t0.timeIntervalSinceNow
        t.total = -tAll.timeIntervalSinceNow
        return Result(text: generation.text,
                      tokenIDs: generation.tokenIDs,
                      hitStop: generation.hitStop,
                      timings: t)
    }
}