import Testing
@testable import SstWhisper

struct PromptTokensTests {
    /// Fake tokenizer: one token (id 7) per word, plus one special token (id 2000) that must be filtered out.
    let encode: (String) -> [Int] = { text in
        text.split(whereSeparator: { !$0.isLetter }).map { _ in 7 } + [2000]
    }

    @Test func allTermsFitAndSpecialTokensAreFiltered() {
        let tokens = WhisperKitTranscriber.promptTokens(
            vocabulary: ["Anthropic", "Vercel", "Supabase"], maxTokens: 64, specialTokenBegin: 1000, encode: encode)
        #expect(tokens == [7, 7, 7])
    }

    @Test func termsThatDoNotFitAreDroppedFromTheEnd() {
        let vocab = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
        let tokens = WhisperKitTranscriber.promptTokens(
            vocabulary: vocab, maxTokens: 4, specialTokenBegin: 1000, encode: encode)
        #expect(tokens?.count == 4)
    }

    @Test func emptyVocabularyGivesNoPrompt() {
        #expect(WhisperKitTranscriber.promptTokens(
            vocabulary: [], maxTokens: 64, specialTokenBegin: 1000, encode: encode) == nil)
    }

    @Test func firstTermAloneTooLongGivesNoPrompt() {
        #expect(WhisperKitTranscriber.promptTokens(
            vocabulary: ["one two three"], maxTokens: 2, specialTokenBegin: 1000, encode: encode) == nil)
    }
}