import Foundation
import os

struct TranscriptionOutputFilter {
    private static let logger = Logger(subsystem: "com.noaber.proat", category: "TranscriptionOutputFilter")
    private static let removePunctuationKey = "RemovePunctuation"
    private static let lowercaseTranscriptionKey = "LowercaseTranscription"
    private static let apostropheLikeCharacters = CharacterSet(charactersIn: "'’‘ʼ＇")

    private static let hallucinationPatterns = [
        #"\[.*?\]"#,            // []
        #"\(.*?\)"#,            // ()
        #"\{.*?\}"#,            // {}
        #"<\|[^|>]*\|>"#        // Whisper special tokens: <|nl|>, <|notimestamps|>, <|0.00|>
    ]

    static func filter(_ text: String) -> String {
        var filteredText = text

        // Remove <TAG>...</TAG> blocks
        let tagBlockPattern = #"<([A-Za-z][A-Za-z0-9:_-]*)[^>]*>[\s\S]*?</\1>"#
        if let regex = try? NSRegularExpression(pattern: tagBlockPattern) {
            let range = NSRange(filteredText.startIndex..., in: filteredText)
            filteredText = regex.stringByReplacingMatches(in: filteredText, options: [], range: range, withTemplate: "")
        }

        // Remove bracketed hallucinations
        for pattern in hallucinationPatterns {
            if let regex = try? NSRegularExpression(pattern: pattern) {
                let range = NSRange(filteredText.startIndex..., in: filteredText)
                filteredText = regex.stringByReplacingMatches(in: filteredText, options: [], range: range, withTemplate: "")
            }
        }

        // Remove filler words (if enabled)
        if FillerWordManager.shared.isEnabled {
            for fillerWord in FillerWordManager.shared.fillerWords {
                let pattern = "\\b\(NSRegularExpression.escapedPattern(for: fillerWord))\\b[,.]?"
                if let regex = try? NSRegularExpression(pattern: pattern, options: .caseInsensitive) {
                    let range = NSRange(filteredText.startIndex..., in: filteredText)
                    filteredText = regex.stringByReplacingMatches(in: filteredText, options: [], range: range, withTemplate: "")
                }
            }
        }

        // Collapse repetition loops (hallucination defense — upstream PR #565)
        filteredText = collapseRepetitionLoops(filteredText)

        // Clean whitespace
        filteredText = filteredText.replacingOccurrences(of: #"\s{2,}"#, with: " ", options: .regularExpression)
        filteredText = filteredText.trimmingCharacters(in: .whitespacesAndNewlines)

        // Laag 1 taal-gate: Voxtral's language=nl-hint is geen harde forcering.
        // Bij korte/ambigue audio levert het model soms een compleet anderstalig
        // resultaat in een niet-Latijns schrift (bv. Cyrillisch "Снапки варить
        // буду"). Nederlands gebruikt uitsluitend Latijns schrift, dus zulke
        // output is met zekerheid fout → afkeuren (leeg = niets plakken).
        if isPredominantlyNonLatinScript(filteredText) {
            logger.notice("🚫 Taal-gate: niet-Latijns schrift afgekeurd")
            ProatDebugLog.write("TAAL-GATE | niet-Latijns schrift afgekeurd | \(String(filteredText.prefix(80)))")
            return ""
        }

        return filteredText
    }

    // MARK: - Laag 1 taal-gate (niet-Latijns schrift)

    /// Boven dit aandeel niet-Latijnse letters wordt de output afgekeurd.
    /// Bewust hoog (0.5) zodat een enkel los symbool (α, β) in een verder
    /// Nederlandse zin nooit een goede transcriptie weggooit — alleen output
    /// die overwégend uit een ander schrift bestaat sneuvelt.
    private static let nonLatinRejectThreshold = 0.5

    /// True als de tekst overwegend uit niet-Latijnse létters bestaat.
    /// Telt uitsluitend alfabetische karakters: cijfers, leestekens, spaties en
    /// emoji tellen niet mee (die zijn schrift-neutraal). Geeft false bij tekst
    /// zonder letters, zodat de bestaande filters daarover beslissen.
    static func isPredominantlyNonLatinScript(_ text: String) -> Bool {
        var latinLetters = 0
        var nonLatinLetters = 0
        for scalar in text.unicodeScalars where scalar.properties.isAlphabetic {
            if isLatinLetter(scalar) {
                latinLetters += 1
            } else {
                nonLatinLetters += 1
            }
        }
        let totalLetters = latinLetters + nonLatinLetters
        guard totalLetters > 0 else { return false }
        return Double(nonLatinLetters) / Double(totalLetters) >= nonLatinRejectThreshold
    }

    /// True voor letters uit het Latijnse schrift, inclusief alle Nederlandse
    /// diakrieten (é ë ï ö ü ç …). Basic Latin t/m Latin Extended-B (0041–024F)
    /// plus Latin Extended Additional (1E00–1EFF). Alfabetische scalars daarbuiten
    /// zijn een ander schrift: Grieks, Cyrillisch, Hebreeuws, Arabisch, CJK,
    /// Hangul, Thai, Devanagari, enzovoort.
    private static func isLatinLetter(_ scalar: Unicode.Scalar) -> Bool {
        let v = scalar.value
        return (0x0041...0x024F).contains(v) || (0x1E00...0x1EFF).contains(v)
    }

    // MARK: - Repetition Loop Collapse (from upstream PR #565 by cmbays)

    /// Collapses hallucination loops where whisper repeats the same phrase many times.
    /// Two passes: sentence-level dedup, then n-gram run collapse for unpunctuated output.
    private static func collapseRepetitionLoops(_ text: String) -> String {
        let words = text.split(separator: " ")
        guard words.count >= 10 else { return text }

        var result = text
        result = collapseSentenceRuns(result)
        result = collapseNgramRuns(result)

        if result != text {
            logger.notice("🔄 Collapsed repetition loop in transcription output")
        }
        return result
    }

    /// Collapse runs of 3+ consecutive identical sentences.
    private static func collapseSentenceRuns(_ text: String) -> String {
        let pattern = #"(?<=[.!?])\s+"#
        guard let regex = try? NSRegularExpression(pattern: pattern) else { return text }
        let range = NSRange(text.startIndex..., in: text)
        let parts = regex.matches(in: text, range: range)
        guard !parts.isEmpty else { return text }

        var sentences: [String] = []
        var lastEnd = text.startIndex
        for match in parts {
            guard let matchRange = Range(match.range, in: text) else { continue }
            sentences.append(String(text[lastEnd..<matchRange.lowerBound]))
            lastEnd = matchRange.upperBound
        }
        let remainder = String(text[lastEnd...])
        if !remainder.trimmingCharacters(in: .whitespaces).isEmpty {
            sentences.append(remainder)
        }
        guard sentences.count >= 3 else { return text }

        var collapsed: [String] = []
        var i = 0
        while i < sentences.count {
            let normalized = normalize(sentences[i])
            var runLength = 1
            while i + runLength < sentences.count && normalize(sentences[i + runLength]) == normalized {
                runLength += 1
            }
            collapsed.append(sentences[i])
            if runLength >= 3 {
                logger.notice("🔄 Collapsed \(runLength, privacy: .public) repeated sentences")
            }
            i += runLength
        }
        return collapsed.joined(separator: " ")
    }

    /// Collapse consecutive repeated n-grams (5-12 words) for unpunctuated loops.
    private static func collapseNgramRuns(_ text: String) -> String {
        var words = text.split(separator: " ").map(String.init)
        guard words.count >= 10 else { return text }

        let minN = 5
        let maxN = min(12, words.count / 2)
        guard minN <= maxN else { return text }

        for n in stride(from: maxN, through: minN, by: -1) {
            let minReps = n >= 8 ? 2 : 3
            var i = 0
            var collapsed = false
            var newWords: [String] = []

            while i <= words.count - n {
                let gram = words[i..<(i + n)].map { $0.lowercased() }
                var reps = 1
                var j = i + n
                while j + n <= words.count {
                    let candidate = words[j..<(j + n)].map { $0.lowercased() }
                    if candidate == gram {
                        reps += 1
                        j += n
                    } else {
                        break
                    }
                }
                if reps >= minReps {
                    newWords.append(contentsOf: words[i..<(i + n)])
                    logger.notice("🔄 Collapsed \(reps, privacy: .public)x repeated \(n, privacy: .public)-gram")
                    i = j
                    collapsed = true
                } else {
                    newWords.append(words[i])
                    i += 1
                }
            }
            if i < words.count {
                newWords.append(contentsOf: words[i...])
            }
            if collapsed {
                words = newWords
            }
        }
        return words.joined(separator: " ")
    }

    /// Normalize a string for comparison: lowercase, strip punctuation and extra whitespace.
    private static func normalize(_ s: String) -> String {
        s.lowercased()
            .replacingOccurrences(of: #"[^\w\s]"#, with: "", options: .regularExpression)
            .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression)
            .trimmingCharacters(in: .whitespaces)
    }

    static func applyUserCleanupPreferences(_ text: String) -> String {
        let shouldRemovePunctuation = UserDefaults.standard.bool(forKey: removePunctuationKey)
        let shouldLowercase = UserDefaults.standard.bool(forKey: lowercaseTranscriptionKey)

        guard shouldRemovePunctuation || shouldLowercase else {
            return text
        }

        var cleanedText = text
        if shouldRemovePunctuation {
            cleanedText = removePunctuation(from: cleanedText)
        }
        if shouldLowercase {
            cleanedText = cleanedText.lowercased()
        }

        return cleanedText
    }

    static func removePunctuation(from text: String) -> String {
        guard !text.isEmpty else { return text }

        let punctuationSeparators = CharacterSet.punctuationCharacters.subtracting(apostropheLikeCharacters)
        let cleanedScalars = text.unicodeScalars.map { scalar -> String in
            if apostropheLikeCharacters.contains(scalar) {
                return ""
            }

            if punctuationSeparators.contains(scalar) {
                return " "
            }

            return String(scalar)
        }

        return normalizeWhitespace(cleanedScalars.joined())
    }

    private static func normalizeWhitespace(_ text: String) -> String {
        text
            .replacingOccurrences(of: #"[^\S\r\n]{2,}"#, with: " ", options: .regularExpression)
            .replacingOccurrences(of: #"[ \t]+\n"#, with: "\n", options: .regularExpression)
            .replacingOccurrences(of: #"\n[ \t]+"#, with: "\n", options: .regularExpression)
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }
} 
