#pragma once

#include "gguf.h"

#include <string>
#include <vector>
#include <unordered_map>
#include <map>

namespace qwen3_tts {

// BPE tokenizer configuration
struct tokenizer_config {
    int32_t vocab_size = 151936;
    int32_t pad_token_id = 151643;
    int32_t eos_token_id = 151645;  // <|im_end|>
    int32_t bos_token_id = 151644;  // <|im_start|>
};

// Text tokenizer class (BPE-based, GPT-2 style with Qwen2 pre-tokenization)
class TextTokenizer {
public:
    TextTokenizer();
    ~TextTokenizer();
    
    // Load tokenizer from GGUF file
    bool load_from_gguf(struct gguf_context * ctx);
    
    // Encode text to token IDs
    std::vector<int32_t> encode(const std::string & text) const;
    
    // Encode with TTS text format: <|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n
    std::vector<int32_t> encode_for_tts(const std::string & text) const;

    // Encode reference text for ICL voice cloning:
    // <|im_start|>assistant\n{text}<|im_end|>\n
    std::vector<int32_t> encode_reference_for_tts(const std::string & text) const;
    
    // Encode instruction for TTS (e.g. voice steering)
    std::vector<int32_t> encode_instruct(const std::string & instruct) const;

    // Decode token IDs to text
    std::string decode(const std::vector<int32_t> & tokens) const;
    
    // Decode single token
    std::string decode_token(int32_t token_id) const;
    
    // Get configuration
    const tokenizer_config & get_config() const { return config_; }
    
    // Get error message
    const std::string & get_error() const { return error_msg_; }
    
    // Check if loaded
    bool is_loaded() const { return loaded_; }
    
    // Get special token IDs
    int32_t bos_token_id() const { return config_.bos_token_id; }
    int32_t eos_token_id() const { return config_.eos_token_id; }
    int32_t pad_token_id() const { return config_.pad_token_id; }
    
private:
    tokenizer_config config_;
    std::string error_msg_;
    bool loaded_ = false;
    
    // Vocabulary: token string -> token ID
    std::unordered_map<std::string, int32_t> vocab_;
    
    // Reverse vocabulary: token ID -> token string
    std::vector<std::string> id_to_token_;
    
    // BPE merges: pair -> rank (lower rank = higher priority)
    std::map<std::pair<std::string, std::string>, int32_t> bpe_ranks_;
    
    // Special tokens for "assistant", "user" and newline
    int32_t assistant_token_id_ = 77091;
    int32_t user_token_id_ = 872;
    int32_t newline_token_id_ = 198;  // '\n' encoded
    
    // Helper: convert bytes to unicode (GPT-2 style byte encoding)
    static std::string bytes_to_unicode(const std::string & text);
    static std::string unicode_to_bytes(const std::string & text);
    
    // Helper: get UTF-8 character length
    static size_t utf8_len(char c);
    
    // BPE encoding for a single word
    std::vector<std::string> bpe(const std::string & token) const;
    
    // Find the pair with lowest rank in a sequence
    std::pair<std::string, std::string> get_min_pair(
        const std::vector<std::string> & word) const;
};

} // namespace qwen3_tts