import Foundation

// ─────────────────────────────────────────────────────────────
// SentencePiece Tokenizer (tokenizer.model format)
// Simplified implementation for Gemma models
// ─────────────────────────────────────────────────────────────

public final class SentencePieceTokenizer: Tokenizer {
    private let vocab: [String: Int]
    private let reverseVocab: [Int: String]
    private let pieceToId: [String: Int]
    
    public let vocabSize: Int
    public let bosTokenId: Int
    public let eosTokenId: Int
    public let eosTokenIds: Set<Int>
    public let padTokenId: Int
    
    public init(modelPath: String) throws {
        // Load SentencePiece model file
        // Note: This is simplified implementation
        // Full implementation requires protobuf parsing
        
        let data = try Data(contentsOf: URL(fileURLWithPath: modelPath))
        
        // Parse vocab from model (simplified)
        // SentencePiece .model is protobuf format with vocab embedded
        self.vocab = try Self.parseVocabFromModel(data)
        self.reverseVocab = Dictionary(uniqueKeysWithValues: vocab.map { ($1, $0) })
        self.vocabSize = vocab.count
        
        // Special tokens for Gemma
        self.bosTokenId = vocab["<bos>"] ?? vocab["<start_of_turn>"] ?? 2
        self.eosTokenId = vocab["<eos>"] ?? vocab["<end_of_turn>"] ?? 1
        var eosIds = Set<Int>([eosTokenId])
        if let t = vocab["<turn|>"] { eosIds.insert(t) }
        if let t = vocab["<|tool_response>"] { eosIds.insert(t) }
        self.eosTokenIds = eosIds
        self.padTokenId = vocab["<pad>"] ?? 0
        
        self.pieceToId = vocab
    }
    
    public func rawToken(for id: Int) -> String? {
        reverseVocab[id]
    }
    
    public func encode(text: String) -> [Int] {
        var tokens: [Int] = [bosTokenId]
        
        // SentencePiece encoding algorithm (simplified)
        // Full algorithm: find longest matching pieces
        
        var remaining = text
        while !remaining.isEmpty {
            // Find longest matching piece in vocab
            var found = false
            for length in stride(from: min(remaining.count, 20), through: 1, by: -1) {
                let piece = String(remaining.prefix(length))
                
                // Check vocab (with SentencePiece space marker)
                let spPiece = piece.hasPrefix(" ") ? "▁" + piece.dropFirst() : piece
                
                if let tokenId = vocab[spPiece] ?? vocab[piece] {
                    tokens.append(tokenId)
                    remaining = String(remaining.dropFirst(length))
                    found = true
                    break
                }
            }
            
            if !found {
                // Unknown character
                if let unkId = vocab["<unk>"] {
                    tokens.append(unkId)
                    remaining = String(remaining.dropFirst())
                } else {
                    // Skip unknown
                    remaining = String(remaining.dropFirst())
                }
            }
        }
        
        tokens.append(eosTokenId)
        return tokens
    }
    
    public func decode(tokens: [Int]) -> String {
        var text = ""
        
        for tokenId in tokens {
            // Skip special tokens
            if tokenId == bosTokenId || tokenId == eosTokenId || tokenId == padTokenId {
                continue
            }
            
            // Look up piece
            if let piece = reverseVocab[tokenId] {
                // Convert SentencePiece space marker back to space
                let decodedPiece = piece.replacingOccurrences(of: "▁", with: " ")
                text += decodedPiece
            }
        }
        
        return text.trimmingCharacters(in: .whitespaces)
    }
    
    // ─────────────────────────────────────────────────────────────
    // Model Parsing (Simplified)
    // ─────────────────────────────────────────────────────────────
    
    private static func parseVocabFromModel(_ data: Data) throws -> [String: Int] {
        // Simplified vocab parsing
        // Full implementation requires protobuf decoder
        
        // For prototype, try to extract vocab from text representation
        // SentencePiece models sometimes have text vocab embedded
        
        var vocab: [String: Int] = [:]
        
        // Add common Gemma tokens
        vocab["<bos>"] = 2
        vocab["<eos>"] = 1
        vocab["<pad>"] = 0
        vocab["<unk>"] = 3
        
        // Try to parse vocab entries (simplified)
        if let text = String(data: data, encoding: .utf8) {
            let lines = text.split(separator: "\n")
            for line in lines {
                // Parse vocab entries: piece <space> id
                let parts = line.split(separator: "\t")
                if parts.count >= 2 {
                    let piece = String(parts[0])
                    let id = Int(parts[1]) ?? vocab.count
                    vocab[piece] = id
                }
            }
        }
        
        // Fallback: create character-level vocab if parsing failed
        if vocab.count < 100 {
            var idx = vocab.count
            for char in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 ,.!?;:'\"-()[]{}" {
                vocab[String(char)] = idx
                vocab["▁" + String(char)] = idx + 100 // Space marker variant
                idx += 1
            }
        }
        
        return vocab
    }
}