import Foundation
import Speech
import Testing
@testable import Notes

struct LanguagePickerTests {
    /// Build an attributed transcript out of `(text, confidence?)` fragments, so a
    /// test can hand-craft the confidence attributes the picker reads.
    private func transcript(_ fragments: [(String, Double?)]) -> AttributedString {
        var result = AttributedString()
        for (text, confidence) in fragments {
            var fragment = AttributedString(text)
            if let confidence { fragment.transcriptionConfidence = confidence }
            result.append(fragment)
        }
        return result
    }

    // MARK: - Confidence extraction

    @Test func meanConfidenceIsCharacterWeighted() {
        // "hello" (5 chars @ 0.9) + " world" (6 chars @ 0.4) → 0.627…
        let attributed = transcript([("hello", 0.9), (" world", 0.4)])
        let mean = LanguagePicker.meanConfidence(of: attributed)
        #expect(abs(mean - ((0.9 * 5 + 0.4 * 6) / 11)) < 0.0001)
    }

    @Test func meanConfidenceIgnoresRunsWithoutConfidence() {
        // Only the confident run counts toward the mean.
        let attributed = transcript([("sure", 1.0), (" maybe", nil)])
        #expect(abs(LanguagePicker.meanConfidence(of: attributed) - 1.0) < 0.0001)
    }

    @Test func meanConfidenceOfEmptyOrUnattributedIsZero() {
        #expect(LanguagePicker.meanConfidence(of: AttributedString("")) == 0)
        #expect(LanguagePicker.meanConfidence(of: transcript([("no attributes", nil)])) == 0)
    }

    // MARK: - Ranking confidence (mixed-engine graceful degradation)

    @Test func hasConfidenceDetectsAttributedRuns() {
        #expect(LanguagePicker.hasConfidence(of: transcript([("sure", 0.9)])))
        #expect(LanguagePicker.hasConfidence(of: transcript([("a", nil), ("b", 0.2)])))
        #expect(!LanguagePicker.hasConfidence(of: transcript([("none", nil)])))
        #expect(!LanguagePicker.hasConfidence(of: AttributedString("")))
    }

    @Test func rankingConfidenceUsesRealMeanWhenPresent() {
        // With any confidence attribute, ranking == the character-weighted mean.
        let attributed = transcript([("hello", 0.9), (" world", 0.4)])
        #expect(abs(LanguagePicker.rankingConfidence(of: attributed)
            - LanguagePicker.meanConfidence(of: attributed)) < 0.0001)
    }

    @Test func rankingConfidenceFallsBackToNeutralForUnattributedText() {
        // A DictationTranscriber transcript with text but no confidence attribute
        // scores the documented neutral value, not 0, so it can still compete.
        let dictation = transcript([("привет мир", nil)])
        #expect(LanguagePicker.rankingConfidence(of: dictation) == LanguagePicker.neutralConfidence)
    }

    @Test func rankingConfidenceOfEmptyIsZero() {
        #expect(LanguagePicker.rankingConfidence(of: AttributedString("")) == 0)
    }

    @Test func neutralDictationBeatsLowConfidenceSpeech() {
        // Audio is Russian: the dictation locale produces good (neutral-scored)
        // text; the mismatched speech locale produces low-confidence garbage.
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "en-US", meanConfidence: 0.2, transcriptLength: 12),
            LanguageCandidate(localeID: "ru-RU",
                              meanConfidence: LanguagePicker.neutralConfidence,
                              transcriptLength: 18),
        ])
        #expect(winner?.localeID == "ru-RU")
    }

    @Test func highConfidenceSpeechBeatsNeutralDictation() {
        // Audio is English: the speech locale's real high confidence outranks the
        // mismatched dictation locale's neutral score.
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "en-US", meanConfidence: 0.85, transcriptLength: 20),
            LanguageCandidate(localeID: "ru-RU",
                              meanConfidence: LanguagePicker.neutralConfidence,
                              transcriptLength: 22),
        ])
        #expect(winner?.localeID == "en-US")
    }

    @Test func twoNeutralDictationCandidatesBreakByLength() {
        // Both engines dictation-only (e.g. ru-RU vs uk-UA), both neutral-scored:
        // the longer transcript wins the tie-break.
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "uk-UA",
                              meanConfidence: LanguagePicker.neutralConfidence,
                              transcriptLength: 8),
            LanguageCandidate(localeID: "ru-RU",
                              meanConfidence: LanguagePicker.neutralConfidence,
                              transcriptLength: 24),
        ])
        #expect(winner?.localeID == "ru-RU")
    }

    // MARK: - Argmax

    @Test func pickChoosesHighestConfidence() {
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "en-US", meanConfidence: 0.42, transcriptLength: 20),
            LanguageCandidate(localeID: "ru-RU", meanConfidence: 0.88, transcriptLength: 18),
            LanguageCandidate(localeID: "fr-FR", meanConfidence: 0.31, transcriptLength: 25),
        ])
        #expect(winner?.localeID == "ru-RU")
    }

    @Test func pickReturnsNilForNoCandidates() {
        #expect(LanguagePicker.pick(from: []) == nil)
    }

    // MARK: - Empty-transcript zeroing

    @Test func emptyTranscriptNeverBeatsRealText() {
        // A zero-length, zero-confidence candidate must lose to any real result.
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "en-US", meanConfidence: 0, transcriptLength: 0),
            LanguageCandidate(localeID: "ru-RU", meanConfidence: 0.05, transcriptLength: 3),
        ])
        #expect(winner?.localeID == "ru-RU")
    }

    // MARK: - Tie-break determinism

    @Test func equalConfidencePrefersLongerTranscript() {
        let winner = LanguagePicker.pick(from: [
            LanguageCandidate(localeID: "en-US", meanConfidence: 0.7, transcriptLength: 10),
            LanguageCandidate(localeID: "ru-RU", meanConfidence: 0.7, transcriptLength: 30),
        ])
        #expect(winner?.localeID == "ru-RU")
    }

    @Test func fullTieBreaksByLocaleID() {
        // Identical confidence and length → smallest localeID wins, deterministically.
        let candidates = [
            LanguageCandidate(localeID: "ru-RU", meanConfidence: 0.5, transcriptLength: 10),
            LanguageCandidate(localeID: "en-US", meanConfidence: 0.5, transcriptLength: 10),
            LanguageCandidate(localeID: "fr-FR", meanConfidence: 0.5, transcriptLength: 10),
        ]
        #expect(LanguagePicker.pick(from: candidates)?.localeID == "en-US")
        // Order of input must not change the result.
        #expect(LanguagePicker.pick(from: candidates.reversed())?.localeID == "en-US")
    }

    @Test func allEmptyCandidatesStillPickDeterministically() {
        let candidates = [
            LanguageCandidate(localeID: "ru-RU", meanConfidence: 0, transcriptLength: 0),
            LanguageCandidate(localeID: "en-US", meanConfidence: 0, transcriptLength: 0),
        ]
        #expect(LanguagePicker.pick(from: candidates)?.localeID == "en-US")
    }
}
