"""Tokenizer setup for Sharona 32B² triple-context, Qwen2.5-Coder vocab 152064.

Uses existing Qwen special tokens that have pre-trained embeddings as
context markers. These visual/code tokens are unused in our text-only context.
"""

import os
import logging
from pathlib import Path
from typing import List, Optional

from tokenizers import Tokenizer

log = logging.getLogger(__name__)

BOS_TOKEN_ID = 151643
EOS_TOKEN_ID = 151645
PAD_TOKEN_ID = 151643

M1_TOKEN_ID = 151646
M2_TOKEN_ID = 151648
M3_TOKEN_ID = 151650
NL_TOKEN_ID = 151652
SEP_TOKEN_ID = 151657

SPECIAL_TOKENS = {
    "<|bos|>": BOS_TOKEN_ID,
    "<|eos|>": EOS_TOKEN_ID,
    "<|pad|>": PAD_TOKEN_ID,
    "<|m1|>": M1_TOKEN_ID,
    "<|m2|>": M2_TOKEN_ID,
    "<|m3|>": M3_TOKEN_ID,
    "<|nl|>": NL_TOKEN_ID,
    "<|sep|>": SEP_TOKEN_ID,
}

VOCAB_SIZE = 152_064

DEFAULT_TOKENIZER_PATHS = [
    str(Path(__file__).resolve().parent / "tokenizer.json"),
    "/workspace/models/Qwen2.5-Coder-32B-Instruct/tokenizer.json",
]


class SharonaTokenizer:
    def __init__(self, tokenizer: Tokenizer):
        self._tok = tokenizer
        self.n_vocab = VOCAB_SIZE

    def encode(self, text: str, allowed_special=None, disallowed_special=None) -> List[int]:
        return self._tok.encode(text, add_special_tokens=False).ids

    def decode(self, ids) -> str:
        return self._tok.decode(list(ids), skip_special_tokens=True)

    def decode_with_special(self, ids) -> str:
        return self._tok.decode(list(ids), skip_special_tokens=False)


def _find_tokenizer_json(explicit_path: Optional[str] = None) -> str:
    candidates = ([explicit_path] if explicit_path else []) + DEFAULT_TOKENIZER_PATHS
    for path in candidates:
        if path and os.path.exists(path):
            return path
    raise FileNotFoundError(
        f"tokenizer.json not found. Searched: {candidates}. "
        "Place tokenizer.json in data/ or pass an explicit path."
    )


def get_tokenizer(tokenizer_json_path: Optional[str] = None) -> "SharonaTokenizer":
    path = _find_tokenizer_json(tokenizer_json_path)
    log.info("Loading tokenizer from %s", path)
    tok = Tokenizer.from_file(path)
    tok.no_truncation()
    tok.no_padding()
    return SharonaTokenizer(tok)


def get_tokenizer_from_str(tokenizer_json_str: str) -> "SharonaTokenizer":
    tok = Tokenizer.from_str(tokenizer_json_str)
    tok.no_truncation()
    tok.no_padding()
    return SharonaTokenizer(tok)