#!/usr/bin/env python3
"""
Self-contained tokenizer for Ather 1 model.
Zero external tokenizer library dependencies (pure Python byte/vocab mapping).
"""

import os
import json
from typing import List, Dict

class StandaloneAtherTokenizer:
    def __init__(self, vocab_file: str):
        with open(vocab_file, "r", encoding="utf-8") as f:
            self.token_to_id = json.load(f)
        self.id_to_token = {v: k for k, v in self.token_to_id.items()}

        self.pad_token_id = self.token_to_id.get("<|pad|>", 0)
        self.bos_token_id = self.token_to_id.get("<|bos|>", 1)
        self.eos_token_id = self.token_to_id.get("<|endoftext|>", 2)
        self.unk_token_id = self.token_to_id.get("<|unk|>", 3)
        self.vocab_size = len(self.token_to_id)

    def encode(self, text: str, add_bos: bool = True) -> List[int]:
        tokens = [self.bos_token_id] if add_bos else []
        i = 0
        n = len(text)
        while i < n:
            matched = False
            for length in range(min(32, n - i), 0, -1):
                sub = text[i : i + length]
                if sub in self.token_to_id:
                    tokens.append(self.token_to_id[sub])
                    i += length
                    matched = True
                    break
            if not matched:
                tokens.append(self.unk_token_id)
                i += 1
        return tokens

    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
        parts = []
        for tid in token_ids:
            if skip_special_tokens and tid in (
                self.pad_token_id,
                self.bos_token_id,
                self.eos_token_id,
                self.unk_token_id,
            ):
                continue
            parts.append(self.id_to_token.get(tid, ""))
        return "".join(parts)