import Foundation

/// The exact fixed-shape tokenizer contract exported with `purpose-lite-v1`.
///
/// Keeping this implementation in Swift avoids a Python/Transformers runtime dependency in the
/// app. The checked-in tokenizer goldens pin Unicode normalization, WordPiece splitting, and the
/// 63-head/62-tail truncation layout byte-for-byte against the training pipeline.
struct PurposeModelInput: Equatable, Sendable {
    let inputIDs: [Int32]
    let attentionMask: [Int32]
    let tokenTypeIDs: [Int32]
}

enum PurposeTokenizerError: Error {
    case emptyVocabulary
}

struct PurposeWordPieceTokenizer: Sendable {
    static let sequenceLength = 128
    static let headTokens = 63
    static let tailTokens = 62

    private static let padID: Int32 = 0
    private static let unknownID: Int32 = 100
    private static let classificationID: Int32 = 101
    private static let separatorID: Int32 = 102

    private let vocabulary: [String: Int32]

    init(vocabularyText: String) throws {
        var vocabulary: [String: Int32] = [:]
        for (index, token) in vocabularyText.split(
            separator: "\n", omittingEmptySubsequences: false
        ).enumerated() {
            var value = String(token)
            if value.last == "\r" { value.removeLast() }
            guard !value.isEmpty else { continue }
            vocabulary[value] = Int32(index)
        }
        guard !vocabulary.isEmpty else { throw PurposeTokenizerError.emptyVocabulary }
        self.vocabulary = vocabulary
    }

    func encode(_ prompt: String) -> PurposeModelInput {
        let content = contentIDs(for: Self.normalizePrompt(prompt))
        let inputIDs: [Int32]
        let tokenTypes: [Int32]

        if content.count <= Self.sequenceLength - 2 {
            inputIDs = [Self.classificationID] + content + [Self.separatorID]
            tokenTypes = Array(repeating: 0, count: inputIDs.count)
        } else {
            let head = Array(content.prefix(Self.headTokens))
            let tail = Array(content.suffix(Self.tailTokens))
            inputIDs = [Self.classificationID] + head + [Self.separatorID]
                + tail + [Self.separatorID]
            tokenTypes = Array(repeating: 0, count: head.count + 2)
                + Array(repeating: 1, count: tail.count + 1)
        }

        let padding = Self.sequenceLength - inputIDs.count
        return PurposeModelInput(
            inputIDs: inputIDs + Array(repeating: Self.padID, count: padding),
            attentionMask: Array(repeating: 1, count: inputIDs.count)
                + Array(repeating: 0, count: padding),
            tokenTypeIDs: tokenTypes + Array(repeating: 0, count: padding))
    }

    /// Matches `purpose_data.normalize_prompt`: Unicode NFKC, Unicode whitespace collapse, trim.
    static func normalizePrompt(_ prompt: String) -> String {
        prompt.precomposedStringWithCompatibilityMapping
            .split(whereSeparator: \Character.isWhitespace)
            .joined(separator: " ")
    }

    private func contentIDs(for text: String) -> [Int32] {
        bertPretokens(text).flatMap(wordPieceIDs(forToken:))
    }

    /// Mirrors Hugging Face's lowercase BERT normalizer and `BertPreTokenizer`.
    private func bertPretokens(_ text: String) -> [String] {
        var cleaned = ""
        for scalar in text.unicodeScalars {
            if scalar.value == 0 || scalar.value == 0xFFFD { continue }
            if scalar.properties.isWhitespace {
                cleaned.unicodeScalars.append(" ")
                continue
            }
            if Self.isControl(scalar) { continue }
            if Self.isChinese(scalar.value) {
                cleaned.unicodeScalars.append(" ")
                cleaned.unicodeScalars.append(scalar)
                cleaned.unicodeScalars.append(" ")
            } else {
                cleaned.unicodeScalars.append(scalar)
            }
        }

        // In the uncased BERT tokenizer, `strip_accents == nil` follows `lowercase == true`.
        let decomposed = cleaned.lowercased().decomposedStringWithCanonicalMapping
        var normalized = ""
        for scalar in decomposed.unicodeScalars
        where scalar.properties.generalCategory != .nonspacingMark {
            normalized.unicodeScalars.append(scalar)
        }

        var result: [String] = []
        var current = ""
        func flush() {
            guard !current.isEmpty else { return }
            result.append(current)
            current.removeAll(keepingCapacity: true)
        }
        for scalar in normalized.unicodeScalars {
            if scalar.properties.isWhitespace {
                flush()
            } else if Self.isPunctuation(scalar) {
                flush()
                result.append(String(scalar))
            } else {
                current.unicodeScalars.append(scalar)
            }
        }
        flush()
        return result
    }

    private func wordPieceIDs(forToken token: String) -> [Int32] {
        let characters = Array(token)
        guard characters.count <= 100 else { return [Self.unknownID] }

        var result: [Int32] = []
        var start = 0
        while start < characters.count {
            var end = characters.count
            var match: Int32?
            while start < end {
                let body = String(characters[start..<end])
                let candidate = start == 0 ? body : "##" + body
                if let identifier = vocabulary[candidate] {
                    match = identifier
                    break
                }
                end -= 1
            }
            guard let match else { return [Self.unknownID] }
            result.append(match)
            start = end
        }
        return result
    }

    private static func isControl(_ scalar: Unicode.Scalar) -> Bool {
        switch scalar.properties.generalCategory {
        case .control, .format: return true
        default: return false
        }
    }

    private static func isPunctuation(_ scalar: Unicode.Scalar) -> Bool {
        let value = scalar.value
        if (33...47).contains(value) || (58...64).contains(value)
            || (91...96).contains(value) || (123...126).contains(value) {
            return true
        }
        switch scalar.properties.generalCategory {
        case .connectorPunctuation, .dashPunctuation, .openPunctuation,
             .closePunctuation, .initialPunctuation, .finalPunctuation,
             .otherPunctuation:
            return true
        default:
            return false
        }
    }

    private static func isChinese(_ value: UInt32) -> Bool {
        (0x4E00...0x9FFF).contains(value)
            || (0x3400...0x4DBF).contains(value)
            || (0x20000...0x2A6DF).contains(value)
            || (0x2A700...0x2B73F).contains(value)
            || (0x2B740...0x2B81F).contains(value)
            || (0x2B820...0x2CEAF).contains(value)
            || (0xF900...0xFAFF).contains(value)
            || (0x2F800...0x2FA1F).contains(value)
    }
}
