import Foundation
import SstCore
import WhisperKit

public final class WhisperKitTranscriber: Transcriber {
    public enum TranscriberError: Error, LocalizedError {
        case notLoaded

        public var errorDescription: String? { "Spraakmodel is nog niet geladen." }
    }

    private let modelName: String
    private let storage: ModelStorage
    private var pipeline: WhisperKit?

    /// Download progress (0...1) during the one-time model download. May be called from any thread.
    public var onProgress: ((Double) -> Void)?

    public init(model: String = ModelCatalog.defaultVariant, storage: ModelStorage = ModelStorage()) {
        self.modelName = model
        self.storage = storage
    }

    public var variant: String { modelName }

    private func loadPipeline(from folder: URL) async throws -> WhisperKit {
        try await WhisperKit(WhisperKitConfig(
            downloadBase: storage.base,
            modelFolder: folder.path,
            tokenizerFolder: storage.base,
            prewarm: true,
            load: true,
            download: false
        ))
    }

    /// Loads the model into memory. Downloads it (≈1.5 GB) only if it is not stored locally yet;
    /// once downloaded, launching needs no network.
    public func load() async throws {
        guard pipeline == nil else { return }
        try FileManager.default.createDirectory(at: storage.base, withIntermediateDirectories: true)

        if let folder = storage.folder(for: modelName),
           let loaded = try? await loadPipeline(from: folder) {
            pipeline = loaded
            return
        }
        let folder = try await WhisperKit.download(
            variant: modelName,
            downloadBase: storage.base,
            from: ModelStorage.repo
        ) { [onProgress] progress in
            onProgress?(progress.fractionCompleted)
        }
        pipeline = try await loadPipeline(from: folder)
    }

    public func transcribe(_ samples: [Float]) async throws -> String {
        try await transcribe(samples, vocabulary: [])
    }

    /// Terms are added one by one while the encoded prompt stays within `maxTokens`.
    /// Returns nil when nothing fits or the list is empty.
    static func promptTokens(
        vocabulary: [String],
        maxTokens: Int = 64,
        specialTokenBegin: Int,
        encode: (String) -> [Int]
    ) -> [Int]? {
        var kept: [String] = []
        var tokens: [Int] = []
        for term in vocabulary {
            let candidate = kept + [term]
            let encoded = encode(" " + candidate.joined(separator: ", ")).filter { $0 < specialTokenBegin }
            if encoded.count > maxTokens { break }
            kept = candidate
            tokens = encoded
        }
        return tokens.isEmpty ? nil : tokens
    }

    public func transcribe(_ samples: [Float], vocabulary: [String]) async throws -> String {
        guard let pipeline else { throw TranscriberError.notLoaded }
        // language nil + detectLanguage: Whisper picks NL or EN per utterance.
        var options = DecodingOptions(
            task: .transcribe,
            language: nil,
            detectLanguage: true,
            skipSpecialTokens: true,
            withoutTimestamps: true
        )
        if !vocabulary.isEmpty, let tokenizer = pipeline.tokenizer {
            options.promptTokens = Self.promptTokens(
                vocabulary: vocabulary,
                specialTokenBegin: tokenizer.specialTokens.specialTokenBegin,
                encode: { tokenizer.encode(text: $0) }
            )
        }
        let results = try await pipeline.transcribe(audioArray: samples, decodeOptions: options)
        return results
            .map(\.text)
            .joined(separator: " ")
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }
}