# -*- coding: utf-8 -*-
from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
import os

def train_tokenizer(iterator, vocab_size=50000, min_frequency=3, output_dir="tokenizer_model"):
    """
    Εκπαιδεύει έναν Tokenizer τύπου BPE και αποθηκεύει το αποτέλεσμα.
    """
    print("🔄 Ξεκινάει η εκπαίδευση του tokenizer...")

    # Δημιουργία του Tokenizer με μοντέλο BPE και ορισμό token για άγνωστα
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    
    # Normalization για ελληνικά και Unicode (διατηρεί τους τόνους)
    tokenizer.normalizer = normalizers.NFC()
    
    # Προ-tokenizer για μικτά κείμενα (ελληνικά και αγγλικά)
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.WhitespaceSplit(),
        pre_tokenizers.Punctuation(),
        pre_tokenizers.Digits(individual_digits=True)
    ])
    
    # Ορισμός ειδικών tokens
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=["<|endoftext|>", "<pad>", "<unk>", "<mask>", "[CITATION]"],
        continuing_subword_prefix=""
    )

    # Εκπαίδευση του tokenizer
    tokenizer.train_from_iterator(iterator, trainer=trainer)
    tokenizer.decoder = decoders.ByteLevel()

    # Δημιουργία φακέλου αποθήκευσης αν δεν υπάρχει
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Αποθήκευση του tokenizer
    tokenizer.save(os.path.join(output_dir, "tokenizer.json"))
    tokenizer.model.save(output_dir)  # Αποθηκεύει vocab.json και merges.txt

    print(f"✅ Ο tokenizer αποθηκεύτηκε στον φάκελο '{output_dir}'!")

    return tokenizer