import os
from typing import List, Optional
import tiktoken
from transformers import PreTrainedTokenizer

class BPETokenizer:
    def __init__(self, model_path: Optional[str] = None):
        if model_path:
            # In a real scenario, we would load the trained tiktoken model
            # For this prototype, we'll use the cl100k_base encoding as a base
            self.encoder = tiktoken.get_encoding("cl100k_base")
        else:
            self.encoder = tiktoken.get_encoding("cl100k_base")
            
        self.special_tokens = {
            "<pad>": 100001,
            "<eos>": 100002,
            "<bos>": 100003,
            "<unk>": 100004,
            "<sys>": 100005,
            "<user>": 100006,
            "<assistant>": 100007,
        }
        
        # In a real implementation, we'd extend the tiktoken vocab
        # For now, we'll just map them.
        
    def encode(self, text: str, bos: bool = False, eos: bool = False) -> List[int]:
        tokens = self.encoder.encode(text)
        if bos:
            tokens = [self.special_tokens["<bos>"]] + tokens
        if eos:
            tokens = tokens + [self.special_tokens["<eos>"]]
        return tokens

    def decode(self, tokens: List[int]) -> str:
        # Filter out special tokens for decoding
        valid_tokens = [t for t in tokens if t < 100001]
        return self.encoder.decode(valid_tokens)

    @property
    def vocab_size(self) -> int:
        return self.encoder.n_vocab + len(self.special_tokens)

if __name__ == "__main__":
    tokenizer = BPETokenizer()
    text = "Hello, how are you today?"
    encoded = tokenizer.encode(text, bos=True, eos=True)
    decoded = tokenizer.decode(encoded)
    print(f"Text: {text}")
    print(f"Encoded: {encoded}")
    print(f"Decoded: {decoded}")
    print(f"Vocab size: {tokenizer.vocab_size}")