# Türkçe BPE Tokenizer - HuggingFace tokenizers kullanır (hızlı)
import os
import pickle
from typing import List

class TurkishTokenizer:
    def __init__(self, vocab_size=50000):
        self.vocab_size = vocab_size
        self._hf_tokenizer = None

        # Special tokens
        self.PAD = "<pad>"
        self.UNK = "<unk>"
        self.BOS = "<bos>"
        self.EOS = "<eos>"
        self.USER = "<|user|>"
        self.ASSISTANT = "<|assistant|>"
        self.SYSTEM = "<|system|>"

        self.special_tokens = [self.PAD, self.UNK, self.BOS, self.EOS,
                               self.USER, self.ASSISTANT, self.SYSTEM]

        # stoi/itos için uyumluluk (train sonrası doldurulur)
        self.stoi = {}
        self.itos = {}

    def train(self, text_files: list):
        """BPE training using HuggingFace tokenizers (fast C++ backend)"""
        from tokenizers import Tokenizer
        from tokenizers.models import BPE
        from tokenizers.trainers import BpeTrainer
        from tokenizers.pre_tokenizers import ByteLevel
        from tokenizers.processors import TemplateProcessing

        print(f"Tokenizer eğitiliyor ({self.vocab_size} vocab, {len(text_files)} dosya)...")

        tokenizer = Tokenizer(BPE(unk_token=self.UNK))
        tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)

        trainer = BpeTrainer(
            vocab_size=self.vocab_size,
            special_tokens=self.special_tokens,
            show_progress=True,
            min_frequency=2,
        )

        tokenizer.train(text_files, trainer)
        self._hf_tokenizer = tokenizer
        self._sync_vocab()
        print(f"Vocab size: {len(self.stoi)}")

    def _sync_vocab(self):
        """HF tokenizer'dan stoi/itos senkronize et"""
        vocab = self._hf_tokenizer.get_vocab()
        self.stoi = vocab
        self.itos = {v: k for k, v in vocab.items()}

    def encode(self, text: str) -> List[int]:
        if self._hf_tokenizer is None:
            raise RuntimeError("Tokenizer eğitilmemiş veya yüklenmemiş!")
        return self._hf_tokenizer.encode(text).ids

    def decode(self, tokens: List[int]) -> str:
        if self._hf_tokenizer is None:
            raise RuntimeError("Tokenizer eğitilmemiş veya yüklenmemiş!")
        special_ids = {self.stoi.get(t, -1) for t in self.special_tokens}
        filtered = [t for t in tokens if t not in special_ids]
        return self._hf_tokenizer.decode(filtered, skip_special_tokens=True)

    def save(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self._hf_tokenizer.save(path.replace(".pkl", ".json"))
        # pkl de kaydet (uyumluluk için)
        with open(path, 'wb') as f:
            pickle.dump({'vocab_size': self.vocab_size}, f)
        print(f"Tokenizer saved to {path}")

    def load(self, path):
        from tokenizers import Tokenizer
        json_path = path.replace(".pkl", ".json")
        if os.path.exists(json_path):
            self._hf_tokenizer = Tokenizer.from_file(json_path)
            self._sync_vocab()
            print(f"Tokenizer loaded from {json_path} ({len(self.stoi)} vocab)")
        else:
            raise FileNotFoundError(f"Tokenizer dosyası bulunamadı: {json_path}")